diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index a7bc97a..ead1292 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -36,11 +36,24 @@ jobs: - uses: Swatinem/rust-cache@v2 + # The extension directory is DuckDB-version-specific (v//), so the + # key is derived from Cargo.lock, which pins the duckdb crate. A hardcoded version in + # the key would keep hitting after a crate bump and, because an exact hit skips the + # save, the new extensions would be re-downloaded on every run and never cached. - name: Cache DuckDB Extensions + id: duckdb-extensions uses: actions/cache@v4 with: path: ~/.duckdb/extensions - key: duckdb-ext-${{ matrix.os }}-v1.5.5 + key: duckdb-ext-${{ matrix.os }}-${{ hashFiles('Cargo.lock') }} + + # Populates the cache on a miss: opening a DuckDB store installs vss and fts when + # they are absent, and the post-job save then has something to store. Runs the + # DuckDB tests only, so a download failure is reported here rather than as a wall of + # unrelated-looking test failures. + - name: Install DuckDB Extensions + if: steps.duckdb-extensions.outputs.cache-hit != 'true' + run: cargo test --all duckdb - name: Test run: cargo test --all diff --git a/src/rag/providers/duckdb.rs b/src/rag/providers/duckdb.rs index 5c877fe..64b3eee 100644 --- a/src/rag/providers/duckdb.rs +++ b/src/rag/providers/duckdb.rs @@ -44,11 +44,12 @@ impl DuckDbProvider { // "Setting with name ... is not in the catalog, but it exists in the vss // extension". Omitting it entirely makes CREATE INDEX ... USING HNSW on a // file-backed database fail with "HNSW index persistence is not yet supported - // by default". LOAD vss -> LOAD fts -> SET -> CREATE INDEX. + // by default". ensure vss (installing it if missing) -> ensure fts -> SET -> + // CREATE INDEX. + Self::ensure_extension(&conn, "vss")?; + Self::ensure_extension(&conn, "fts")?; conn.execute_batch(&format!( - "LOAD vss; - LOAD fts; - SET hnsw_enable_experimental_persistence = true; + "SET hnsw_enable_experimental_persistence = true; CREATE TABLE IF NOT EXISTS vectors ( doc_id UBIGINT PRIMARY KEY, embedding FLOAT[{dim}] @@ -73,6 +74,29 @@ impl DuckDbProvider { }) } + /// Make a DuckDB extension available on `conn`, installing it if this machine does + /// not have it yet. `LOAD` is attempted first so an extension that is already + /// installed costs nothing and never touches the network; `INSTALL` is only reached + /// once, on a machine seeing the extension for the first time. + fn ensure_extension(conn: &Connection, name: &str) -> Result<()> { + if conn.execute_batch(&format!("LOAD {name};")).is_ok() { + return Ok(()); + } + conn.execute_batch(&format!("INSTALL {name};")) + .with_context(|| { + format!( + "Failed to install the DuckDB `{name}` extension. The duckdb RAG driver \ + needs it, and downloading it needs network access the first time. If this \ + machine is offline, connect once and retry, or run `INSTALL {name};` \ + yourself from a DuckDB shell." + ) + })?; + conn.execute_batch(&format!("LOAD {name};")) + .with_context(|| { + format!("Failed to load the DuckDB `{name}` extension after installing it.") + }) + } + /// Read all `(doc_id, embedding)` pairs so `create()` can hydrate `data.vectors` /// from disk. This is what makes the next incremental sync non-destructive, and is /// mandatory rather than an optimization.