[bun:sqlite] Support multiple statements in db.run() (#8541)

* [bun:sqlite] Support multiple statements in db.run() * Update sqlite.test.js * Update JSSQLStatement.cpp * Another test --------- Co-authored-by: Jarred Sumner <[email protected]>
oven-sh · Jan 29, 2024 · 413aaaf · 413aaaf
1 parent c538bf8
commit 413aaaf
Show file tree

Hide file tree

Showing 2 changed files with 156 additions and 40 deletions.
diff --git a/src/bun.js/bindings/sqlite/JSSQLStatement.cpp b/src/bun.js/bindings/sqlite/JSSQLStatement.cpp
@@ -96,6 +96,16 @@ static void enableFastMallocForSQLite()
 #endif
 }
 
+class AutoDestructingSQLiteStatement {
+public:
+    sqlite3_stmt* stmt { nullptr };
+
+    ~AutoDestructingSQLiteStatement()
+    {
+        sqlite3_finalize(stmt);
+    }
+};
+
 static void initializeSQLite()
 {
     static std::once_flag onceFlag;
@@ -899,6 +909,11 @@ JSC_DEFINE_HOST_FUNCTION(jsSQLStatementLoadExtensionFunction, (JSC::JSGlobalObje
     RELEASE_AND_RETURN(scope, JSValue::encode(JSC::jsUndefined()));
 }
 
+static bool isSkippedInSQLiteQuery(const char c)
+{
+    return c == ' ' || c == ';' || (c >= '\t' && c <= '\r');
+}
+
 // This runs a query one-off
 // without the overhead of a long-lived statement object
 // does not return anything
@@ -945,65 +960,92 @@ JSC_DEFINE_HOST_FUNCTION(jsSQLStatementExecuteFunction, (JSC::JSGlobalObject * l
         return JSValue::encode(JSC::jsUndefined());
     }
 
-    // TODO: trim whitespace & newlines before sending
-    // we don't because webkit doesn't expose a function that makes this super
-    // easy without using unicode whitespace definition the
-    // StringPrototype.trim() implementation GC allocates a new JSString* and
-    // potentially re-allocates the string (not 100% sure if reallocates) so we
-    // can't use that here
-    sqlite3_stmt* statement = nullptr;
+    CString utf8;
+
+    const char* sqlStringHead;
+    const char* end;
+    bool didSetBindings = false;
 
-    int rc = SQLITE_OK;
     if (
         // fast path: ascii latin1 string is utf8
         sqlString.is8Bit() && simdutf::validate_ascii(reinterpret_cast<const char*>(sqlString.characters8()), sqlString.length())) {
-        rc = sqlite3_prepare_v3(db, reinterpret_cast<const char*>(sqlString.characters8()), sqlString.length(), 0, &statement, nullptr);
+
+        sqlStringHead = reinterpret_cast<const char*>(sqlString.characters8());
+        end = sqlStringHead + sqlString.length();
     } else {
         // slow path: utf16 or latin1 string with supplemental characters
-        CString utf8 = sqlString.utf8();
-        rc = sqlite3_prepare_v3(db, utf8.data(), utf8.length(), 0, &statement, nullptr);
+        utf8 = sqlString.utf8();
+        sqlStringHead = utf8.data();
+        end = sqlStringHead + utf8.length();
     }
 
-    if (UNLIKELY(rc != SQLITE_OK || statement == nullptr)) {
-        throwException(lexicalGlobalObject, scope, createError(lexicalGlobalObject, rc == SQLITE_OK ? "Query contained no valid SQL statement; likely empty query."_s : WTF::String::fromUTF8(sqlite3_errmsg(db))));
-        // sqlite3 handles when the pointer is null
-        sqlite3_finalize(statement);
-        return JSValue::encode(JSC::jsUndefined());
-    }
+    bool didExecuteAny = false;
+
+    int rc = SQLITE_OK;
+
+#if ASSERT_ENABLED
+    int maxSqlStringBytes = end - sqlStringHead;
+#endif
 
-    if (!bindingsAliveScope.value().isUndefinedOrNull()) {
-        if (bindingsAliveScope.value().isObject()) {
-            JSC::JSValue reb = rebindStatement(lexicalGlobalObject, bindingsAliveScope.value(), scope, db, statement, false);
-            if (UNLIKELY(!reb.isNumber())) {
-                sqlite3_finalize(statement);
-                return JSValue::encode(reb); /* this means an error */
+    while (sqlStringHead && sqlStringHead < end) {
+        if (UNLIKELY(isSkippedInSQLiteQuery(*sqlStringHead))) {
+            sqlStringHead++;
+
+            while (sqlStringHead < end && isSkippedInSQLiteQuery(*sqlStringHead))
+                sqlStringHead++;
+        }
+
+        AutoDestructingSQLiteStatement sql;
+        const char* tail = nullptr;
+
+        // Bounds checks
+        ASSERT(end >= sqlStringHead);
+        ASSERT(end - sqlStringHead >= 0);
+        ASSERT(end - sqlStringHead <= maxSqlStringBytes);
+
+        rc = sqlite3_prepare_v3(db, sqlStringHead, end - sqlStringHead, 0, &sql.stmt, &tail);
+
+        if (rc != SQLITE_OK)
+            break;
+
+        if (!sql.stmt) {
+            // this happens for an empty statement
+            sqlStringHead = tail;
+            continue;
+        }
+
+        // First statement gets the bindings.
+        if (!didSetBindings && !bindingsAliveScope.value().isUndefinedOrNull()) {
+            if (bindingsAliveScope.value().isObject()) {
+                JSC::JSValue reb = rebindStatement(lexicalGlobalObject, bindingsAliveScope.value(), scope, db, sql.stmt, false);
+                if (UNLIKELY(!reb.isNumber())) {
+                    return JSValue::encode(reb); /* this means an error */
+                }
+            } else {
+                throwException(lexicalGlobalObject, scope, createTypeError(lexicalGlobalObject, "Expected bindings to be an object or array"_s));
+                return JSValue::encode(jsUndefined());
             }
-        } else {
-            throwException(lexicalGlobalObject, scope, createTypeError(lexicalGlobalObject, "Expected bindings to be an object or array"_s));
-            sqlite3_finalize(statement);
-            return JSValue::encode(jsUndefined());
+            didSetBindings = true;
         }
-    }
 
-    rc = sqlite3_step(statement);
-    if (!sqlite3_stmt_readonly(statement)) {
-        databases()[handle]->version++;
-    }
+        do {
+            rc = sqlite3_step(sql.stmt);
+        } while (rc == SQLITE_ROW);
 
-    while (rc == SQLITE_ROW) {
-        rc = sqlite3_step(statement);
+        didExecuteAny = true;
+        sqlStringHead = tail;
     }
 
-    if (UNLIKELY(rc != SQLITE_DONE && rc != SQLITE_OK)) {
+    if (UNLIKELY(rc != SQLITE_OK && rc != SQLITE_DONE)) {
         throwException(lexicalGlobalObject, scope, createSQLiteError(lexicalGlobalObject, db));
-        // we finalize after just incase something about error messages in
-        // sqlite depends on the existence of the most recent statement i don't
-        // think that's actually how this works - just being cautious
-        sqlite3_finalize(statement);
         return JSValue::encode(JSC::jsUndefined());
     }
 
-    sqlite3_finalize(statement);
+    if (!didExecuteAny) {
+        throwException(lexicalGlobalObject, scope, createError(lexicalGlobalObject, "Query contained no valid SQL statement; likely empty query."_s));
+        return JSValue::encode(JSC::jsUndefined());
+    }
+
     return JSValue::encode(jsUndefined());
 }
 

diff --git a/test/js/bun/sqlite/sqlite.test.js b/test/js/bun/sqlite/sqlite.test.js
@@ -669,3 +669,77 @@ it("empty blob", () => {
     },
   ]);
 });
+
+it("multiple statements with a schema change", () => {
+  const db = new Database(":memory:");
+  db.run(
+    `
+    CREATE TABLE foo (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT); 
+    CREATE TABLE bar (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT);
+
+    INSERT INTO foo (name) VALUES ('foo');
+    INSERT INTO foo (name) VALUES ('bar');
+
+    INSERT INTO bar (name) VALUES ('foo');
+    INSERT INTO bar (name) VALUES ('bar');
+  `,
+  );
+
+  expect(db.query("SELECT * FROM foo").all()).toEqual([
+    {
+      id: 1,
+      name: "foo",
+    },
+    {
+      id: 2,
+      name: "bar",
+    },
+  ]);
+
+  expect(db.query("SELECT * FROM bar").all()).toEqual([
+    {
+      id: 1,
+      name: "foo",
+    },
+    {
+      id: 2,
+      name: "bar",
+    },
+  ]);
+});
+
+it("multiple statements", () => {
+  const fixtures = [
+    "INSERT INTO foo (name) VALUES ('foo')",
+    "INSERT INTO foo (name) VALUES ('barabc')",
+    "INSERT INTO foo (name) VALUES ('!bazaspdok')",
+  ];
+  for (let separator of [";", ";\n", "\n;", "\r\n;", ";\r\n", ";\t", "\t;", "\r\n;"]) {
+    for (let spaceOffset of [1, 0, -1]) {
+      for (let spacesCount = 0; spacesCount < 8; spacesCount++) {
+        const db = new Database(":memory:");
+        db.run("CREATE TABLE foo (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT)");
+
+        const prefix = spaceOffset < 0 ? " ".repeat(spacesCount) : "";
+        const suffix = spaceOffset > 0 ? " ".repeat(spacesCount) : "";
+        const query = fixtures.join(prefix + separator + suffix);
+        db.run(query);
+
+        expect(db.query("SELECT * FROM foo").all()).toEqual([
+          {
+            id: 1,
+            name: "foo",
+          },
+          {
+            id: 2,
+            name: "barabc",
+          },
+          {
+            id: 3,
+            name: "!bazaspdok",
+          },
+        ]);
+      }
+    }
+  }
+});