59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
"""PostgreSQL access: a connection pool, one connection per request, and two query helpers."""
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import psycopg2.extras
|
|
import psycopg2.pool
|
|
from flask import Flask, current_app, g
|
|
|
|
DEFAULT_DATABASE_URL = "postgresql://postgres:[email protected]:5432/postgres"
|
|
SCHEMA = Path(__file__).with_name("schema.sql")
|
|
|
|
|
|
def init_app(app: Flask) -> None:
|
|
pool = psycopg2.pool.ThreadedConnectionPool(1, 5, os.environ.get("DATABASE_URL", DEFAULT_DATABASE_URL))
|
|
conn = pool.getconn()
|
|
try:
|
|
with conn, conn.cursor() as cur:
|
|
cur.execute(SCHEMA.read_text())
|
|
finally:
|
|
pool.putconn(conn)
|
|
app.extensions["db_pool"] = pool
|
|
app.after_request(_finish_transaction)
|
|
app.teardown_appcontext(_release_connection)
|
|
|
|
|
|
def query(sql: str, params: tuple = ()) -> list[dict]:
|
|
with _connection().cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
|
cur.execute(sql, params)
|
|
return cur.fetchall() if cur.description else []
|
|
|
|
|
|
def query_one(sql: str, params: tuple = ()) -> dict | None:
|
|
rows = query(sql, params)
|
|
return rows[0] if rows else None
|
|
|
|
|
|
def _connection():
|
|
if "db" not in g:
|
|
g.db = current_app.extensions["db_pool"].getconn()
|
|
return g.db
|
|
|
|
|
|
def _finish_transaction(response):
|
|
# Commit only successful responses so a 4xx/5xx never leaves partial writes behind.
|
|
conn = g.get("db")
|
|
if conn is not None:
|
|
if response.status_code < 400:
|
|
conn.commit()
|
|
else:
|
|
conn.rollback()
|
|
return response
|
|
|
|
|
|
def _release_connection(_exc):
|
|
conn = g.pop("db", None)
|
|
if conn is not None:
|
|
conn.rollback() # no-op after commit; discards work if after_request never ran
|
|
current_app.extensions["db_pool"].putconn(conn)
|