docsif 0.1.0__py3-none-any.whl

sif/database/database.py

@@ -0,0 +1,136 @@
+"""Database connection and management for SIF."""
+
+from __future__ import annotations
+
+import sqlite3
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+
+import sqlite_vec
+
+from sif.database.schema import SchemaManager
+
+
+class Database:
+    """Main database class for SIF."""
+
+    def __init__(self, db_path: str | Path) -> None:
+        self.db_path = Path(db_path).expanduser().resolve()
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._connection: sqlite3.Connection | None = None
+        self._schema: SchemaManager | None = None
+
+    @property
+    def connection(self) -> sqlite3.Connection:
+        """Get or create database connection."""
+        if self._connection is None:
+            self._connection = sqlite3.connect(
+                str(self.db_path),
+                detect_types=sqlite3.PARSE_DECLTYPES,
+                check_same_thread=False,
+            )
+            self._connection.row_factory = sqlite3.Row
+            self._enable_extensions()
+        return self._connection
+
+    def _enable_extensions(self) -> None:
+        """Enable SQLite extensions."""
+        # Enable FTS5
+        self._connection.execute("PRAGMA foreign_keys = ON")
+
+        # Load sqlite-vec extension
+        try:
+            self._connection.enable_load_extension(True)
+            sqlite_vec.load(self._connection)
+            self._connection.enable_load_extension(False)
+        except Exception:
+            pass  # sqlite-vec is optional
+
+    def init_schema(self) -> None:
+        """Initialize database schema."""
+        from sif.config.settings import get_settings  # noqa: PLC0415
+
+        settings = get_settings()
+        self._schema = SchemaManager(self.connection, embedding_dim=settings.embedding_dim)
+        self._schema.create_all()
+
+    @contextmanager
+    def transaction(self) -> Generator[sqlite3.Connection, None, None]:
+        """Execute operations in a transaction."""
+        conn = self.connection
+        try:
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+
+    def close(self) -> None:
+        """Close database connection."""
+        if self._connection:
+            self._connection.close()
+            self._connection = None
+
+    def __enter__(self) -> Database:
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Context manager exit."""
+        self.close()
+
+    def get_stats(self) -> dict:
+        """Get database statistics."""
+        if self._schema is None:
+            from sif.config.settings import get_settings  # noqa: PLC0415
+
+            settings = get_settings()
+            self._schema = SchemaManager(self.connection, embedding_dim=settings.embedding_dim)
+        return self._schema.get_stats()
+
+    def reset(self) -> None:
+        """Reset database (drop all tables)."""
+        if self._schema is None:
+            from sif.config.settings import get_settings  # noqa: PLC0415
+
+            settings = get_settings()
+            self._schema = SchemaManager(self.connection, embedding_dim=settings.embedding_dim)
+        self._schema.drop_all()
+        self._schema.create_all()
+
+
+class DatabaseConnection:
+    """Lightweight database connection wrapper."""
+
+    def __init__(self, db_path: str | Path) -> None:
+        self.db_path = Path(db_path).expanduser().resolve()
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
+    @contextmanager
+    def connect(self) -> Generator[sqlite3.Connection, None, None]:
+        """Create a connection context."""
+        conn = sqlite3.connect(
+            str(self.db_path),
+            detect_types=sqlite3.PARSE_DECLTYPES,
+            check_same_thread=False,
+        )
+        conn.row_factory = sqlite3.Row
+        conn.execute("PRAGMA foreign_keys = ON")
+
+        # Load sqlite-vec extension
+        try:
+            conn.enable_load_extension(True)
+            sqlite_vec.load(conn)
+            conn.enable_load_extension(False)
+        except Exception:
+            pass  # sqlite-vec is optional
+
+        try:
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+        finally:
+            conn.close()

sif/database/migrations.py

@@ -0,0 +1,187 @@
+"""Database migration management."""
+
+from sif.config.constants import CURRENT_SCHEMA_VERSION
+from sif.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+# Schema definition
+SCHEMA_SQL = """
+-- Collections table
+CREATE TABLE IF NOT EXISTS collections (
+    id TEXT PRIMARY KEY,
+    name TEXT UNIQUE NOT NULL,
+    description TEXT,
+    paths TEXT NOT NULL DEFAULT '[]',
+    document_count INTEGER NOT NULL DEFAULT 0,
+    chunk_count INTEGER NOT NULL DEFAULT 0,
+    metadata TEXT NOT NULL DEFAULT '{}',
+    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    last_indexed_at DATETIME
+);
+
+-- Documents table
+CREATE TABLE IF NOT EXISTS documents (
+    id TEXT PRIMARY KEY,
+    collection_id TEXT NOT NULL,
+    path TEXT NOT NULL,
+    content TEXT NOT NULL,
+    checksum TEXT NOT NULL,
+    file_size INTEGER NOT NULL,
+    metadata TEXT NOT NULL DEFAULT '{}',
+    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    indexed_at DATETIME,
+    FOREIGN KEY (collection_id) REFERENCES collections(id) ON DELETE CASCADE,
+    UNIQUE(collection_id, path)
+);
+
+-- Chunks table
+CREATE TABLE IF NOT EXISTS chunks (
+    id TEXT PRIMARY KEY,
+    document_id TEXT NOT NULL,
+    content TEXT NOT NULL,
+    start_line INTEGER NOT NULL,
+    end_line INTEGER NOT NULL,
+    token_count INTEGER NOT NULL,
+    embedding_id TEXT,
+    FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE
+);
+
+-- Contexts table
+CREATE TABLE IF NOT EXISTS contexts (
+    id TEXT PRIMARY KEY,
+    content TEXT NOT NULL,
+    context_type TEXT NOT NULL,
+    target_id TEXT NOT NULL,
+    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    UNIQUE(context_type, target_id)
+);
+
+-- Schema version tracking
+CREATE TABLE IF NOT EXISTS schema_version (
+    version INTEGER PRIMARY KEY,
+    applied_at DATETIME DEFAULT CURRENT_TIMESTAMP
+);
+
+-- Indexes
+CREATE INDEX IF NOT EXISTS idx_documents_collection ON documents(collection_id);
+CREATE INDEX IF NOT EXISTS idx_documents_path ON documents(path);
+CREATE INDEX IF NOT EXISTS idx_chunks_document ON chunks(document_id);
+CREATE INDEX IF NOT EXISTS idx_contexts_type ON contexts(context_type);
+CREATE INDEX IF NOT EXISTS idx_contexts_target ON contexts(target_id);
+"""
+
+# FTS5 virtual table
+FTS_SQL = """
+CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
+    content,
+    content_rowid=rowid,
+    tokenize='porter'
+);
+
+-- Trigger to keep FTS index in sync
+CREATE TRIGGER IF NOT EXISTS documents_fts_insert
+AFTER INSERT ON documents
+BEGIN
+    INSERT INTO documents_fts(rowid, content) VALUES (new.rowid, new.content);
+END;
+
+CREATE TRIGGER IF NOT EXISTS documents_fts_update
+AFTER UPDATE ON documents
+BEGIN
+    UPDATE documents_fts SET content = new.content WHERE rowid = old.rowid;
+END;
+
+CREATE TRIGGER IF NOT EXISTS documents_fts_delete
+AFTER DELETE ON documents
+BEGIN
+    DELETE FROM documents_fts WHERE rowid = old.rowid;
+END;
+"""
+
+# Vector table (sqlite-vec)
+VECTOR_SQL_TEMPLATE = """
+CREATE VIRTUAL TABLE IF NOT EXISTS chunk_embeddings USING vec0(
+    embedding_id TEXT PRIMARY KEY,
+    embedding FLOAT[{dim}]
+);
+"""
+
+
+class MigrationManager:
+    """Manages database schema migrations."""
+
+    def __init__(self, connection: "DatabaseConnection") -> None:  # noqa: F821
+        self._connection = connection
+
+    def get_current_version(self) -> int:
+        """Get the current schema version."""
+        try:
+            row = self._connection.fetchone(
+                "SELECT version FROM schema_version ORDER BY version DESC LIMIT 1"
+            )
+            return row["version"] if row else 0
+        except Exception:
+            return 0
+
+    def migrate(self, embedding_dim: int = 384) -> None:
+        """Run all pending migrations."""
+        current_version = self.get_current_version()
+
+        if current_version >= CURRENT_SCHEMA_VERSION:
+            logger.info("Database schema is up to date")
+            return
+
+        logger.info(
+            f"Migrating database from version {current_version} to {CURRENT_SCHEMA_VERSION}"
+        )
+
+        # Apply schema
+        self._apply_schema(embedding_dim)
+
+        # Update version
+        self._connection.execute(
+            "INSERT INTO schema_version (version) VALUES (?)",
+            (CURRENT_SCHEMA_VERSION,),
+        )
+
+        logger.info("Database migration completed successfully")
+
+    def _apply_schema(self, embedding_dim: int) -> None:
+        """Apply the database schema."""
+        # Apply main schema
+        self._connection.execute(SCHEMA_SQL)
+
+        # Apply FTS schema
+        self._connection.execute(FTS_SQL)
+
+        # Apply vector schema
+        vector_sql = VECTOR_SQL_TEMPLATE.format(dim=embedding_dim)
+        self._connection.execute(vector_sql)
+
+    def reset(self) -> None:
+        """Reset the database (drop all tables)."""
+        logger.warning("Resetting database - all data will be lost")
+
+        tables = [
+            "documents_fts",
+            "chunk_embeddings",
+            "chunks",
+            "documents",
+            "contexts",
+            "collections",
+            "schema_version",
+        ]
+
+        for table in tables:
+            try:
+                self._connection.execute(f"DROP TABLE IF EXISTS {table}")
+            except Exception as e:
+                logger.warning(f"Error dropping table {table}: {e}")
+
+        logger.info("Database reset completed")