sqlite: add StatementSync.prototype.iterate method

PR-URL: #54213
Reviewed-By: Zeyu "Alex" Yang <[email protected]>
Reviewed-By: Colin Ihrig <[email protected]>

Author: tpoisseau

doc/api/sqlite.md

@@ -287,6 +287,25 @@ object. If the prepared statement does not return any results, this method
 returns `undefined`. The prepared statement [parameters are bound][] using the
 values in `namedParameters` and `anonymousParameters`.
 
+### `statement.iterate([namedParameters][, ...anonymousParameters])`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+* `namedParameters` {Object} An optional object used to bind named parameters.
+  The keys of this object are used to configure the mapping.
+* `...anonymousParameters` {null|number|bigint|string|Buffer|Uint8Array} Zero or
+  more values to bind to anonymous parameters.
+* Returns: {Iterator} An iterable iterator of objects. Each object corresponds to a row
+  returned by executing the prepared statement. The keys and values of each
+  object correspond to the column names and values of the row.
+
+This method executes a prepared statement and returns an iterator of
+objects. If the prepared statement does not return any results, this method
+returns an empty iterator. The prepared statement [parameters are bound][] using
+the values in `namedParameters` and `anonymousParameters`.
+
 ### `statement.run([namedParameters][, ...anonymousParameters])`
 
 <!-- YAML

src/env_properties.h

@@ -194,8 +194,10 @@
   V(ipv4_string, "IPv4")                                                       \
   V(ipv6_string, "IPv6")                                                       \
   V(isclosing_string, "isClosing")                                             \
+  V(isfinished_string, "isFinished")                                           \
   V(issuer_string, "issuer")                                                   \
   V(issuercert_string, "issuerCertificate")                                    \
+  V(iterator_string, "Iterator")                                               \
   V(jwk_crv_string, "crv")                                                     \
   V(jwk_d_string, "d")                                                         \
   V(jwk_dp_string, "dp")                                                       \
@@ -241,6 +243,7 @@
   V(nistcurve_string, "nistCurve")                                             \
   V(node_string, "node")                                                       \
   V(nsname_string, "nsname")                                                   \
+  V(num_cols_string, "num_cols")                                               \
   V(object_string, "Object")                                                   \
   V(ocsp_request_string, "OCSPRequest")                                        \
   V(oncertcb_string, "oncertcb")                                               \
@@ -288,6 +291,7 @@
   V(priority_string, "priority")                                               \
   V(process_string, "process")                                                 \
   V(promise_string, "promise")                                                 \
+  V(prototype_string, "prototype")                                             \
   V(psk_string, "psk")                                                         \
   V(pubkey_string, "pubkey")                                                   \
   V(public_exponent_string, "publicExponent")                                  \
@@ -309,6 +313,7 @@
   V(require_string, "require")                                                 \
   V(resource_string, "resource")                                               \
   V(retry_string, "retry")                                                     \
+  V(return_string, "return")                                                   \
   V(salt_length_string, "saltLength")                                          \
   V(scheme_string, "scheme")                                                   \
   V(scopeid_string, "scopeid")                                                 \
@@ -332,6 +337,7 @@
   V(standard_name_string, "standardName")                                      \
   V(start_time_string, "startTime")                                            \
   V(state_string, "state")                                                     \
+  V(statement_string, "statement")                                             \
   V(stats_string, "stats")                                                     \
   V(status_string, "status")                                                   \
   V(stdio_string, "stdio")                                                     \

src/node_sqlite.cc

@@ -22,6 +22,7 @@ using v8::ConstructorBehavior;
 using v8::Context;
 using v8::DontDelete;
 using v8::Exception;
+using v8::External;
 using v8::Function;
 using v8::FunctionCallback;
 using v8::FunctionCallbackInfo;
@@ -790,6 +791,180 @@ void StatementSync::All(const FunctionCallbackInfo<Value>& args) {
   args.GetReturnValue().Set(Array::New(isolate, rows.data(), rows.size()));
 }
 
