Skip to content

sqlite: support `db.loadExtension`

Rodrigo Muino Tomonari requested to merge github/fork/himself65/himself65/20240717/load_extension into main

Closes https://github.com/nodejs/node/issues/53898

This PR will implement loadExtension API to align with other SQLite3 JS library like better-sqlite3, node-sqlite3, jsr:@db/sqlite, bun:sqlite

We used https://github.com/nodejs/sqlite-extension-sample to generate the sample SQLite extension into our fixtures

Example Code:

import { load } from 'sqlite-vec'
import { DatabaseSync } from 'node:sqlite'

const db = new DatabaseSync(':memory:', {
  allowLoadExtension: true
})
load(db)  // db.loadExtension() support
const { sqlite_version, vec_version } = db.prepare(
  'select sqlite_version() as sqlite_version, vec_version() as vec_version;'
).get()
console.log(`SQLite version: ${sqlite_version}, vector version ${vec_version}`)

const items = [
  [1, [0.1, 0.1, 0.1, 0.1]],
  [2, [0.2, 0.2, 0.2, 0.2]],
  [3, [0.3, 0.3, 0.3, 0.3]],
  [4, [0.4, 0.4, 0.4, 0.4]],
  [5, [0.5, 0.5, 0.5, 0.5]]
]
const query = [0.3, 0.3, 0.3, 0.3]

db.exec('CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])')

const insertStmt = db.prepare(
  'INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)'
)

for (const [id, vector] of items) {
  const rowId = BigInt(id)
  const embedding = new Uint8Array(Float32Array.from(vector).buffer)
  insertStmt.run(rowId, embedding)
}

const rows = db.prepare(
  `
  SELECT
    rowid,
    distance
  FROM vec_items
  WHERE embedding MATCH ?
  ORDER BY distance
  LIMIT 3
`
).all(new Uint8Array(Float32Array.from(query).buffer))

console.log(rows)
> sqlite-vec /Users/himself65/Code/node/out/Debug/node --experimental-sqlite ./index.mjs
(node:72169) ExperimentalWarning: SQLite is an experimental feature and might change at any time
(Use `node --trace-warnings ...` to show where the warning was created)
SQLite version: 3.46.0, vector version v0.1.1
[
  { rowid: 3, distance: 0 },
  { rowid: 4, distance: 0.19999998807907104 },
  { rowid: 2, distance: 0.20000001788139343 }
]

Merge request reports