first commit
This commit is contained in:
264
core/postgres.py
Normal file
264
core/postgres.py
Normal file
@@ -0,0 +1,264 @@
|
||||
"""
|
||||
postgres.py - Generic PostgreSQL CRUD module
|
||||
|
||||
Requires: pip install psycopg2-binary
|
||||
|
||||
Connection config from environment:
|
||||
DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASS
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
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 re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", name):
|
||||
raise ValueError(f"Invalid SQL identifier: {name}")
|
||||
return f'"{name}"'
|
||||
|
||||
|
||||
def _build_where(where, prefix=""):
|
||||
clauses = []
|
||||
params = {}
|
||||
for i, (col, val) in enumerate(where.items()):
|
||||
param_name = f"{prefix}{col}_{i}"
|
||||
safe_col = _safe_id(col)
|
||||
|
||||
if isinstance(val, tuple) and len(val) == 2:
|
||||
op, operand = val
|
||||
op = op.upper()
|
||||
allowed = {
|
||||
"=",
|
||||
"!=",
|
||||
"<",
|
||||
">",
|
||||
"<=",
|
||||
">=",
|
||||
"LIKE",
|
||||
"ILIKE",
|
||||
"IN",
|
||||
"IS",
|
||||
"IS NOT",
|
||||
}
|
||||
if op not in allowed:
|
||||
raise ValueError(f"Unsupported operator: {op}")
|
||||
if op == "IN":
|
||||
ph = ", ".join(f"%({param_name}_{j})s" for j in range(len(operand)))
|
||||
clauses.append(f"{safe_col} IN ({ph})")
|
||||
for j, item in enumerate(operand):
|
||||
params[f"{param_name}_{j}"] = item
|
||||
elif op in ("IS", "IS NOT"):
|
||||
clauses.append(f"{safe_col} {op} NULL")
|
||||
else:
|
||||
clauses.append(f"{safe_col} {op} %({param_name})s")
|
||||
params[param_name] = operand
|
||||
elif val is None:
|
||||
clauses.append(f"{safe_col} IS NULL")
|
||||
else:
|
||||
clauses.append(f"{safe_col} = %({param_name})s")
|
||||
params[param_name] = val
|
||||
|
||||
return " AND ".join(clauses), params
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_connection():
|
||||
conn = psycopg2.connect(**_get_config())
|
||||
try:
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_cursor(dict_cursor=True):
|
||||
with get_connection() as conn:
|
||||
factory = psycopg2.extras.RealDictCursor if dict_cursor else None
|
||||
cur = conn.cursor(cursor_factory=factory)
|
||||
try:
|
||||
yield cur
|
||||
finally:
|
||||
cur.close()
|
||||
|
||||
|
||||
def insert(table, data):
|
||||
columns = list(data.keys())
|
||||
placeholders = [f"%({col})s" for col in columns]
|
||||
safe_cols = [_safe_id(c) for c in columns]
|
||||
|
||||
query = f"""
|
||||
INSERT INTO {_safe_id(table)}
|
||||
({", ".join(safe_cols)})
|
||||
VALUES ({", ".join(placeholders)})
|
||||
RETURNING *
|
||||
"""
|
||||
with get_cursor() as cur:
|
||||
cur.execute(query, data)
|
||||
return dict(cur.fetchone()) if cur.rowcount 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}"
|
||||
if order_by:
|
||||
if isinstance(order_by, list):
|
||||
order_by = ", ".join(order_by)
|
||||
query += f" ORDER BY {order_by}"
|
||||
if limit is not None:
|
||||
query += f" LIMIT {int(limit)}"
|
||||
if offset is not None:
|
||||
query += f" OFFSET {int(offset)}"
|
||||
|
||||
with get_cursor() as cur:
|
||||
cur.execute(query, params)
|
||||
return [dict(row) for row in cur.fetchall()]
|
||||
|
||||
|
||||
def select_one(table, where):
|
||||
results = select(table, where=where, limit=1)
|
||||
return results[0] if results else None
|
||||
|
||||
|
||||
def update(table, data, where):
|
||||
set_columns = list(data.keys())
|
||||
set_clause = ", ".join(f"{_safe_id(col)} = %(set_{col})s" for col in set_columns)
|
||||
params = {f"set_{col}": val for col, val in data.items()}
|
||||
|
||||
where_clause, where_params = _build_where(where, prefix="where_")
|
||||
params.update(where_params)
|
||||
|
||||
query = f"""
|
||||
UPDATE {_safe_id(table)}
|
||||
SET {set_clause}
|
||||
WHERE {where_clause}
|
||||
RETURNING *
|
||||
"""
|
||||
with get_cursor() as cur:
|
||||
cur.execute(query, params)
|
||||
return [dict(row) for row in cur.fetchall()]
|
||||
|
||||
|
||||
def delete(table, where):
|
||||
where_clause, params = _build_where(where)
|
||||
query = f"""
|
||||
DELETE FROM {_safe_id(table)}
|
||||
WHERE {where_clause}
|
||||
RETURNING *
|
||||
"""
|
||||
with get_cursor() as cur:
|
||||
cur.execute(query, params)
|
||||
return [dict(row) for row in cur.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 cur:
|
||||
cur.execute(query, params)
|
||||
return cur.fetchone()["count"]
|
||||
|
||||
|
||||
def exists(table, where):
|
||||
return count(table, where) > 0
|
||||
|
||||
|
||||
def upsert(table, data, conflict_columns):
|
||||
columns = list(data.keys())
|
||||
placeholders = [f"%({col})s" for col in columns]
|
||||
safe_cols = [_safe_id(c) for c in columns]
|
||||
conflict_cols = [_safe_id(c) for c in conflict_columns]
|
||||
|
||||
update_cols = [c for c in columns if c not in conflict_columns]
|
||||
update_clause = ", ".join(
|
||||
f"{_safe_id(c)} = EXCLUDED.{_safe_id(c)}" for c in update_cols
|
||||
)
|
||||
|
||||
query = f"""
|
||||
INSERT INTO {_safe_id(table)}
|
||||
({", ".join(safe_cols)})
|
||||
VALUES ({", ".join(placeholders)})
|
||||
ON CONFLICT ({", ".join(conflict_cols)})
|
||||
DO UPDATE SET {update_clause}
|
||||
RETURNING *
|
||||
"""
|
||||
with get_cursor() as cur:
|
||||
cur.execute(query, data)
|
||||
return dict(cur.fetchone()) if cur.rowcount else None
|
||||
|
||||
|
||||
def insert_many(table, rows):
|
||||
if not rows:
|
||||
return 0
|
||||
columns = list(rows[0].keys())
|
||||
safe_cols = [_safe_id(c) for c in columns]
|
||||
query = f"""
|
||||
INSERT INTO {_safe_id(table)}
|
||||
({", ".join(safe_cols)})
|
||||
VALUES %s
|
||||
"""
|
||||
template = f"({', '.join(f'%({col})s' for col in columns)})"
|
||||
with get_cursor() as cur:
|
||||
psycopg2.extras.execute_values(
|
||||
cur, query, rows, template=template, page_size=100
|
||||
)
|
||||
return cur.rowcount
|
||||
|
||||
|
||||
def execute(query, params=None):
|
||||
with get_cursor() as cur:
|
||||
cur.execute(query, params or {})
|
||||
if cur.description:
|
||||
return [dict(row) for row in cur.fetchall()]
|
||||
return cur.rowcount
|
||||
|
||||
|
||||
def table_exists(table):
|
||||
with get_cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT EXISTS (
|
||||
SELECT FROM information_schema.tables
|
||||
WHERE table_schema = 'public' AND table_name = %(table)s
|
||||
)
|
||||
""",
|
||||
{"table": table},
|
||||
)
|
||||
return cur.fetchone()["exists"]
|
||||
|
||||
|
||||
def get_columns(table):
|
||||
with get_cursor() as cur:
|
||||
cur.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(row) for row in cur.fetchall()]
|
||||
Reference in New Issue
Block a user