poimen-memory/crates/mem-store/src/schema.rs

/// Database schema initialization.
use sqlx::PgPool;
use anyhow::Result;

/// Initialize database schema. Idempotent — safe to call multiple times.
pub async fn init_schema(pool: &PgPool) -> Result<()> {
    // Enable pgvector
    sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
        .execute(pool)
        .await?;

    // Event log — source of truth
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS events (
          id BIGSERIAL PRIMARY KEY,
          project VARCHAR NOT NULL,
          query_id VARCHAR NOT NULL,
          run_id VARCHAR NOT NULL,
          turn INT NOT NULL,
          event_type VARCHAR NOT NULL,
          created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          data JSONB NOT NULL,
          UNIQUE(project, query_id, run_id, turn)
        )
        "#,
    )
    .execute(pool)
    .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_events_project_query ON events(project, query_id)")
        .execute(pool)
        .await?;
    sqlx::query("CREATE INDEX IF NOT EXISTS idx_events_run ON events(run_id)")
        .execute(pool)
        .await?;
    sqlx::query("CREATE INDEX IF NOT EXISTS idx_events_type ON events(event_type)")
        .execute(pool)
        .await?;

    // L0: Evidence chunks
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS chunks_l0 (
          id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
          project VARCHAR NOT NULL,
          query_id VARCHAR NOT NULL,
          source VARCHAR NOT NULL,
          content TEXT NOT NULL,
          tokens INT NOT NULL,
          created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        "#,
    )
    .execute(pool)
    .await?;

    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_chunks_l0_project_query ON chunks_l0(project, query_id)",
    )
    .execute(pool)
    .await?;

    // L1: Per-query memories
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS memories_l1 (
          id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
          project VARCHAR NOT NULL,
          query_id VARCHAR NOT NULL,
          content TEXT NOT NULL,
          tokens INT NOT NULL,
          embedding vector(768),
          chunks_seen INT NOT NULL,
          chunks_used INT NOT NULL,
          updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          run_id VARCHAR NOT NULL,
          UNIQUE(project, query_id)
        )
        "#,
    )
    .execute(pool)
    .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_memories_l1_project ON memories_l1(project)")
        .execute(pool)
        .await?;
    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_memories_l1_embedding ON memories_l1 USING ivfflat (embedding vector_cosine_ops)",
    )
    .execute(pool)
    .await?;

    // L1 -> L0 provenance
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS l1_l0_edges (
          l1_id UUID REFERENCES memories_l1(id) ON DELETE CASCADE,
          l0_id UUID REFERENCES chunks_l0(id) ON DELETE CASCADE,
          PRIMARY KEY (l1_id, l0_id)
        )
        "#,
    )
    .execute(pool)
    .await?;

    // L2: Project synthesis
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS memories_l2 (
          id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
          project VARCHAR NOT NULL UNIQUE,
          content TEXT NOT NULL,
          tokens INT NOT NULL,
          embedding vector(768),
          l1_count INT NOT NULL,
          updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          run_id VARCHAR NOT NULL
        )
        "#,
    )
    .execute(pool)
    .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_memories_l2_project ON memories_l2(project)")
        .execute(pool)
        .await?;
    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_memories_l2_embedding ON memories_l2 USING ivfflat (embedding vector_cosine_ops)",
    )
    .execute(pool)
    .await?;

    // L2 -> L1 provenance
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS l2_l1_edges (
          l2_id UUID REFERENCES memories_l2(id) ON DELETE CASCADE,
          l1_id UUID REFERENCES memories_l1(id) ON DELETE CASCADE,
          PRIMARY KEY (l2_id, l1_id)
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Reference corpus
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS reference_corpus (
          id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
          project VARCHAR NOT NULL,
          name VARCHAR NOT NULL,
          content TEXT NOT NULL,
          embedding vector(768),
          created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          UNIQUE(project, name)
        )
        "#,
    )
    .execute(pool)
    .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_corpus_project ON reference_corpus(project)")
        .execute(pool)
        .await?;
    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_corpus_embedding ON reference_corpus USING ivfflat (embedding vector_cosine_ops)",
    )
    .execute(pool)
    .await?;

    // Ingest jobs
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS ingest_jobs (
          id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
          project VARCHAR NOT NULL,
          ingest_id VARCHAR NOT NULL UNIQUE,
          status VARCHAR NOT NULL DEFAULT 'pending',
          created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          started_at TIMESTAMP,
          completed_at TIMESTAMP,
          error TEXT
        )
        "#,
    )
    .execute(pool)
    .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_ingest_jobs_project ON ingest_jobs(project)")
        .execute(pool)
        .await?;
    sqlx::query("CREATE INDEX IF NOT EXISTS idx_ingest_jobs_status ON ingest_jobs(status)")
        .execute(pool)
        .await?;

    // Skills
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS skills (
          id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
          project VARCHAR NOT NULL,
          name VARCHAR NOT NULL,
          description TEXT NOT NULL,
          when_to_use TEXT,
          examples TEXT,
          l1_source UUID NOT NULL REFERENCES memories_l1(id) ON DELETE CASCADE,
          created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          UNIQUE(project, name)
        )
        "#,
    )
    .execute(pool)
    .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_skills_project ON skills(project)")
        .execute(pool)
        .await?;

    tracing::info!("Database schema initialized");
    Ok(())
}