Add initial SQLmem package structure with SQL parser, cache manager, column registry, and tests

src/sqlmem/cache.py

@@ -0,0 +1,158 @@
+import atexit
+import signal
+import sqlite3
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+
+from loguru import logger
+
+import sqlmem._meta as _meta
+
+SCHEMA_VERSION = 1
+
+
+class CacheManager:
+    def __init__(self, db_path: Path, backup_interval: int) -> None:
+        self._db_path = db_path
+        self._backup_interval = backup_interval
+        self._mem_conn = sqlite3.connect(":memory:", check_same_thread=False)
+        self._lock = threading.Lock()
+        self._closed = False
+
+        self._ensure_meta_tables()
+        self._load_from_disk()
+        self._start_backup_thread()
+
+        atexit.register(self._backup_to_disk)
+        signal.signal(signal.SIGTERM, self._on_sigterm)
+
+    @property
+    def connection(self) -> sqlite3.Connection:
+        return self._mem_conn
+
+    def _ensure_meta_tables(self) -> None:
+        self._mem_conn.executescript("""
+            CREATE TABLE IF NOT EXISTS _sqlmem_meta (
+                key   TEXT PRIMARY KEY,
+                value TEXT NOT NULL
+            );
+            CREATE TABLE IF NOT EXISTS _sqlmem_tables (
+                table_name      TEXT PRIMARY KEY,
+                last_refresh_at TEXT NOT NULL,
+                row_count       INTEGER
+            );
+            CREATE TABLE IF NOT EXISTS _sqlmem_columns (
+                table_name  TEXT NOT NULL,
+                column_name TEXT NOT NULL,
+                PRIMARY KEY (table_name, column_name)
+            );
+        """)
+        self._mem_conn.execute(
+            "INSERT OR IGNORE INTO _sqlmem_meta (key, value) VALUES (?, ?)",
+            ("app_version", _meta.__version__),
+        )
+        self._mem_conn.execute(
+            "INSERT OR IGNORE INTO _sqlmem_meta (key, value) VALUES (?, ?)",
+            ("schema_version", str(SCHEMA_VERSION)),
+        )
+        self._mem_conn.execute(
+            "INSERT OR IGNORE INTO _sqlmem_meta (key, value) VALUES (?, ?)",
+            ("created_at", _now()),
+        )
+        self._mem_conn.commit()
+
+    def _load_from_disk(self) -> None:
+        if not self._db_path.exists():
+            logger.info(f"No cache file found at {self._db_path}, starting fresh.")
+            return
+
+        logger.info(f"Loading cache from {self._db_path}")
+        disk_conn = sqlite3.connect(self._db_path)
+        try:
+            schema_version = disk_conn.execute(
+                "SELECT value FROM _sqlmem_meta WHERE key = 'schema_version'"
+            ).fetchone()
+            if schema_version is None or int(schema_version[0]) != SCHEMA_VERSION:
+                logger.warning("Cache schema version mismatch — discarding cache file, starting fresh.")
+                disk_conn.close()
+                return
+
+            disk_conn.backup(self._mem_conn)
+            logger.info("Cache loaded from disk successfully.")
+        except Exception as e:
+            logger.error(f"Failed to load cache from disk: {e} — starting fresh.")
+        finally:
+            disk_conn.close()
+
+    def _backup_to_disk(self) -> None:
+        if self._closed:
+            return
+        logger.info(f"Backing up cache to {self._db_path}")
+        try:
+            with self._lock:
+                disk_conn = sqlite3.connect(self._db_path)
+                self._mem_conn.backup(disk_conn)
+                disk_conn.close()
+            logger.info("Cache backup complete.")
+        except Exception as e:
+            logger.error(f"Cache backup failed: {e}")
+
+    def _start_backup_thread(self) -> None:
+        def loop() -> None:
+            event = threading.Event()
+            while not event.wait(self._backup_interval):
+                self._backup_to_disk()
+
+        t = threading.Thread(target=loop, daemon=True, name="sqlmem-backup")
+        t.start()
+        logger.debug(f"Backup thread started (interval={self._backup_interval}s)")
+
+    def _on_sigterm(self, signum, frame) -> None:
+        logger.info("SIGTERM received — flushing cache to disk.")
+        self._backup_to_disk()
+
+    def mark_table_refreshed(self, table: str, row_count: int) -> None:
+        with self._lock:
+            self._mem_conn.execute(
+                """
+                INSERT INTO _sqlmem_tables (table_name, last_refresh_at, row_count)
+                VALUES (?, ?, ?)
+                ON CONFLICT(table_name) DO UPDATE SET
+                    last_refresh_at = excluded.last_refresh_at,
+                    row_count = excluded.row_count
+                """,
+                (table, _now(), row_count),
+            )
+            self._mem_conn.commit()
+
+    def is_table_cached(self, table: str) -> bool:
+        row = self._mem_conn.execute(
+            "SELECT 1 FROM _sqlmem_tables WHERE table_name = ?", (table,)
+        ).fetchone()
+        return row is not None
+
+    def load_table(self, table: str, columns: list[str], source_conn: sqlite3.Connection) -> None:
+        cols = ", ".join(columns)
+        logger.info(f"Fetching {table!r} columns [{cols}] from source DB")
+        rows = source_conn.execute(f"SELECT {cols} FROM {table}").fetchall()
+
+        with self._lock:
+            self._mem_conn.execute(f"DROP TABLE IF EXISTS {table}")
+            col_defs = ", ".join(f"{c} TEXT" for c in columns)
+            self._mem_conn.execute(f"CREATE TABLE {table} ({col_defs})")
+            placeholders = ", ".join("?" * len(columns))
+            self._mem_conn.executemany(f"INSERT INTO {table} VALUES ({placeholders})", rows)
+            self._mem_conn.commit()
+
+        self.mark_table_refreshed(table, len(rows))
+        logger.info(f"Table {table!r} cached ({len(rows)} rows, columns: {columns})")
+
+    def close(self) -> None:
+        self._backup_to_disk()
+        self._closed = True
+        self._mem_conn.close()
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()