Upload mcp_database_universal/engines/sqlite.py with huggingface_hub

mcp_database_universal/engines/sqlite.py

@@ -0,0 +1,263 @@
+"""SQLite engine — built-in, no external dependencies."""
+
+import sqlite3
+import os
+import time
+from mcp_database_universal.engines.base import (
+    BaseEngine, DBInfo, ColumnInfo, TableInfo, TableDetail,
+    IndexInfo, ForeignKeyInfo, TableStats, QueryResult,
+)
+
+
+class SQLiteEngine(BaseEngine):
+    def __init__(self, path: str = ":memory:", read_only: bool = True):
+        self.path = path
+        self._read_only = read_only
+        self._conn: sqlite3.Connection | None = None
+
+    async def connect(self) -> None:
+        if self.path == ":memory:":
+            self._conn = sqlite3.connect(":memory:")
+        else:
+            dir_path = os.path.dirname(self.path)
+            if dir_path:
+                os.makedirs(dir_path, exist_ok=True)
+            if self._read_only:
+                uri = f"file:{self.path}?mode=ro"
+                self._conn = sqlite3.connect(uri, uri=True)
+            else:
+                self._conn = sqlite3.connect(self.path)
+        self._conn.row_factory = sqlite3.Row
+        if not self._read_only:
+            try:
+                self._conn.execute("PRAGMA journal_mode=WAL")
+            except sqlite3.OperationalError:
+                pass
+        self._conn.execute("PRAGMA foreign_keys=ON")
+
+    async def disconnect(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    def _ensure_conn(self) -> sqlite3.Connection:
+        if self._conn is None:
+            raise RuntimeError("Not connected. Call connect() first.")
+        return self._conn
+
+    async def get_db_info(self) -> DBInfo:
+        conn = self._ensure_conn()
+        version = conn.execute("SELECT sqlite_version()").fetchone()[0]
+        if self.path == ":memory:":
+            name = ":memory:"
+            size = "~0 KB"
+        else:
+            name = self.path
+            try:
+                size_bytes = os.path.getsize(self.path)
+                if size_bytes < 1024:
+                    size = f"~{size_bytes} B"
+                elif size_bytes < 1024 * 1024:
+                    size = f"~{size_bytes // 1024} KB"
+                else:
+                    size = f"~{size_bytes // (1024 * 1024)} MB"
+            except OSError:
+                size = "unknown"
+        return DBInfo(engine="sqlite", version=version, name=name, size_approx=size)
+
+    async def get_tables(self) -> list[TableInfo]:
+        conn = self._ensure_conn()
+        rows = conn.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
+        ).fetchall()
+
+        tables = []
+        for row in rows:
+            name = row["name"]
+            try:
+                count = conn.execute(f'SELECT COUNT(*) FROM "{name}"').fetchone()[0]
+            except Exception:
+                count = 0
+
+            try:
+                cols = conn.execute(f'PRAGMA table_info("{name}")').fetchall()
+                col_count = len(cols)
+                has_pk = any(c["pk"] for c in cols)
+            except Exception:
+                col_count = 0
+                has_pk = False
+
+            try:
+                fks = conn.execute(f'PRAGMA foreign_key_list("{name}")').fetchall()
+                fk_out = len(fks)
+            except Exception:
+                fk_out = 0
+
+            tables.append(TableInfo(
+                name=name,
+                row_count=count,
+                column_count=col_count,
+                has_primary_key=has_pk,
+                foreign_keys_out=fk_out,
+                foreign_keys_in=0,
+            ))
+
+        for t in tables:
+            for other in tables:
+                if other.name == t.name:
+                    continue
+                try:
+                    fks = conn.execute(f'PRAGMA foreign_key_list("{other.name}")').fetchall()
+                    for fk in fks:
+                        if fk["table"] == t.name:
+                            t.foreign_keys_in += 1
+                except Exception:
+                    pass
+
+        return tables
+
+    async def get_table_detail(self, table: str) -> TableDetail:
+        conn = self._ensure_conn()
+
+        col_rows = conn.execute(f'PRAGMA table_info("{table}")').fetchall()
+        columns = []
+        pk_name = None
+        for c in col_rows:
+            is_pk = bool(c["pk"])
+            if is_pk:
+                pk_name = c["name"]
+            columns.append(ColumnInfo(
+                name=c["name"],
+                type=c["type"] or "TEXT",
+                nullable=not c["notnull"],
+                default=c["dflt_value"],
+                is_primary_key=is_pk,
+            ))
+
+        fk_rows = conn.execute(f'PRAGMA foreign_key_list("{table}")').fetchall()
+        foreign_keys = []
+        fk_cols = set()
+        for fk in fk_rows:
+            foreign_keys.append(ForeignKeyInfo(
+                column=fk["from"],
+                references_table=fk["table"],
+                references_column=fk["to"],
+            ))
+            fk_cols.add(fk["from"])
+
+        for col in columns:
+            if col.name in fk_cols:
+                col.is_foreign_key = True
+                for fk in foreign_keys:
+                    if fk.column == col.name:
+                        col.foreign_key_table = fk.references_table
+                        col.foreign_key_column = fk.references_column
+                        break
+
+        idx_rows = conn.execute(f'PRAGMA index_list("{table}")').fetchall()
+        indexes = []
+        for idx in idx_rows:
+            idx_info = conn.execute(f'PRAGMA index_info("{idx["name"]}")').fetchall()
+            idx_cols = [i["name"] for i in idx_info]
+            indexes.append(IndexInfo(
+                name=idx["name"],
+                columns=idx_cols,
+                unique=bool(idx["unique"]),
+            ))
+
+        try:
+            sample = conn.execute(f'SELECT * FROM "{table}" LIMIT 5').fetchall()
+            sample_data = [dict(row) for row in sample]
+        except Exception:
+            sample_data = []
+
+        stats = await self.get_table_stats(table)
+
+        return TableDetail(
+            name=table,
+            columns=columns,
+            indexes=indexes,
+            primary_key=pk_name,
+            foreign_keys=foreign_keys,
+            sample_data=sample_data,
+            stats=stats,
+        )
+
+    async def execute_query(self, sql: str, params: dict | None = None) -> QueryResult:
+        conn = self._ensure_conn()
+        start = time.monotonic()
+        try:
+            if params:
+                cursor = conn.execute(sql, params)
+            else:
+                cursor = conn.execute(sql)
+            rows = cursor.fetchall()
+            columns = [desc[0] for desc in cursor.description] if cursor.description else []
+            elapsed = int((time.monotonic() - start) * 1000)
+            result_rows = [dict(row) for row in rows]
+            return QueryResult(
+                columns=columns,
+                rows=result_rows,
+                row_count=len(result_rows),
+                truncated=False,
+                execution_time_ms=elapsed,
+                sql=sql,
+            )
+        except Exception as e:
+            elapsed = int((time.monotonic() - start) * 1000)
+            return QueryResult(
+                sql=sql,
+                execution_time_ms=elapsed,
+                warning=f"Error: {str(e)}",
+            )
+
+    async def get_sample_data(self, table: str, limit: int = 5) -> QueryResult:
+        sql = f'SELECT * FROM "{table}" LIMIT {limit}'
+        return await self.execute_query(sql)
+
+    async def get_table_stats(self, table: str) -> TableStats:
+        conn = self._ensure_conn()
+        try:
+            row_count = conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]
+        except Exception:
+            row_count = 0
+
+        try:
+            page_count = conn.execute("PRAGMA page_count").fetchone()[0]
+            page_size = conn.execute("PRAGMA page_size").fetchone()[0]
+            total_bytes = page_count * page_size
+            if total_bytes < 1024:
+                total_size = f"~{total_bytes} B"
+            elif total_bytes < 1024 * 1024:
+                total_size = f"~{total_bytes // 1024} KB"
+            else:
+                total_size = f"~{total_bytes // (1024 * 1024)} MB"
+            avg_row_size = f"~{total_bytes // max(row_count, 1)} B" if row_count > 0 else "~0 B"
+        except Exception:
+            total_size = "unknown"
+            avg_row_size = "unknown"
+
+        null_counts: dict[str, int] = {}
+        try:
+            col_rows = conn.execute(f'PRAGMA table_info("{table}")').fetchall()
+            for c in col_rows:
+                try:
+                    cnt = conn.execute(
+                        f'SELECT COUNT(*) FROM "{table}" WHERE "{c["name"]}" IS NULL'
+                    ).fetchone()[0]
+                    if cnt > 0:
+                        null_counts[c["name"]] = cnt
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+        return TableStats(
+            row_count=row_count,
+            avg_row_size=avg_row_size,
+            total_size=total_size,
+            null_counts=null_counts,
+        )
+
+    def is_read_only(self) -> bool:
+        return self._read_only