sqlite: support db.loadExtension

Author: himself65

doc/api/errors.md

@@ -2140,6 +2140,16 @@ An attempt was made to open an IPC communication channel with a synchronously
 forked Node.js process. See the documentation for the [`child_process`][] module
 for more information.
 
+<a id="ERR_LOAD_SQLITE_EXTENSION"></a>
+
+### `ERR_LOAD_SQLITE_EXTENSION`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+An error occurred while loading a SQLite extension.
+
 <a id="ERR_LOADER_CHAIN_INCOMPLETE"></a>
 
 ### `ERR_LOADER_CHAIN_INCOMPLETE`

doc/api/sqlite.md

@@ -107,6 +107,8 @@ added: v22.5.0
   * `open` {boolean} If `true`, the database is opened by the constructor. When
     this value is `false`, the database must be opened via the `open()` method.
     **Default:** `true`.
+  * `allowLoadExtension` {boolean} If `true`, the `loadExtension` SQL function
+    is enabled. **Default:** `false`.
 
 Constructs a new `DatabaseSync` instance.
 
@@ -119,6 +121,17 @@ added: v22.5.0
 Closes the database connection. An exception is thrown if the database is not
 open. This method is a wrapper around [`sqlite3_close_v2()`][].
 
+### `database.loadExtension(path)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+* `path` {string} The path to the shared library to load.
+
+Loads a shared library into the database connection. The `allowLoadExtension` option must be
+enabled when constructing the `DatabaseSync` instance.
+
 ### `database.exec(sql)`
 
 <!-- YAML

src/node_errors.h

@@ -91,6 +91,7 @@ void OOMErrorHandler(const char* location, const v8::OOMDetails& details);
   V(ERR_INVALID_THIS, TypeError)                                               \
   V(ERR_INVALID_URL, TypeError)                                                \
   V(ERR_INVALID_URL_SCHEME, TypeError)                                         \
+  V(ERR_LOAD_SQLITE_EXTENSION, Error)                                          \
   V(ERR_MEMORY_ALLOCATION_FAILED, Error)                                       \
   V(ERR_MESSAGE_TARGET_CONTEXT_UNAVAILABLE, Error)                             \
   V(ERR_MISSING_ARGS, TypeError)                                               \
@@ -186,6 +187,7 @@ ERRORS_WITH_CODE(V)
   V(ERR_INVALID_STATE, "Invalid state")                                        \
   V(ERR_INVALID_THIS, "Value of \"this\" is the wrong type")                   \
   V(ERR_INVALID_URL_SCHEME, "The URL must be of scheme file:")                 \
