55 lines
1.2 KiB
Python
55 lines
1.2 KiB
Python
from contextlib import contextmanager
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
from config import Config
|
|
|
|
|
|
def get_connection():
|
|
return psycopg2.connect(
|
|
host=Config.DB_HOST,
|
|
dbname=Config.DB_NAME,
|
|
user=Config.DB_USER,
|
|
password=Config.DB_PASSWORD,
|
|
port=Config.DB_PORT,
|
|
cursor_factory=psycopg2.extras.RealDictCursor,
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def get_cursor(commit=False):
|
|
conn = get_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
yield cur
|
|
if commit:
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
cur.close()
|
|
conn.close()
|
|
|
|
|
|
def fetch_one(query, params=None):
|
|
with get_cursor() as cur:
|
|
cur.execute(query, params or ())
|
|
return cur.fetchone()
|
|
|
|
|
|
def fetch_all(query, params=None):
|
|
with get_cursor() as cur:
|
|
cur.execute(query, params or ())
|
|
return cur.fetchall()
|
|
|
|
|
|
def execute(query, params=None):
|
|
with get_cursor(commit=True) as cur:
|
|
cur.execute(query, params or ())
|
|
|
|
|
|
def execute_returning(query, params=None):
|
|
with get_cursor(commit=True) as cur:
|
|
cur.execute(query, params or ())
|
|
return cur.fetchone()
|