30 lines
665 B
Python
30 lines
665 B
Python
|
from psycopg_pool import ConnectionPool
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from app.core.config import settings
|
||
|
|
|
||
|
|
connection_pool = ConnectionPool(settings.DATABASE_URL)
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def get_db_connection():
|
||
|
|
conn = connection_pool.getconn()
|
||
|
|
try:
|
||
|
|
yield conn
|
||
|
|
finally:
|
||
|
|
connection_pool.putconn(conn)
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def get_db_cursor(commit=False):
|
||
|
|
with get_db_connection() as conn:
|
||
|
|
cursor = conn.cursor()
|
||
|
|
try:
|
||
|
|
yield cursor
|
||
|
|
if commit:
|
||
|
|
conn.commit()
|
||
|
|
except Exception as e:
|
||
|
|
conn.rollback()
|
||
|
|
raise e
|
||
|
|
finally:
|
||
|
|
cursor.close()
|