von/pgbot
1
0
Fork 0
Code Issues 1 Pull requests Projects Releases Packages Wiki Activity Actions

update populate script, fuck if I know why

Browse source
This commit is contained in:
Von Random 2024-08-10 04:51:58 +03:00
parent 836a72fa8e
commit 85ee606788
1 changed files with 23 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
populate
View file

@ -7,31 +7,32 @@ with open('config.yml', 'r', encoding='UTF-8') as data:
config = yaml.safe_load(data.read()) config = yaml.safe_load(data.read())
with open(config['response_tokens'], 'r', encoding='UTF-8') as data: with open(config['response_tokens'], 'r', encoding='UTF-8') as data:
valid_tokens = {item for item in yaml.safe_load(data.read())} valid_tokens = {item for item in yaml.safe_load(data.read())}
#with open(sys.argv[1], 'r', encoding='UTF-8') as data:
# phrases = yaml.safe_load(data.read())
with open(sys.argv[1], 'r', encoding='UTF-8') as data: with open(sys.argv[1], 'r', encoding='UTF-8') as data:
phrases = yaml.safe_load(data.read()) names = yaml.safe_load(data.read())
with psycopg.connect(config['db_spec']) as conn: with psycopg.connect(config['db_spec']) as conn:
query_phrases = 'INSERT INTO responses (tokens, response) VALUES (%s, %s)' query_phrases = 'INSERT INTO responses (tokens, response) VALUES (%s, %s)'
query_names = 'INSERT INTO names (tg_id, name) VALUES (%s, %s)'
with conn.cursor() as cur: with conn.cursor() as cur:
for regexref, responses in phrases.items(): #for regexref, responses in phrases.items():
tokens = set(regexref.split(',')) # tokens = set(regexref.split(','))
if tokens != tokens & valid_tokens: # if tokens != tokens & valid_tokens:
print(f'{str(tokens)} failed to add!') # print(f'{str(tokens)} failed to add!')
continue # continue
token_string = ','.join(sorted(tokens)) # token_string = ','.join(sorted(tokens))
for response in responses: # for response in responses:
cur.execute('SAVEPOINT sp1') # cur.execute('SAVEPOINT sp1')
try: # try:
cur.execute( query_phrases, (token_string, response)) # cur.execute(query_phrases, (token_string, response))
except psycopg.errors.UniqueViolation as err: # except psycopg.errors.UniqueViolation as err:
cur.execute('ROLLBACK TO SAVEPOINT sp1') # cur.execute('ROLLBACK TO SAVEPOINT sp1')
# print(err) # # print(err)
continue # continue
cur.execute('RELEASE SAVEPOINT sp1') # cur.execute('RELEASE SAVEPOINT sp1')
conn.commit()
#for item in names:
# usernames = names[item]
# for username in usernames:
# cur.execute('INSERT INTO names (tg_id, name) VALUES (%s, %s)',
# (item, username))
#conn.commit() #conn.commit()
for tg_id, usernames in names.items():
for username in usernames:
cur.execute(query_names, (tg_id, username))
conn.commit()