"""Database abstraction layer with MySQL and file-based backends. Usage: from refcurator.db import get_backend db = get_backend() rows = db.fetch_all("SELECT * FROM documents WHERE crawl_status = %s", ("pending",)) doc_id = db.insert_returning_id("INSERT INTO documents (...) VALUES (...)", params) """ from __future__ import annotations import json import logging from contextlib import contextmanager from pathlib import Path from typing import Any, Optional, Protocol, Sequence from refcurator.config import get_db_config, get_state_backend, get_state_dir logger = logging.getLogger("refcurator.db") class DatabaseBackend(Protocol): """Protocol for database backends.""" def execute(self, sql: str, params: Sequence = ()) -> int: """Execute a statement, return affected row count.""" ... def fetch_one(self, sql: str, params: Sequence = ()) -> Optional[dict]: """Fetch a single row as dict.""" ... def fetch_all(self, sql: str, params: Sequence = ()) -> list[dict]: """Fetch all rows as list of dicts.""" ... def insert_returning_id(self, sql: str, params: Sequence = ()) -> int: """Insert a row and return the auto-generated ID.""" ... def close(self) -> None: """Close the connection.""" ... class MySQLBackend: """MySQL backend using PyMySQL.""" def __init__(self, config: dict | None = None): import pymysql import pymysql.cursors if config is None: config = get_db_config().get("mysql", {}) self._conn = pymysql.connect( host=config.get("host", "localhost"), port=int(config.get("port", 3306)), user=config.get("user", "root"), password=config.get("password", ""), database=config.get("database", "reference_library"), charset="utf8mb4", cursorclass=pymysql.cursors.DictCursor, autocommit=True, ) def execute(self, sql: str, params: Sequence = ()) -> int: with self._conn.cursor() as cur: return cur.execute(sql, params) def fetch_one(self, sql: str, params: Sequence = ()) -> Optional[dict]: with self._conn.cursor() as cur: cur.execute(sql, params) return cur.fetchone() def fetch_all(self, sql: str, params: Sequence = ()) -> list[dict]: with self._conn.cursor() as cur: cur.execute(sql, params) return cur.fetchall() def insert_returning_id(self, sql: str, params: Sequence = ()) -> int: with self._conn.cursor() as cur: cur.execute(sql, params) return cur.lastrowid def close(self) -> None: self._conn.close() class FileBackend: """JSON file-based backend for use without MySQL. Stores data as JSON arrays in the state directory. Supports basic CRUD but not complex queries or JOINs. """ def __init__(self, state_dir: Path | None = None): self._dir = state_dir or get_state_dir() self._dir.mkdir(parents=True, exist_ok=True) self._cache: dict[str, list[dict]] = {} self._counters: dict[str, int] = {} self._load_counters() def _table_path(self, table: str) -> Path: return self._dir / f"{table}.json" def _load_table(self, table: str) -> list[dict]: if table not in self._cache: path = self._table_path(table) if path.is_file(): self._cache[table] = json.loads(path.read_text()) else: self._cache[table] = [] return self._cache[table] def _save_table(self, table: str) -> None: path = self._table_path(table) path.write_text(json.dumps(self._cache.get(table, []), indent=2, default=str)) def _load_counters(self) -> None: counter_path = self._dir / "_counters.json" if counter_path.is_file(): self._counters = json.loads(counter_path.read_text()) def _save_counters(self) -> None: counter_path = self._dir / "_counters.json" counter_path.write_text(json.dumps(self._counters, indent=2)) def _next_id(self, table: str) -> int: current = self._counters.get(table, 0) self._counters[table] = current + 1 self._save_counters() return current + 1 # --- Protocol methods --- # These provide basic support for the most common operations. # Complex SQL is not supported; use MySQL for full functionality. def execute(self, sql: str, params: Sequence = ()) -> int: """Basic INSERT/UPDATE/DELETE support via SQL pattern matching.""" sql_lower = sql.strip().lower() table = _extract_table_name(sql) if sql_lower.startswith("insert"): return self._handle_insert(table, sql, params) elif sql_lower.startswith("update"): return self._handle_update(table, sql, params) elif sql_lower.startswith("delete"): return self._handle_delete(table, sql, params) logger.warning("FileBackend: unsupported SQL operation: %s", sql[:60]) return 0 def fetch_one(self, sql: str, params: Sequence = ()) -> Optional[dict]: rows = self.fetch_all(sql, params) return rows[0] if rows else None def fetch_all(self, sql: str, params: Sequence = ()) -> list[dict]: table = _extract_table_name(sql) if not table: logger.warning("FileBackend: cannot extract table from: %s", sql[:60]) return [] # Hard-fail on SQL patterns that FileBackend cannot handle correctly _reject_unsupported_sql(sql) rows = self._load_table(table) # Basic WHERE clause filtering conditions = _extract_where_conditions(sql, params) if conditions: rows = [r for r in rows if _matches_conditions(r, conditions)] # Basic ORDER BY order_col = _extract_order_by(sql) if order_col: desc = "desc" in sql.lower().split("order by")[-1].lower() rows = sorted(rows, key=lambda r: r.get(order_col, ""), reverse=desc) # Basic LIMIT limit = _extract_limit(sql) if limit is not None: rows = rows[:limit] return rows def insert_returning_id(self, sql: str, params: Sequence = ()) -> int: table = _extract_table_name(sql) self._handle_insert(table, sql, params) return self._counters.get(table, 0) def close(self) -> None: pass # No connection to close # --- Internal handlers --- def _handle_insert(self, table: str, sql: str, params: Sequence) -> int: columns = _extract_insert_columns(sql) if not columns or len(columns) != len(params): logger.warning("FileBackend: column/param mismatch for INSERT into %s", table) return 0 row = dict(zip(columns, params)) pk = _primary_key_for(table) if pk and pk not in row: row[pk] = self._next_id(table) rows = self._load_table(table) rows.append(row) self._cache[table] = rows self._save_table(table) return 1 def _handle_update(self, table: str, sql: str, params: Sequence) -> int: rows = self._load_table(table) set_cols = _extract_set_columns(sql) conditions = _extract_where_conditions(sql, params[len(set_cols):]) set_values = list(params[:len(set_cols)]) count = 0 for row in rows: if _matches_conditions(row, conditions): for col, val in zip(set_cols, set_values): row[col] = val count += 1 if count > 0: self._save_table(table) return count def _handle_delete(self, table: str, sql: str, params: Sequence) -> int: rows = self._load_table(table) conditions = _extract_where_conditions(sql, params) before = len(rows) self._cache[table] = [r for r in rows if not _matches_conditions(r, conditions)] self._save_table(table) return before - len(self._cache[table]) class UnsupportedQueryError(Exception): """Raised when FileBackend encounters SQL it cannot handle correctly.""" pass def _reject_unsupported_sql(sql: str) -> None: """Raise UnsupportedQueryError if the SQL uses patterns FileBackend cannot handle. FileBackend only supports single-table SELECT with simple WHERE col = %s. JOINs, subqueries, aggregates, and GROUP BY would return wrong results silently. """ import re sql_upper = sql.upper() unsupported = [] if re.search(r"\bJOIN\b", sql_upper): unsupported.append("JOIN") if re.search(r"\bGROUP\s+BY\b", sql_upper): unsupported.append("GROUP BY") if re.search(r"\b(MAX|MIN|SUM|AVG|COUNT)\s*\(", sql_upper): unsupported.append("aggregate functions") if re.search(r"\bLEFT\s+JOIN\b", sql_upper): unsupported.append("LEFT JOIN") if re.search(r"\(\s*SELECT\b", sql_upper): unsupported.append("subquery") if unsupported: raise UnsupportedQueryError( f"FileBackend cannot execute queries with {', '.join(unsupported)}. " f"Configure MySQL or use 'backend: mysql' in pipeline_config.yaml. " f"Query: {sql[:80]}..." ) # --- SQL parsing helpers (minimal, covers common patterns) --- def _extract_table_name(sql: str) -> str: """Extract the primary table name from a SQL statement.""" import re sql_clean = sql.strip() # INSERT INTO table m = re.search(r"(?:insert\s+into|update|delete\s+from|from)\s+(\w+)", sql_clean, re.I) if m: return m.group(1) return "" def _extract_insert_columns(sql: str) -> list[str]: """Extract column names from INSERT INTO table (col1, col2, ...).""" import re m = re.search(r"\(([^)]+)\)\s*VALUES", sql, re.I) if m: return [c.strip() for c in m.group(1).split(",")] return [] def _extract_set_columns(sql: str) -> list[str]: """Extract column names from UPDATE ... SET col1 = %s, col2 = %s.""" import re m = re.search(r"SET\s+(.+?)(?:\s+WHERE|$)", sql, re.I | re.S) if m: return [c.strip().split("=")[0].strip() for c in m.group(1).split(",")] return [] def _extract_where_conditions(sql: str, params: Sequence) -> list[tuple[str, Any]]: """Extract simple col = %s conditions from WHERE clause.""" import re m = re.search(r"WHERE\s+(.+?)(?:\s+ORDER|\s+LIMIT|$)", sql, re.I | re.S) if not m: return [] where_clause = m.group(1) cols = re.findall(r"(\w+)\s*=\s*%s", where_clause) # Map to params (taking from the end of params for UPDATE, all for SELECT) return list(zip(cols, params[-len(cols):] if cols else [])) def _extract_order_by(sql: str) -> str: """Extract the first ORDER BY column.""" import re m = re.search(r"ORDER\s+BY\s+(\w+)", sql, re.I) return m.group(1) if m else "" def _extract_limit(sql: str) -> int | None: """Extract LIMIT value.""" import re m = re.search(r"LIMIT\s+(\d+)", sql, re.I) return int(m.group(1)) if m else None def _matches_conditions(row: dict, conditions: list[tuple[str, Any]]) -> bool: """Check if a row matches all WHERE conditions.""" return all(str(row.get(col)) == str(val) for col, val in conditions) def _primary_key_for(table: str) -> str: """Return the auto-increment primary key name for a table.""" pk_map = { "sources": "source_id", "documents": "doc_id", "distilled_content": "distill_id", "review_logs": "review_id", "topics": "topic_id", "export_jobs": "export_id", "pipeline_runs": "run_id", "pipeline_iteration_tracker": "tracker_id", "crawl_schedule": "schedule_id", "change_detection": "change_id", } return pk_map.get(table, "") # --- Factory --- def get_backend(backend_type: str | None = None) -> DatabaseBackend: """Create and return the appropriate database backend. Args: backend_type: 'mysql' or 'file'. If None, reads from pipeline config. Returns: A DatabaseBackend instance. """ if backend_type is None: backend_type = get_state_backend() if backend_type == "mysql": return MySQLBackend() return FileBackend() @contextmanager def db_session(backend_type: str | None = None): """Context manager for database sessions. Usage: with db_session() as db: rows = db.fetch_all("SELECT * FROM documents") """ db = get_backend(backend_type) try: yield db finally: db.close()