Build reusable bot framework

core/postgres.py

@@ -0,0 +1,289 @@
+"""
+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()]