claudia-docs-api/app/database.py

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()