+  V(ERR_LOAD_SQLITE_EXTENSION, "Failed to load SQLite extension")              \
   V(ERR_MEMORY_ALLOCATION_FAILED, "Failed to allocate memory")                 \
   V(ERR_OSSL_EVP_INVALID_DIGEST, "Invalid digest used")                        \
   V(ERR_MESSAGE_TARGET_CONTEXT_UNAVAILABLE,                                    \

src/node_sqlite.cc

@@ -6,6 +6,7 @@
 #include "node.h"
 #include "node_errors.h"
 #include "node_mem-inl.h"
+#include "path.h"
 #include "sqlite3.h"
 #include "util-inl.h"
 
@@ -78,12 +79,14 @@ inline void THROW_ERR_SQLITE_ERROR(Isolate* isolate, sqlite3* db) {
 DatabaseSync::DatabaseSync(Environment* env,
                            Local<Object> object,
                            Local<String> location,
-                           bool open)
+                           bool open,
+                           bool allow_load_extension)
     : BaseObject(env, object) {
   MakeWeak();
   node::Utf8Value utf8_location(env->isolate(), location);
   location_ = utf8_location.ToString();
   connection_ = nullptr;
+  allow_load_extension_ = allow_load_extension;
 
   if (open) {
     Open();
@@ -109,6 +112,12 @@ bool DatabaseSync::Open() {
   int flags = SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE;
   int r = sqlite3_open_v2(location_.c_str(), &connection_, flags, nullptr);
   CHECK_ERROR_OR_THROW(env()->isolate(), connection_, r, SQLITE_OK, false);
+  if (allow_load_extension_) {
+    int load_extension_ret = sqlite3_db_config(
+        connection_, SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION, 1, nullptr);
+    CHECK_ERROR_OR_THROW(
+        env()->isolate(), connection_, load_extension_ret, SQLITE_OK, false);
+  }
   return true;
 }
 
@@ -127,6 +136,7 @@ void DatabaseSync::New(const FunctionCallbackInfo<Value>& args) {
   }
 
   bool open = true;
+  bool allow_load_extension = false;
 
   if (args.Length() > 1) {
     if (!args[1]->IsObject()) {
@@ -137,10 +147,17 @@ void DatabaseSync::New(const FunctionCallbackInfo<Value>& args) {
 
     Local<Object> options = args[1].As<Object>();
     Local<String> open_string = FIXED_ONE_BYTE_STRING(env->isolate(), "open");
+    Local<String> allow_load_extension_string =
+        FIXED_ONE_BYTE_STRING(env->isolate(), "allowLoadExtension");
     Local<Value> open_v;
+    Local<Value> allow_load_extension_v;
     if (!options->Get(env->context(), open_string).ToLocal(&open_v)) {
       return;
     }
+    if (!options->Get(env->context(), allow_load_extension_string)
+             .ToLocal(&allow_load_extension_v)) {
+      return;
+    }
     if (!open_v->IsUndefined()) {
       if (!open_v->IsBoolean()) {
         node::THROW_ERR_INVALID_ARG_TYPE(
@@ -149,9 +166,19 @@ void DatabaseSync::New(const FunctionCallbackInfo<Value>& args) {
       }
       open = open_v.As<Boolean>()->Value();
     }
+    if (!allow_load_extension_v->IsUndefined()) {
+      if (!allow_load_extension_v->IsBoolean()) {
+        node::THROW_ERR_INVALID_ARG_TYPE(
+            env->isolate(),
+            "The \"options.allowLoadExtension\" argument must be a boolean.");
+        return;
+      }
+      allow_load_extension = allow_load_extension_v.As<Boolean>()->Value();
+    }
   }
 
-  new DatabaseSync(env, args.This(), args[0].As<String>(), open);
+  new DatabaseSync(
+      env, args.This(), args[0].As<String>(), open, allow_load_extension);
 }
 
 void DatabaseSync::Open(const FunctionCallbackInfo<Value>& args) {
@@ -211,6 +238,35 @@ void DatabaseSync::Exec(const FunctionCallbackInfo<Value>& args) {
   CHECK_ERROR_OR_THROW(env->isolate(), db->connection_, r, SQLITE_OK, void());
 }
 
+void DatabaseSync::LoadExtension(const FunctionCallbackInfo<Value>& args) {
+  DatabaseSync* db;
+  ASSIGN_OR_RETURN_UNWRAP(&db, args.This());
+  Environment* env = Environment::GetCurrent(args);
+  THROW_AND_RETURN_ON_BAD_STATE(
+      env, db->connection_ == nullptr, "database is not open");
+  THROW_AND_RETURN_ON_BAD_STATE(
+      env, !db->allow_load_extension_, "load extension is not allowed");
+
+  if (!args[0]->IsString()) {
+    node::THROW_ERR_INVALID_ARG_TYPE(env->isolate(),
+                                     "The \"path\" argument must be a string.");
+    return;
+  }
+
+  auto isolate = env->isolate();
+
+  BufferValue path(isolate, args[0]);
+  CHECK_NOT_NULL(*path);
+  ToNamespacedPath(env, &path);
+  THROW_IF_INSUFFICIENT_PERMISSIONS(
+      env, permission::PermissionScope::kFileSystemRead, path.ToStringView());
+  char* errmsg = nullptr;
+  int r = sqlite3_load_extension(db->connection_, *path, nullptr, &errmsg);
+  if (r != SQLITE_OK) {
+    isolate->ThrowException(node::ERR_LOAD_SQLITE_EXTENSION(isolate, errmsg));
+  }
+}
+
 StatementSync::StatementSync(Environment* env,
                              Local<Object> object,
                              sqlite3* db,
@@ -668,6 +724,8 @@ static void Initialize(Local<Object> target,
   SetProtoMethod(isolate, db_tmpl, "close", DatabaseSync::Close);
   SetProtoMethod(isolate, db_tmpl, "prepare", DatabaseSync::Prepare);
   SetProtoMethod(isolate, db_tmpl, "exec", DatabaseSync::Exec);
+  SetProtoMethod(
+      isolate, db_tmpl, "loadExtension", DatabaseSync::LoadExtension);
   SetConstructorFunction(context, target, "DatabaseSync", db_tmpl);
   SetConstructorFunction(context,
                          target,

src/node_sqlite.h

@@ -18,13 +18,15 @@ class DatabaseSync : public BaseObject {
   DatabaseSync(Environment* env,
                v8::Local<v8::Object> object,
                v8::Local<v8::String> location,
-               bool open);
+               bool open,
+               bool allow_load_extension);
   void MemoryInfo(MemoryTracker* tracker) const override;
   static void New(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void Open(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void Close(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void Prepare(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void Exec(const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void LoadExtension(const v8::FunctionCallbackInfo<v8::Value>& args);
 
   SET_MEMORY_INFO_NAME(DatabaseSync)
   SET_SELF_SIZE(DatabaseSync)
@@ -34,6 +36,7 @@ class DatabaseSync : public BaseObject {
 
   ~DatabaseSync() override;
   std::string location_;
+  bool allow_load_extension_;
   sqlite3* connection_;
 };
 

test/fixtures/sqlite/vec0.aarch64.dylib

@@ -1,9 +1,12 @@
 // Flags: --experimental-sqlite
 'use strict';
 const { spawnPromisified } = require('../common');
+const assert = require('node:assert');
 const tmpdir = require('../common/tmpdir');
 const { existsSync } = require('node:fs');
 const { join } = require('node:path');
+const os = require('node:os');
+const { path } = require('../common/fixtures');
 const { DatabaseSync, StatementSync } = require('node:sqlite');
 const { suite, test } = require('node:test');
 let cnt = 0;
@@ -78,6 +81,132 @@ suite('DatabaseSync() constructor', () => {
       message: /The "options\.open" argument must be a boolean/,
     });
   });
+
+  test('throws if options.allowLoadExtension is provided but is not a boolean', (t) => {
+    t.assert.throws(() => {
+      new DatabaseSync('foo', { allowLoadExtension: 5 });
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "options\.allowLoadExtension" argument must be a boolean/,
+    });
+  });
+});
+
+suite('DatabaseSync.prototype.loadExtension()', () => {
+  test('throws if database is not open', (t) => {
+    const db = new DatabaseSync(nextDb(), { open: false });
+    t.assert.throws(() => {
+      db.loadExtension();
+    }, {
+      code: 'ERR_INVALID_STATE',
+      message: /database is not open/,
+    });
+  });
+
+  test('throws if path is not a valid sqlite extension', (t) => {
+    const db = new DatabaseSync(nextDb(), {
+      allowLoadExtension: true,
+    });
+    // Try to load a non-existent file
+    const files = [
+      '/dev/null',
+      path('a.js'),
+      path('shared-memory.wasm'),
+      path('crash.wat'),
+      path('doc_inc_1.md'),
+      path('utf8-bom.json'),
+      path('x.txt'),
+    ];
+    for (const file of files) {
+      t.assert.throws(() => {
+        db.loadExtension(file);
+      }, {
+        code: 'ERR_LOAD_SQLITE_EXTENSION',
+      }, `loadExtension("${file}") should throw an error`);
+    }
+  });
+
+  test('should load sqlite extension successfully', (t) => {
+    const dbPath = nextDb();
+    const db = new DatabaseSync(dbPath, { allowLoadExtension: true });
+    const supportedPlatforms = [
+      ['macos', 'x86_64'],
+      ['windows', 'x86_64'],
+      ['linux', 'x86_64'],
+      ['macos', 'aarch64'],
+    ];
+    function validPlatform(platform, arch) {
+      return (
+        supportedPlatforms.find(([p, a]) => platform === p && arch === a) !== null
+      );
+    }
+
+    function getExtension(platform, arch) {
+      switch (platform) {
+        case 'darwin':
+          return arch === 'arm64' ? '.aarch64.dylib' : '.x86_64.dylib';
+        case 'windows':
+          return '.dll';
+        case 'linux':
+          return '.so';
+        default:
+          return null;
+      }
+    }
+    const platform = os.platform();
+    const arch = process.arch;
+    if (!validPlatform(platform, arch)) {
+      t.skip('Unsupported platform');
+      return;
+    }
+    const ext = getExtension(platform, arch);
+    const filePath = path('sqlite', `vec0${ext}`);
+    t.assert.strictEqual(db.loadExtension(filePath), undefined);
+
+    const { vec_version } = db.prepare(
+      'select vec_version() as vec_version;'
+    ).get();
+    assert.strictEqual(vec_version, 'v0.1.1');
+
+    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));
+
+    assert.deepStrictEqual(rows, [
+      { rowid: 3, distance: 0 },
+      { rowid: 4, distance: 0.19999998807907104 },
+      { rowid: 2, distance: 0.20000001788139343 },
+    ]);
+  });
 });
 
 suite('DatabaseSync.prototype.open()', () => {