context_harness/
migrate.rs

1//! Database schema migrations.
2//!
3//! Creates all required tables and ensures idempotent execution.
4//! Designed to be run via `ctx init`.
5//!
6//! # Schema
7//!
8//! ```text
9//! ┌──────────────┐     ┌──────────────┐     ┌──────────────┐
10//! │  documents   │──┐  │   chunks     │──┐  │  embeddings  │
11//! │              │  │  │              │  │  │              │
12//! │ id (PK)      │  │  │ id (PK)      │  │  │ chunk_id(PK) │
13//! │ source       │  └──│ document_id  │  └──│ model        │
14//! │ source_id    │     │ chunk_index  │     │ dims         │
15//! │ source_url   │     │ text         │     │ created_at   │
16//! │ title        │     │ hash         │     │ hash         │
17//! │ author       │     └──────────────┘     └──────────────┘
18//! │ created_at   │
19//! │ updated_at   │     ┌──────────────┐     ┌──────────────┐
20//! │ content_type │     │  chunks_fts  │     │chunk_vectors │
21//! │ body         │     │  (FTS5)      │     │              │
22//! │ metadata_json│     │ chunk_id     │     │ chunk_id(PK) │
23//! │ raw_json     │     │ document_id  │     │ document_id  │
24//! │ dedup_hash   │     │ text         │     │ embedding    │
25//! └──────────────┘     └──────────────┘     └──────────────┘
26//!
27//! ┌──────────────┐
28//! │ checkpoints  │
29//! │              │
30//! │ source (PK)  │
31//! │ cursor       │
32//! │ updated_at   │
33//! └──────────────┘
34//! ```
35//!
36//! # Tables
37//!
38//! | Table | Purpose |
39//! |-------|---------|
40//! | `documents` | Normalized document metadata and body text |
41//! | `chunks` | Text segments with content hashes |
42//! | `checkpoints` | Incremental sync cursors per connector |
43//! | `chunks_fts` | FTS5 full-text index over chunk text (BM25) |
44//! | `embeddings` | Embedding metadata (model, dims, hash) |
45//! | `chunk_vectors` | Embedding vectors stored as BLOBs |
46//!
47//! # Indexes
48//!
49//! - `idx_chunks_document_id` — fast chunk lookup by document
50//! - `idx_documents_source` — fast document filtering by connector
51//! - `idx_documents_updated_at` — efficient date range queries
52//! - `idx_chunk_vectors_document_id` — fast vector lookup by document
53//!
54//! # Idempotency
55//!
56//! All operations use `CREATE TABLE IF NOT EXISTS` or check for existing
57//! objects before creation. Running `ctx init` multiple times is safe.
58
59use anyhow::Result;
60
61use crate::config::Config;
62use crate::db;
63
64/// Run all database migrations.
65///
66/// Creates all tables, indexes, and virtual tables required by Context
67/// Harness. Safe to call multiple times — all operations are idempotent.
68///
69/// # Tables Created
70///
71/// - `documents` — normalized document storage
72/// - `chunks` — text segments with content hashes
73/// - `checkpoints` — incremental sync cursors
74/// - `chunks_fts` — FTS5 full-text search index
75/// - `embeddings` — embedding metadata (model, dims, staleness hash)
76/// - `chunk_vectors` — embedding vector BLOBs
77///
78/// # Errors
79///
80/// Returns an error if the database connection fails or any SQL statement
81/// cannot be executed.
82pub async fn run_migrations(config: &Config) -> Result<()> {
83    let pool = db::connect(config).await?;
84
85    // Create documents table
86    sqlx::query(
87        r#"
88        CREATE TABLE IF NOT EXISTS documents (
89            id TEXT PRIMARY KEY,
90            source TEXT NOT NULL,
91            source_id TEXT NOT NULL,
92            source_url TEXT,
93            title TEXT,
94            author TEXT,
95            created_at INTEGER NOT NULL,
96            updated_at INTEGER NOT NULL,
97            content_type TEXT NOT NULL DEFAULT 'text/plain',
98            body TEXT NOT NULL,
99            metadata_json TEXT NOT NULL DEFAULT '{}',
100            raw_json TEXT,
101            dedup_hash TEXT NOT NULL,
102            UNIQUE(source, source_id)
103        )
104        "#,
105    )
106    .execute(&pool)
107    .await?;
108
109    // Create chunks table
110    sqlx::query(
111        r#"
112        CREATE TABLE IF NOT EXISTS chunks (
113            id TEXT PRIMARY KEY,
114            document_id TEXT NOT NULL,
115            chunk_index INTEGER NOT NULL,
116            text TEXT NOT NULL,
117            hash TEXT NOT NULL,
118            UNIQUE(document_id, chunk_index),
119            FOREIGN KEY (document_id) REFERENCES documents(id)
120        )
121        "#,
122    )
123    .execute(&pool)
124    .await?;
125
126    // Create checkpoints table
127    sqlx::query(
128        r#"
129        CREATE TABLE IF NOT EXISTS checkpoints (
130            source TEXT PRIMARY KEY,
131            cursor TEXT NOT NULL,
132            updated_at INTEGER NOT NULL
133        )
134        "#,
135    )
136    .execute(&pool)
137    .await?;
138
139    // Create FTS5 virtual table over chunks (not idempotent natively, check first)
140    let fts_exists: bool = sqlx::query_scalar(
141        "SELECT COUNT(*) > 0 FROM sqlite_master WHERE type='table' AND name='chunks_fts'",
142    )
143    .fetch_one(&pool)
144    .await?;
145
146    if !fts_exists {
147        sqlx::query(
148            r#"
149            CREATE VIRTUAL TABLE chunks_fts USING fts5(
150                chunk_id UNINDEXED,
151                document_id UNINDEXED,
152                text
153            )
154            "#,
155        )
156        .execute(&pool)
157        .await?;
158    }
159
160    // Embeddings metadata table
161    sqlx::query(
162        r#"
163        CREATE TABLE IF NOT EXISTS embeddings (
164            chunk_id TEXT PRIMARY KEY,
165            model TEXT NOT NULL,
166            dims INTEGER NOT NULL,
167            created_at INTEGER NOT NULL,
168            hash TEXT NOT NULL,
169            FOREIGN KEY (chunk_id) REFERENCES chunks(id)
170        )
171        "#,
172    )
173    .execute(&pool)
174    .await?;
175
176    // Chunk vectors table (stores embedding BLOBs)
177    sqlx::query(
178        r#"
179        CREATE TABLE IF NOT EXISTS chunk_vectors (
180            chunk_id TEXT PRIMARY KEY,
181            document_id TEXT NOT NULL,
182            embedding BLOB NOT NULL,
183            FOREIGN KEY (chunk_id) REFERENCES chunks(id),
184            FOREIGN KEY (document_id) REFERENCES documents(id)
185        )
186        "#,
187    )
188    .execute(&pool)
189    .await?;
190
191    // Create indexes for common query patterns
192    sqlx::query("CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)")
193        .execute(&pool)
194        .await?;
195    sqlx::query("CREATE INDEX IF NOT EXISTS idx_documents_source ON documents(source)")
196        .execute(&pool)
197        .await?;
198    sqlx::query(
199        "CREATE INDEX IF NOT EXISTS idx_documents_updated_at ON documents(updated_at DESC)",
200    )
201    .execute(&pool)
202    .await?;
203    sqlx::query(
204        "CREATE INDEX IF NOT EXISTS idx_chunk_vectors_document_id ON chunk_vectors(document_id)",
205    )
206    .execute(&pool)
207    .await?;
208
209    // Install runtime tables atomically and idempotently in the existing DB.
210    let mut tx = pool.begin().await?;
211    // This DDL contains only simple statements (no triggers or embedded
212    // semicolons). Execute separately to retain a Send migration future.
213    for statement in include_str!("agent_store/schema.sql").split(';') {
214        if !statement.trim().is_empty() {
215            sqlx::query(statement).execute(&mut *tx).await?;
216        }
217    }
218    let columns: Vec<String> =
219        sqlx::query_scalar("SELECT name FROM pragma_table_info('agent_runs')")
220            .fetch_all(&mut *tx)
221            .await?;
222    for (name, definition) in [
223        (
224            "lifecycle",
225            "TEXT NOT NULL DEFAULT 'active' CHECK (lifecycle IN ('active', 'suspended', 'terminal'))",
226        ),
227        (
228            "outcome",
229            "TEXT CHECK (outcome IS NULL OR outcome IN ('completed', 'blocked', 'needs_user_input', 'failed', 'limit_exceeded', 'cancelled'))",
230        ),
231        ("reason_code", "TEXT"),
232        (
233            "reason_detail",
234            "TEXT CHECK (reason_detail IS NULL OR json_valid(reason_detail))",
235        ),
236        ("budgets", "TEXT NOT NULL DEFAULT '{}' CHECK (json_valid(budgets))"),
237        (
238            "usage",
239            "TEXT NOT NULL DEFAULT '{\"model_turns\":0,\"input_tokens\":null,\"output_tokens\":null,\"total_tokens\":null,\"responses_with_usage\":0,\"responses_without_usage\":0,\"tool_calls\":0,\"token_accounting\":\"unavailable\"}' CHECK (json_valid(usage))",
240        ),
241        (
242            "suspended_ms",
243            "INTEGER NOT NULL DEFAULT 0 CHECK (suspended_ms >= 0)",
244        ),
245    ] {
246        if !columns.iter().any(|column| column == name) {
247            sqlx::query(&format!("ALTER TABLE agent_runs ADD COLUMN {name} {definition}"))
248                .execute(&mut *tx)
249                .await?;
250        }
251    }
252    let tool_columns: Vec<String> =
253        sqlx::query_scalar("SELECT name FROM pragma_table_info('tool_invocations')")
254            .fetch_all(&mut *tx)
255            .await?;
256    for (name, definition) in [
257        ("result_content", "TEXT"),
258        (
259            "outcome_class",
260            "TEXT CHECK (outcome_class IS NULL OR outcome_class IN ('completed', 'denied', 'recoverable_error', 'terminal_error', 'uncertain_side_effect', 'cancelled'))",
261        ),
262        ("error_code", "TEXT"),
263        ("repetition_key", "TEXT"),
264        (
265            "consecutive_count",
266            "INTEGER CHECK (consecutive_count IS NULL OR consecutive_count > 0)",
267        ),
268    ] {
269        if !tool_columns.iter().any(|column| column == name) {
270            sqlx::query(&format!(
271                "ALTER TABLE tool_invocations ADD COLUMN {name} {definition}"
272            ))
273            .execute(&mut *tx)
274            .await?;
275        }
276    }
277    sqlx::query(
278        "UPDATE agent_runs SET lifecycle = CASE status WHEN 'running' THEN 'active' ELSE 'terminal' END, outcome = CASE status WHEN 'completed' THEN 'completed' WHEN 'failed' THEN 'failed' WHEN 'cancelled' THEN 'cancelled' END, reason_code = CASE status WHEN 'completed' THEN 'legacy_completed' WHEN 'failed' THEN 'legacy_failure' WHEN 'cancelled' THEN 'legacy_cancelled' END, reason_detail = CASE WHEN status = 'running' THEN NULL ELSE '{}' END WHERE reason_code IS NULL",
279    )
280    .execute(&mut *tx)
281    .await?;
282    tx.commit().await?;
283
284    pool.close().await;
285    Ok(())
286}
context_harness/migrate.rs

context_harness/
migrate.rs
