290 lines
9.5 KiB
Python
290 lines
9.5 KiB
Python
"""
|
|
postgres.py - Small parameterized PostgreSQL CRUD layer
|
|
|
|
Connection configuration is read from DB_HOST, DB_PORT, DB_NAME, DB_USER,
|
|
and DB_PASS. Raw SQL remains available through execute() for domain services.
|
|
"""
|
|
|
|
from contextlib import contextmanager
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
from dotenv import load_dotenv
|
|
|
|
|
|
load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)
|
|
|
|
IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
|
|
|
|
|
def _get_config():
|
|
return {
|
|
"host": os.environ.get("DB_HOST", "localhost"),
|
|
"port": int(os.environ.get("DB_PORT", 5432)),
|
|
"dbname": os.environ.get("DB_NAME", "app"),
|
|
"user": os.environ.get("DB_USER", "app"),
|
|
"password": os.environ.get("DB_PASS", ""),
|
|
}
|
|
|
|
|
|
def _safe_id(name):
|
|
if not isinstance(name, str) or not IDENTIFIER.fullmatch(name):
|
|
raise ValueError(f"Invalid SQL identifier: {name}")
|
|
return f'"{name}"'
|
|
|
|
|
|
def _build_where(where, prefix=""):
|
|
if not isinstance(where, dict):
|
|
raise ValueError("where must be a dictionary")
|
|
clauses = []
|
|
params = {}
|
|
for index, (column, value) in enumerate(where.items()):
|
|
paramName = f"{prefix}{column}_{index}"
|
|
safeColumn = _safe_id(column)
|
|
|
|
if isinstance(value, tuple) and len(value) == 2:
|
|
operator, operand = value
|
|
operator = str(operator).upper()
|
|
allowed = {"=", "!=", "<", ">", "<=", ">=", "LIKE", "ILIKE", "IN"}
|
|
if operator not in allowed:
|
|
raise ValueError(f"Unsupported operator: {operator}")
|
|
if operator == "IN":
|
|
values = list(operand)
|
|
if not values:
|
|
clauses.append("FALSE")
|
|
continue
|
|
placeholders = []
|
|
for itemIndex, item in enumerate(values):
|
|
itemName = f"{paramName}_{itemIndex}"
|
|
placeholders.append(f"%({itemName})s")
|
|
params[itemName] = item
|
|
clauses.append(f"{safeColumn} IN ({', '.join(placeholders)})")
|
|
else:
|
|
clauses.append(f"{safeColumn} {operator} %({paramName})s")
|
|
params[paramName] = operand
|
|
elif value is None:
|
|
clauses.append(f"{safeColumn} IS NULL")
|
|
else:
|
|
clauses.append(f"{safeColumn} = %({paramName})s")
|
|
params[paramName] = value
|
|
return " AND ".join(clauses), params
|
|
|
|
|
|
def _order_clause(orderBy):
|
|
if not orderBy:
|
|
return ""
|
|
items = orderBy if isinstance(orderBy, (list, tuple)) else str(orderBy).split(",")
|
|
safeItems = []
|
|
for item in items:
|
|
if isinstance(item, (list, tuple)):
|
|
if len(item) != 2:
|
|
raise ValueError("order tuple must contain column and direction")
|
|
column, direction = item
|
|
else:
|
|
parts = str(item).strip().split()
|
|
if not parts or len(parts) > 2:
|
|
raise ValueError(f"Invalid order expression: {item}")
|
|
column = parts[0]
|
|
direction = parts[1] if len(parts) == 2 else "ASC"
|
|
direction = str(direction).upper()
|
|
if direction not in {"ASC", "DESC"}:
|
|
raise ValueError(f"Invalid order direction: {direction}")
|
|
safeItems.append(f"{_safe_id(column)} {direction}")
|
|
return ", ".join(safeItems)
|
|
|
|
|
|
@contextmanager
|
|
def get_connection():
|
|
connection = psycopg2.connect(**_get_config())
|
|
try:
|
|
yield connection
|
|
connection.commit()
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
@contextmanager
|
|
def get_cursor(dict_cursor=True):
|
|
with get_connection() as connection:
|
|
factory = psycopg2.extras.RealDictCursor if dict_cursor else None
|
|
cursor = connection.cursor(cursor_factory=factory)
|
|
try:
|
|
yield cursor
|
|
finally:
|
|
cursor.close()
|
|
|
|
|
|
def insert(table, data):
|
|
if not isinstance(data, dict) or not data:
|
|
raise ValueError("insert data cannot be empty")
|
|
columns = list(data.keys())
|
|
query = f"""
|
|
INSERT INTO {_safe_id(table)} ({', '.join(_safe_id(col) for col in columns)})
|
|
VALUES ({', '.join(f'%({col})s' for col in columns)})
|
|
RETURNING *
|
|
"""
|
|
with get_cursor() as cursor:
|
|
cursor.execute(query, data)
|
|
record = cursor.fetchone()
|
|
return dict(record) if record else None
|
|
|
|
|
|
def select(table, where=None, order_by=None, limit=None, offset=None):
|
|
query = f"SELECT * FROM {_safe_id(table)}"
|
|
params = {}
|
|
if where:
|
|
clauses, params = _build_where(where)
|
|
query += f" WHERE {clauses}"
|
|
orderClause = _order_clause(order_by)
|
|
if orderClause:
|
|
query += f" ORDER BY {orderClause}"
|
|
if limit is not None:
|
|
limit = int(limit)
|
|
if limit < 0:
|
|
raise ValueError("limit cannot be negative")
|
|
query += " LIMIT %(query_limit)s"
|
|
params["query_limit"] = limit
|
|
if offset is not None:
|
|
offset = int(offset)
|
|
if offset < 0:
|
|
raise ValueError("offset cannot be negative")
|
|
query += " OFFSET %(query_offset)s"
|
|
params["query_offset"] = offset
|
|
with get_cursor() as cursor:
|
|
cursor.execute(query, params)
|
|
return [dict(record) for record in cursor.fetchall()]
|
|
|
|
|
|
def select_one(table, where):
|
|
records = select(table, where=where, limit=1)
|
|
return records[0] if records else None
|
|
|
|
|
|
def update(table, data, where):
|
|
if not isinstance(data, dict) or not data:
|
|
raise ValueError("update data cannot be empty")
|
|
if not isinstance(where, dict) or not where:
|
|
raise ValueError("update requires a non-empty where clause")
|
|
setClause = ", ".join(f"{_safe_id(col)} = %(set_{col})s" for col in data)
|
|
params = {f"set_{col}": value for col, value in data.items()}
|
|
whereClause, whereParams = _build_where(where, prefix="where_")
|
|
params.update(whereParams)
|
|
query = f"""
|
|
UPDATE {_safe_id(table)} SET {setClause}
|
|
WHERE {whereClause}
|
|
RETURNING *
|
|
"""
|
|
with get_cursor() as cursor:
|
|
cursor.execute(query, params)
|
|
return [dict(record) for record in cursor.fetchall()]
|
|
|
|
|
|
def delete(table, where):
|
|
if not isinstance(where, dict) or not where:
|
|
raise ValueError("delete requires a non-empty where clause")
|
|
whereClause, params = _build_where(where)
|
|
query = f"DELETE FROM {_safe_id(table)} WHERE {whereClause} RETURNING *"
|
|
with get_cursor() as cursor:
|
|
cursor.execute(query, params)
|
|
return [dict(record) for record in cursor.fetchall()]
|
|
|
|
|
|
def count(table, where=None):
|
|
query = f"SELECT COUNT(*) AS count FROM {_safe_id(table)}"
|
|
params = {}
|
|
if where:
|
|
clauses, params = _build_where(where)
|
|
query += f" WHERE {clauses}"
|
|
with get_cursor() as cursor:
|
|
cursor.execute(query, params)
|
|
return cursor.fetchone()["count"]
|
|
|
|
|
|
def exists(table, where):
|
|
return count(table, where) > 0
|
|
|
|
|
|
def upsert(table, data, conflict_columns):
|
|
if not isinstance(data, dict) or not data:
|
|
raise ValueError("upsert data cannot be empty")
|
|
if not conflict_columns:
|
|
raise ValueError("conflict columns are required")
|
|
columns = list(data.keys())
|
|
updates = [column for column in columns if column not in conflict_columns]
|
|
action = "DO NOTHING"
|
|
if updates:
|
|
action = "DO UPDATE SET " + ", ".join(
|
|
f"{_safe_id(column)} = EXCLUDED.{_safe_id(column)}" for column in updates
|
|
)
|
|
query = f"""
|
|
INSERT INTO {_safe_id(table)} ({', '.join(_safe_id(col) for col in columns)})
|
|
VALUES ({', '.join(f'%({col})s' for col in columns)})
|
|
ON CONFLICT ({', '.join(_safe_id(col) for col in conflict_columns)})
|
|
{action}
|
|
RETURNING *
|
|
"""
|
|
with get_cursor() as cursor:
|
|
cursor.execute(query, data)
|
|
record = cursor.fetchone()
|
|
return dict(record) if record else None
|
|
|
|
|
|
def insert_many(table, rows):
|
|
if not rows:
|
|
return 0
|
|
columns = list(rows[0].keys())
|
|
if any(list(row.keys()) != columns for row in rows):
|
|
raise ValueError("all rows must use the same columns in the same order")
|
|
query = f"""
|
|
INSERT INTO {_safe_id(table)} ({', '.join(_safe_id(col) for col in columns)})
|
|
VALUES %s
|
|
"""
|
|
template = f"({', '.join(f'%({column})s' for column in columns)})"
|
|
with get_cursor() as cursor:
|
|
psycopg2.extras.execute_values(
|
|
cursor, query, rows, template=template, page_size=100
|
|
)
|
|
return cursor.rowcount
|
|
|
|
|
|
def execute(query, params=None):
|
|
with get_cursor() as cursor:
|
|
cursor.execute(query, params or {})
|
|
if cursor.description:
|
|
return [dict(record) for record in cursor.fetchall()]
|
|
return cursor.rowcount
|
|
|
|
|
|
def table_exists(table):
|
|
with get_cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
SELECT EXISTS (
|
|
SELECT FROM information_schema.tables
|
|
WHERE table_schema = 'public' AND table_name = %(table)s
|
|
)
|
|
""",
|
|
{"table": table},
|
|
)
|
|
return cursor.fetchone()["exists"]
|
|
|
|
|
|
def get_columns(table):
|
|
with get_cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
SELECT column_name, data_type, is_nullable, column_default
|
|
FROM information_schema.columns
|
|
WHERE table_schema = 'public' AND table_name = %(table)s
|
|
ORDER BY ordinal_position
|
|
""",
|
|
{"table": table},
|
|
)
|
|
return [dict(record) for record in cursor.fetchall()]
|