import asyncio import hashlib import uuid from contextlib import asynccontextmanager from pathlib import Path from sqlalchemy import create_engine, event, text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase, sessionmaker from app.config import settings # Async engine for aiosqlite DATABASE_URL = settings.resolved_database_url # Sync engine for migrations and initial setup SYNC_DATABASE_URL = DATABASE_URL.replace("sqlite+aiosqlite:///", "sqlite:///") class Base(DeclarativeBase): pass # Async session factory async_engine = create_async_engine(DATABASE_URL, echo=False) AsyncSessionLocal = async_sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False) # Sync engine for migrations sync_engine = create_engine(SYNC_DATABASE_URL, echo=False) SyncSessionLocal = sessionmaker(sync_engine) async def get_db(): async with AsyncSessionLocal() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() async def get_db_simple(): """Bare async generator for FastAPI dependency injection.""" async with AsyncSessionLocal() as session: yield session async def init_db(): """Initialize database with all tables, views, FTS5, and triggers.""" # Create data directory Path("./data").mkdir(exist_ok=True) async with async_engine.begin() as conn: # Create all tables via SQL (not ORM) to handle SQLite-specific features await conn.run_sync(_create_schema) def _create_schema(sync_conn): """Create all tables, views, FTS5 tables, and triggers synchronously.""" # Agents table sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS agents ( id TEXT PRIMARY KEY, username TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, role TEXT NOT NULL DEFAULT 'agent' CHECK (role IN ('agent', 'admin')), is_deleted INTEGER NOT NULL DEFAULT 0, deleted_at TIMESTAMP NULL, deleted_by TEXT NULL, created_at TIMESTAMP NOT NULL DEFAULT (datetime('now')), updated_at TIMESTAMP NOT NULL DEFAULT (datetime('now')) ) """)) # Projects table sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS projects ( id TEXT PRIMARY KEY, name TEXT NOT NULL, description TEXT, agent_id TEXT NOT NULL REFERENCES agents(id), is_deleted INTEGER NOT NULL DEFAULT 0, deleted_at TIMESTAMP NULL, deleted_by TEXT NULL, created_at TIMESTAMP NOT NULL DEFAULT (datetime('now')), updated_at TIMESTAMP NOT NULL DEFAULT (datetime('now')) ) """)) # Folders table sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS folders ( id TEXT PRIMARY KEY, name TEXT NOT NULL, project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE, parent_id TEXT REFERENCES folders(id) ON DELETE CASCADE, path TEXT NOT NULL, is_deleted INTEGER NOT NULL DEFAULT 0, deleted_at TIMESTAMP NULL, deleted_by TEXT NULL, created_at TIMESTAMP NOT NULL DEFAULT (datetime('now')), updated_at TIMESTAMP NOT NULL DEFAULT (datetime('now')) ) """)) # Documents table sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS documents ( id TEXT PRIMARY KEY, title TEXT NOT NULL, content TEXT NOT NULL DEFAULT '', project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE, folder_id TEXT REFERENCES folders(id) ON DELETE SET NULL, path TEXT NOT NULL, is_deleted INTEGER NOT NULL DEFAULT 0, deleted_at TIMESTAMP NULL, deleted_by TEXT NULL, created_at TIMESTAMP NOT NULL DEFAULT (datetime('now')), updated_at TIMESTAMP NOT NULL DEFAULT (datetime('now')) ) """)) # Tags table sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS tags ( id TEXT PRIMARY KEY, name TEXT UNIQUE NOT NULL, color TEXT NOT NULL DEFAULT '#6366f1', is_deleted INTEGER NOT NULL DEFAULT 0, deleted_at TIMESTAMP NULL, deleted_by TEXT NULL, created_at TIMESTAMP NOT NULL DEFAULT (datetime('now')) ) """)) # Document tags junction sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS document_tags ( document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE, tag_id TEXT NOT NULL REFERENCES tags(id) ON DELETE CASCADE, PRIMARY KEY (document_id, tag_id) ) """)) # Refresh tokens table sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS refresh_tokens ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL REFERENCES agents(id), token_hash TEXT NOT NULL UNIQUE, token_family_id TEXT NOT NULL, token_version INTEGER NOT NULL, user_agent TEXT, ip_address TEXT, created_at TIMESTAMP NOT NULL DEFAULT (datetime('now')), expires_at TIMESTAMP NOT NULL, revoked_at TIMESTAMP NULL, is_global_logout INTEGER NOT NULL DEFAULT 0 ) """)) # JWT blocklist table sync_conn.execute(text(""" CREATE TABLE IF NOT EXISTS jwt_blocklist ( token_id TEXT PRIMARY KEY, revoked_at TIMESTAMP NOT NULL DEFAULT (datetime('now')), expires_at TIMESTAMP NOT NULL ) """)) # Indexes sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_projects_agent ON projects(agent_id)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_folders_project ON folders(project_id)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_folders_parent ON folders(parent_id)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_documents_project ON documents(project_id)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_documents_folder ON documents(folder_id)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_document_tags_doc ON document_tags(document_id)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_document_tags_tag ON document_tags(tag_id)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_refresh_tokens_hash ON refresh_tokens(token_hash)")) sync_conn.execute(text("CREATE INDEX IF NOT EXISTS idx_refresh_tokens_user_family ON refresh_tokens(user_id, token_family_id)")) # --- FTS5 virtual table --- sync_conn.execute(text(""" CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5( document_id, title, content, path, tokenize='unicode61 remove_diacritics 1' ) """)) # --- Active views for soft deletes --- sync_conn.execute(text(""" CREATE VIEW IF NOT EXISTS active_agents AS SELECT * FROM agents WHERE is_deleted = 0 """)) sync_conn.execute(text(""" CREATE VIEW IF NOT EXISTS active_projects AS SELECT * FROM projects WHERE is_deleted = 0 """)) sync_conn.execute(text(""" CREATE VIEW IF NOT EXISTS active_folders AS SELECT * FROM folders WHERE is_deleted = 0 """)) sync_conn.execute(text(""" CREATE VIEW IF NOT EXISTS active_documents AS SELECT * FROM documents WHERE is_deleted = 0 """)) sync_conn.execute(text(""" CREATE VIEW IF NOT EXISTS active_tags AS SELECT * FROM tags WHERE is_deleted = 0 """)) # --- FTS5 Sync Triggers --- # Insert trigger sync_conn.execute(text(""" CREATE TRIGGER IF NOT EXISTS documents_fts_ai AFTER INSERT ON documents BEGIN INSERT INTO documents_fts(document_id, title, content, path) VALUES (new.id, new.title, new.content, new.path); END """)) # Update trigger (delete old + insert new) sync_conn.execute(text(""" CREATE TRIGGER IF NOT EXISTS documents_fts_au AFTER UPDATE ON documents BEGIN DELETE FROM documents_fts WHERE document_id = old.id; INSERT INTO documents_fts(document_id, title, content, path) VALUES (new.id, new.title, new.content, new.path); END """)) # Soft-delete trigger (when is_deleted becomes TRUE) sync_conn.execute(text(""" CREATE TRIGGER IF NOT EXISTS documents_fts_ad AFTER UPDATE ON documents WHEN new.is_deleted = 1 AND old.is_deleted = 0 BEGIN DELETE FROM documents_fts WHERE document_id = old.id; END """)) sync_conn.commit()