32 lines
995 B
Python
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]
|