From 10faf88591fe8c066f95638d03e33da1a68d13be Mon Sep 17 00:00:00 2001
From: Alex Yang <alex@runllama.ai>
Date: Wed, 17 Jul 2024 11:13:13 -0700
Subject: [PATCH] sqlite: support `db.loadExtension`

---
 src/node_sqlite.cc | 23 +++++++++++++++++++++++
 src/node_sqlite.h  |  1 +
 2 files changed, 24 insertions(+)

diff --git a/src/node_sqlite.cc b/src/node_sqlite.cc
index 1d94363bf08907..0a08f48e7bf3bb 100644
--- a/src/node_sqlite.cc
+++ b/src/node_sqlite.cc
@@ -109,6 +109,9 @@ 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);
+  int r2 = sqlite3_db_config(
+      connection_, SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION, 1, NULL);
+  CHECK_ERROR_OR_THROW(env()->isolate(), connection_, r2, SQLITE_OK, false);
   return true;
 }
 
@@ -211,6 +214,24 @@ 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");
+
+  if (!args[0]->IsString()) {
+    node::THROW_ERR_INVALID_ARG_TYPE(env->isolate(),
+                                     "The \"path\" argument must be a string.");
+    return;
+  }
+
+  auto path = node::Utf8Value(env->isolate(), args[0].As<String>());
+  int r = sqlite3_load_extension(db->connection_, *path, nullptr, nullptr);
+  CHECK_ERROR_OR_THROW(env->isolate(), db->connection_, r, SQLITE_OK, void());
+}
+
 StatementSync::StatementSync(Environment* env,
                              Local<Object> object,
                              sqlite3* db,
@@ -668,6 +689,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,
diff --git a/src/node_sqlite.h b/src/node_sqlite.h
index 56b937a719679b..004dd074fad080 100644
--- a/src/node_sqlite.h
+++ b/src/node_sqlite.h
@@ -25,6 +25,7 @@ class DatabaseSync : public BaseObject {
   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)