pgbot/pgbotlib/dbstuff.py

32 lines
995 B
Python

import random
import psycopg
class DBConn:
def __init__(self, *args, **kwargs) -> None:
self.connection = psycopg.connect(*args, **kwargs)
self.cursor = self.connection.cursor()
def update(self, query: str, values: tuple) -> None:
failure = None
try:
self.cursor.execute('SAVEPOINT sp1')
self.cursor.execute(query, values)
except Exception as e:
failure = e
self.cursor.execute('ROLLBACK TO SAVEPOINT sp1')
else:
self.cursor.execute('RELEASE SAVEPOINT sp1')
self.connection.commit()
if failure:
raise failure
def query_raw(self, query: str, values: tuple) -> list:
self.cursor.execute(query, values)
return self.cursor.fetchall()
def query_random(self, query: str, values: tuple) -> str:
result = self.query_raw(query, values)
if not result:
return None
return random.choice(result)[0]