+void StatementSync::IterateReturnCallback(
+    const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  auto isolate = env->isolate();
+  auto context = isolate->GetCurrentContext();
+
+  auto self = args.This();
+  // iterator has fetch all result or break, prevent next func to return result
+  self->Set(context, env->isfinished_string(), Boolean::New(isolate, true))
+      .ToChecked();
+
+  auto external_stmt = Local<External>::Cast(
+      self->Get(context, env->statement_string()).ToLocalChecked());
+  auto stmt = static_cast<StatementSync*>(external_stmt->Value());
+  if (!stmt->IsFinalized()) {
+    sqlite3_reset(stmt->statement_);
+  }
+
+  LocalVector<Name> keys(isolate, {env->done_string(), env->value_string()});
+  LocalVector<Value> values(isolate,
+                            {Boolean::New(isolate, true), Null(isolate)});
+
+  DCHECK_EQ(keys.size(), values.size());
+  Local<Object> result = Object::New(
+      isolate, Null(isolate), keys.data(), values.data(), keys.size());
+  args.GetReturnValue().Set(result);
+}
+
+void StatementSync::IterateNextCallback(
+    const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  auto isolate = env->isolate();
+  auto context = isolate->GetCurrentContext();
+
+  auto self = args.This();
+
+  // skip iteration if is_finished
+  auto is_finished = Local<Boolean>::Cast(
+      self->Get(context, env->isfinished_string()).ToLocalChecked());
+  if (is_finished->Value()) {
+    LocalVector<Name> keys(isolate, {env->done_string(), env->value_string()});
+    LocalVector<Value> values(isolate,
+                              {Boolean::New(isolate, true), Null(isolate)});
+
+    DCHECK_EQ(keys.size(), values.size());
+    Local<Object> result = Object::New(
+        isolate, Null(isolate), keys.data(), values.data(), keys.size());
+    args.GetReturnValue().Set(result);
+    return;
+  }
+
+  auto external_stmt = Local<External>::Cast(
+      self->Get(context, env->statement_string()).ToLocalChecked());
+  auto stmt = static_cast<StatementSync*>(external_stmt->Value());
+  auto num_cols =
+      Local<Integer>::Cast(
+          self->Get(context, env->num_cols_string()).ToLocalChecked())
+          ->Value();
+
+  THROW_AND_RETURN_ON_BAD_STATE(
+      env, stmt->IsFinalized(), "statement has been finalized");
+
+  int r = sqlite3_step(stmt->statement_);
+  if (r != SQLITE_ROW) {
+    CHECK_ERROR_OR_THROW(
+        env->isolate(), stmt->db_->Connection(), r, SQLITE_DONE, void());
+
+    // cleanup when no more rows to fetch
+    sqlite3_reset(stmt->statement_);
+    self->Set(context, env->isfinished_string(), Boolean::New(isolate, true))
+        .ToChecked();
+
+    LocalVector<Name> keys(isolate, {env->done_string(), env->value_string()});
+    LocalVector<Value> values(isolate,
+                              {Boolean::New(isolate, true), Null(isolate)});
+
+    DCHECK_EQ(keys.size(), values.size());
+    Local<Object> result = Object::New(
+        isolate, Null(isolate), keys.data(), values.data(), keys.size());
+    args.GetReturnValue().Set(result);
+    return;
+  }
+
+  LocalVector<Name> row_keys(isolate);
+  row_keys.reserve(num_cols);
+  LocalVector<Value> row_values(isolate);
+  row_values.reserve(num_cols);
+  for (int i = 0; i < num_cols; ++i) {
+    Local<Name> key;
+    if (!stmt->ColumnNameToName(i).ToLocal(&key)) return;
+    Local<Value> val;
+    if (!stmt->ColumnToValue(i).ToLocal(&val)) return;
+    row_keys.emplace_back(key);
+    row_values.emplace_back(val);
+  }
+
+  Local<Object> row = Object::New(
+      isolate, Null(isolate), row_keys.data(), row_values.data(), num_cols);
+
+  LocalVector<Name> keys(isolate, {env->done_string(), env->value_string()});
+  LocalVector<Value> values(isolate, {Boolean::New(isolate, false), row});
+
+  DCHECK_EQ(keys.size(), values.size());
+  Local<Object> result = Object::New(
+      isolate, Null(isolate), keys.data(), values.data(), keys.size());
+  args.GetReturnValue().Set(result);
+}
+
+void StatementSync::Iterate(const FunctionCallbackInfo<Value>& args) {
+  StatementSync* stmt;
+  ASSIGN_OR_RETURN_UNWRAP(&stmt, args.This());
+  Environment* env = Environment::GetCurrent(args);
+  THROW_AND_RETURN_ON_BAD_STATE(
+      env, stmt->IsFinalized(), "statement has been finalized");
+  auto isolate = env->isolate();
+  auto context = env->context();
+  int r = sqlite3_reset(stmt->statement_);
+  CHECK_ERROR_OR_THROW(
+      env->isolate(), stmt->db_->Connection(), r, SQLITE_OK, void());
+
+  if (!stmt->BindParams(args)) {
+    return;
+  }
+
+  Local<Function> next_func =
+      Function::New(context, StatementSync::IterateNextCallback)
+          .ToLocalChecked();
+  Local<Function> return_func =
+      Function::New(context, StatementSync::IterateReturnCallback)
+          .ToLocalChecked();
+
+  LocalVector<Name> keys(isolate, {env->next_string(), env->return_string()});
+  LocalVector<Value> values(isolate, {next_func, return_func});
+
+  Local<Object> global = context->Global();
+  Local<Value> js_iterator;
+  Local<Value> js_iterator_prototype;
+  if (!global->Get(context, env->iterator_string()).ToLocal(&js_iterator))
+    return;
+  if (!js_iterator.As<Object>()
+           ->Get(context, env->prototype_string())
+           .ToLocal(&js_iterator_prototype))
+    return;
+
+  DCHECK_EQ(keys.size(), values.size());
+  Local<Object> iterable_iterator = Object::New(
+      isolate, js_iterator_prototype, keys.data(), values.data(), keys.size());
+
+  auto num_cols_pd = v8::PropertyDescriptor(
+      v8::Integer::New(isolate, sqlite3_column_count(stmt->statement_)), false);
+  num_cols_pd.set_enumerable(false);
+  num_cols_pd.set_configurable(false);
+  iterable_iterator
+      ->DefineProperty(context, env->num_cols_string(), num_cols_pd)
+      .ToChecked();
+
+  auto stmt_pd =
+      v8::PropertyDescriptor(v8::External::New(isolate, stmt), false);
+  stmt_pd.set_enumerable(false);
+  stmt_pd.set_configurable(false);
+  iterable_iterator->DefineProperty(context, env->statement_string(), stmt_pd)
+      .ToChecked();
+
+  auto is_finished_pd =
+      v8::PropertyDescriptor(v8::Boolean::New(isolate, false), true);
+  stmt_pd.set_enumerable(false);
+  stmt_pd.set_configurable(false);
+  iterable_iterator
+      ->DefineProperty(context, env->isfinished_string(), is_finished_pd)
+      .ToChecked();
+
+  args.GetReturnValue().Set(iterable_iterator);
+}
+
 void StatementSync::Get(const FunctionCallbackInfo<Value>& args) {
   StatementSync* stmt;
   ASSIGN_OR_RETURN_UNWRAP(&stmt, args.This());
@@ -987,6 +1162,7 @@ Local<FunctionTemplate> StatementSync::GetConstructorTemplate(
     tmpl->SetClassName(FIXED_ONE_BYTE_STRING(isolate, "StatementSync"));
     tmpl->InstanceTemplate()->SetInternalFieldCount(
         StatementSync::kInternalFieldCount);
+    SetProtoMethod(isolate, tmpl, "iterate", StatementSync::Iterate);
     SetProtoMethod(isolate, tmpl, "all", StatementSync::All);
     SetProtoMethod(isolate, tmpl, "get", StatementSync::Get);
     SetProtoMethod(isolate, tmpl, "run", StatementSync::Run);

src/node_sqlite.h

@@ -93,6 +93,7 @@ class StatementSync : public BaseObject {
                                              DatabaseSync* db,
                                              sqlite3_stmt* stmt);
   static void All(const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void Iterate(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void Get(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void Run(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void SourceSQLGetter(const v8::FunctionCallbackInfo<v8::Value>& args);
@@ -118,6 +119,11 @@ class StatementSync : public BaseObject {
   bool BindValue(const v8::Local<v8::Value>& value, const int index);
   v8::MaybeLocal<v8::Value> ColumnToValue(const int column);
   v8::MaybeLocal<v8::Name> ColumnNameToName(const int column);
+
+  static void IterateNextCallback(
+      const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void IterateReturnCallback(
+      const v8::FunctionCallbackInfo<v8::Value>& args);
 };
 
 using Sqlite3ChangesetGenFunc = int (*)(sqlite3_session*, int*, void**);

test/parallel/test-sqlite-statement-sync.js

@@ -84,6 +84,42 @@ suite('StatementSync.prototype.all()', () => {
   });
 });
 
+suite('StatementSync.prototype.iterate()', () => {
+  test('executes a query and returns an empty iterator on no results', (t) => {
+    const db = new DatabaseSync(nextDb());
+    const stmt = db.prepare('CREATE TABLE storage(key TEXT, val TEXT)');
+    t.assert.deepStrictEqual(stmt.iterate().toArray(), []);
+  });
+
+  test('executes a query and returns all results', (t) => {
+    const db = new DatabaseSync(nextDb());
+    let stmt = db.prepare('CREATE TABLE storage(key TEXT, val TEXT)');
+    t.assert.deepStrictEqual(stmt.run(), { changes: 0, lastInsertRowid: 0 });
+    stmt = db.prepare('INSERT INTO storage (key, val) VALUES (?, ?)');
+    t.assert.deepStrictEqual(
+      stmt.run('key1', 'val1'),
+      { changes: 1, lastInsertRowid: 1 },
+    );
+    t.assert.deepStrictEqual(
+      stmt.run('key2', 'val2'),
+      { changes: 1, lastInsertRowid: 2 },
+    );
+
+    const items = [
+      { __proto__: null, key: 'key1', val: 'val1' },
+      { __proto__: null, key: 'key2', val: 'val2' },
+    ];
+
+    stmt = db.prepare('SELECT * FROM storage ORDER BY key');
+    t.assert.deepStrictEqual(stmt.iterate().toArray(), items);
+
+    const itemsLoop = items.slice();
+    for (const item of stmt.iterate()) {
+      t.assert.deepStrictEqual(item, itemsLoop.shift());
+    }
+  });
+});
+
 suite('StatementSync.prototype.run()', () => {
   test('executes a query and returns change metadata', (t) => {
     const db = new DatabaseSync(nextDb());