"""SQLite database schema and connection management.""" import sqlite3 from contextlib import contextmanager from pathlib import Path from typing import Generator import atocore.config as _config from atocore.observability.logger import get_logger log = get_logger("database") SCHEMA_SQL = """ CREATE TABLE IF NOT EXISTS source_documents ( id TEXT PRIMARY KEY, file_path TEXT UNIQUE NOT NULL, file_hash TEXT NOT NULL, title TEXT, doc_type TEXT DEFAULT 'markdown', tags TEXT DEFAULT '[]', ingested_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS source_chunks ( id TEXT PRIMARY KEY, document_id TEXT NOT NULL REFERENCES source_documents(id) ON DELETE CASCADE, chunk_index INTEGER NOT NULL, content TEXT NOT NULL, heading_path TEXT DEFAULT '', char_count INTEGER NOT NULL, metadata TEXT DEFAULT '{}', created_at DATETIME DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS memories ( id TEXT PRIMARY KEY, memory_type TEXT NOT NULL, content TEXT NOT NULL, project TEXT DEFAULT '', source_chunk_id TEXT REFERENCES source_chunks(id), confidence REAL DEFAULT 1.0, status TEXT DEFAULT 'active', created_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS projects ( id TEXT PRIMARY KEY, name TEXT UNIQUE NOT NULL, description TEXT DEFAULT '', status TEXT DEFAULT 'active', created_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS interactions ( id TEXT PRIMARY KEY, prompt TEXT NOT NULL, context_pack TEXT DEFAULT '{}', response_summary TEXT DEFAULT '', project_id TEXT REFERENCES projects(id), created_at DATETIME DEFAULT CURRENT_TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_chunks_document ON source_chunks(document_id); CREATE INDEX IF NOT EXISTS idx_memories_type ON memories(memory_type); CREATE INDEX IF NOT EXISTS idx_memories_project ON memories(project); CREATE INDEX IF NOT EXISTS idx_memories_status ON memories(status); CREATE INDEX IF NOT EXISTS idx_interactions_project ON interactions(project_id); """ def _ensure_data_dir() -> None: _config.ensure_runtime_dirs() def init_db() -> None: """Initialize the database with schema.""" _ensure_data_dir() with get_connection() as conn: conn.executescript(SCHEMA_SQL) _apply_migrations(conn) log.info("database_initialized", path=str(_config.settings.db_path)) def _apply_migrations(conn: sqlite3.Connection) -> None: """Apply lightweight schema migrations for existing local databases.""" if not _column_exists(conn, "memories", "project"): conn.execute("ALTER TABLE memories ADD COLUMN project TEXT DEFAULT ''") conn.execute("CREATE INDEX IF NOT EXISTS idx_memories_project ON memories(project)") def _column_exists(conn: sqlite3.Connection, table: str, column: str) -> bool: rows = conn.execute(f"PRAGMA table_info({table})").fetchall() return any(row["name"] == column for row in rows) @contextmanager def get_connection() -> Generator[sqlite3.Connection, None, None]: """Get a database connection with row factory.""" _ensure_data_dir() conn = sqlite3.connect(str(_config.settings.db_path)) conn.row_factory = sqlite3.Row conn.execute("PRAGMA foreign_keys = ON") try: yield conn conn.commit() except Exception: conn.rollback() raise finally: conn.close()