feat: add StatementIterator#nextBatch() and #restart()

docs/api.md

@@ -513,6 +513,25 @@ for (const cat of stmt.iterate()) {
 }
 ```
 
+#### iterator.nextBatch(*maxRows*) -> *array*
+
+Steps up to `maxRows` rows in one call, instead of one call per row as `next()` does. In [raw mode](#rawtogglestate---this), it returns one flat array of `rowCount * columnCount` values, in row order. In any other mode, it returns an array of rows, as [`.all()`](#allbindparameters---array-of-rows) would.
+
+A batch with fewer than `maxRows` rows is the last one: the iterator is then finished, exactly as when `next()` returns `done`. `nextBatch()` and `next()` can be mixed, and `return()` closes the iterator early.
+
+```js
+const stmt = db.prepare('SELECT id, name FROM cats').raw();
+const it = stmt.iterate();
+try {
+  for (let batch; (batch = it.nextBatch(64)).length > 0; ) {
+    for (let i = 0; i < batch.length; i += 2) console.log(batch[i], batch[i + 1]);
+    if (batch.length < 64 * 2) break;
+  }
+} finally {
+  it.return();
+}
+```
+
 ### .pluck([toggleState]) -> *this*
 
 **(only on statements that return data)*

lib/index.d.ts

@@ -27,7 +27,7 @@ declare namespace BetterSqlite3 {
         run(...params: BindParameters): Database.RunResult;
         get(...params: BindParameters): Result | undefined;
         all(...params: BindParameters): Result[];
-        iterate(...params: BindParameters): IterableIterator<Result>;
+        iterate(...params: BindParameters): StatementIterator<Result>;
         pluck(toggleState?: boolean): this;
         expand(toggleState?: boolean): this;
         raw(toggleState?: boolean): this;
@@ -47,6 +47,20 @@ declare namespace BetterSqlite3 {
         explainQueryPlan(...params: BindParameters | []): Result[];
     }
 
+    interface StatementIterator<Result = unknown> extends IterableIterator<Result> {
+        /**
+         * Steps up to `maxRows` rows in one call. In raw mode, returns one flat
+         * array of rowCount * columnCount values; otherwise, an array of rows.
+         * A batch with fewer than `maxRows` rows is the last one.
+         */
+        nextBatch(maxRows: number): unknown[];
+        /**
+         * Runs the statement again with new bind parameters, reusing this
+         * iterator, which must be finished. Cheaper than calling iterate() again.
+         */
+        restart(...params: unknown[]): this;
+    }
+
     interface ColumnDefinition {
         name: string;
         column: string | null;
@@ -176,6 +190,7 @@ declare namespace Database {
         unknown[] ? BetterSqlite3.Statement<BindParameters, Result>
         : BetterSqlite3.Statement<[BindParameters], Result>;
     type ColumnDefinition = BetterSqlite3.ColumnDefinition;
+    type StatementIterator<Result = unknown> = BetterSqlite3.StatementIterator<Result>;
     type Transaction<T extends VariableArgFunction = VariableArgFunction> = BetterSqlite3.Transaction<T>;
     type Database = BetterSqlite3.Database;
 }

src/objects/statement-iterator.cpp

@@ -48,6 +48,41 @@ void StatementIterator::Next(NODE_ARGUMENTS info) {
 	}
 }
 
+// Steps up to max_rows rows in one call. In raw mode the result is one flat
+// array of rowCount * columnCount values; in other modes, an array of rows.
+// A batch with fewer than max_rows rows is the last one.
+void StatementIterator::NextBatch(NODE_ARGUMENTS info, int max_rows) {
+	assert(alive == true);
+	db_state->busy = true;
+	if (!logged) {
+		logged = true;
+		if (stmt->db->Log(OnlyIsolate, handle)) {
+			db_state->busy = false;
+			Throw();
+			return;
+		}
+	}
+	UseIsolate;
+	UseContext;
+	v8::LocalVector<v8::Value> values(isolate);
+	const int columns = mode == Data::RAW ? sqlite3_column_count(handle) : 0;
+	int rows = 0;
+	int status = SQLITE_ROW;
+	while (rows < max_rows && (status = sqlite3_step(handle)) == SQLITE_ROW) {
+		if (mode == Data::RAW) {
+			for (int i = 0; i < columns; ++i) values.emplace_back(Data::GetValueJS(isolate, handle, i, safe_ints));
+		} else {
+			values.emplace_back(stmt->GetRowJS(isolate, ctx, safe_ints, mode));
+		}
+		++rows;
+	}
+	db_state->busy = false;
+	v8::Local<v8::Array> batch = v8::Array::New(isolate, values.data(), values.size());
+	if (status == SQLITE_ROW) info.GetReturnValue().Set(batch);
+	else if (status == SQLITE_DONE) { Cleanup(); STATEMENT_RETURN_LOGIC(batch); }
+	else Throw();
+}
+
 void StatementIterator::Return(NODE_ARGUMENTS info) {
 	Cleanup();
 	STATEMENT_RETURN_LOGIC(DoneRecord(OnlyIsolate, db_state->addon));
@@ -70,6 +105,8 @@ void StatementIterator::Cleanup() {
 INIT(StatementIterator::Init) {
 	v8::Local<v8::FunctionTemplate> t = NewConstructorTemplate(isolate, data, JS_new, "StatementIterator");
 	SetPrototypeMethod(isolate, data, t, "next", JS_next);
+	SetPrototypeMethod(isolate, data, t, "nextBatch", JS_nextBatch);
+	SetPrototypeMethod(isolate, data, t, "restart", JS_restart);
 	SetPrototypeMethod(isolate, data, t, "return", JS_return);
 	SetPrototypeSymbolMethod(isolate, data, t, v8::Symbol::GetIterator(isolate), JS_symbolIterator);
 	return t->GetFunction(OnlyContext).ToLocalChecked();
@@ -101,6 +138,41 @@ NODE_METHOD(StatementIterator::JS_next) {
 	else info.GetReturnValue().Set(DoneRecord(OnlyIsolate, iter->db_state->addon));
 }
 
+NODE_METHOD(StatementIterator::JS_nextBatch) {
+	StatementIterator* iter = Unwrap<StatementIterator>(info.This());
+	REQUIRE_DATABASE_NOT_BUSY(iter->db_state);
+	REQUIRE_ARGUMENT_INT32(first, int max_rows);
+	if (max_rows < 1) return ThrowRangeError("Expected first argument to be a positive integer");
+	if (iter->alive) iter->NextBatch(info, max_rows);
+	else info.GetReturnValue().Set(v8::Array::New(OnlyIsolate, 0));
+}
+
+// Runs the statement again, with new bind parameters, reusing this finished
+// iterator instead of creating a new one with iterate().
+NODE_METHOD(StatementIterator::JS_restart) {
+	StatementIterator* iter = Unwrap<StatementIterator>(info.This());
+	if (iter->alive) return ThrowTypeError("The iterator must be finished before it is restarted");
+	Statement* stmt = iter->stmt;
+	sqlite3_stmt* handle = stmt->handle;
+	Database* db = stmt->db;
+	REQUIRE_DATABASE_OPEN(db->GetState());
+	REQUIRE_DATABASE_NOT_BUSY(db->GetState());
+	DOES_ADD_ITERATOR();
+	if (!stmt->bound) {
+		STATEMENT_BIND(handle);
+	} else if (info.Length() > 0) {
+		return ThrowTypeError("This statement already has bound parameters");
+	}
+	iter->bound = stmt->bound;
+	iter->safe_ints = stmt->safe_ints;
+	iter->mode = stmt->mode;
+	iter->alive = true;
+	iter->logged = !iter->db_state->has_logger;
+	stmt->locked = true;
+	iter->db_state->iterators += 1;
+	info.GetReturnValue().Set(info.This());
+}
+
 NODE_METHOD(StatementIterator::JS_return) {
 	StatementIterator* iter = Unwrap<StatementIterator>(info.This());
 	REQUIRE_DATABASE_NOT_BUSY(iter->db_state);

src/objects/statement-iterator.hpp

@@ -13,6 +13,7 @@ class StatementIterator : public node::ObjectWrap {
 	explicit StatementIterator(Statement* stmt, bool bound);
 
 	void Next(NODE_ARGUMENTS info);
+	void NextBatch(NODE_ARGUMENTS info, int max_rows);
 	void Return(NODE_ARGUMENTS info);
 	void Throw();
 	void Cleanup();
@@ -36,15 +37,18 @@ class StatementIterator : public node::ObjectWrap {
 
 	static NODE_METHOD(JS_new);
 	static NODE_METHOD(JS_next);
+	static NODE_METHOD(JS_nextBatch);
+	static NODE_METHOD(JS_restart);
 	static NODE_METHOD(JS_return);
 	static NODE_METHOD(JS_symbolIterator);
 
 	Statement* const stmt;
 	sqlite3_stmt* const handle;
 	Database::State* const db_state;
-	const bool bound;
-	const bool safe_ints;
-	const char mode;
+	// Not const: restart() takes them from the statement again.
+	bool bound;
+	bool safe_ints;
+	char mode;
 	bool alive;
 	bool logged;
 };

test/29.statement.iterate.next-batch.js

@@ -0,0 +1,179 @@
+'use strict';
+const Database = require('../.');
+
+// Drains an iterator with nextBatch(); returns {batches, values}.
+function drain(iterator, maxRows, width) {
+	const batches = [];
+	const values = [];
+	for (;;) {
+		const batch = iterator.nextBatch(maxRows);
+		batches.push(batch.length / width);
+		values.push(...batch);
+		if (batch.length < maxRows * width) return {batches, values};
+	}
+}
+
+describe('StatementIterator#nextBatch()', function () {
+	beforeEach(function () {
+		this.db = new Database(util.next());
+		this.db.prepare('CREATE TABLE entries (a TEXT, b INTEGER, c REAL, d BLOB, e TEXT)').run();
+		this.db.prepare("INSERT INTO entries WITH RECURSIVE temp(a, b, c, d, e) AS (SELECT 'foo', 1, 3.14, x'dddddddd', NULL UNION ALL SELECT a, b + 1, c, d, e FROM temp LIMIT 10) SELECT * FROM temp").run();
+	});
+	afterEach(function () {
+		this.db.close();
+	});
+
+	it('should return flat values in raw mode', function () {
+		this.db.prepare("INSERT INTO entries VALUES ('héllo ☃', 9223372036854775807, -0.5, x'', 'x')").run();
+		const sql = 'SELECT * FROM entries ORDER BY rowid';
+		const want = this.db.prepare(sql).raw().safeIntegers().all().flat();
+		for (const [maxRows, batches] of [[1, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]], [4, [4, 4, 3]], [11, [11, 0]], [12, [11]], [1000, [11]]]) {
+			const stmt = this.db.prepare(sql).raw().safeIntegers();
+			const result = drain(stmt.iterate(), maxRows, 5);
+			expect(result.values).to.deep.equal(want);
+			expect(result.batches).to.deep.equal(batches);
+			expect(stmt.busy).to.be.false;
+		}
+	});
+	it('should honor safeIntegers()', function () {
+		const stmt = this.db.prepare('SELECT b FROM entries ORDER BY rowid LIMIT 2').raw();
+		expect(stmt.safeIntegers(false).iterate().nextBatch(10)).to.deep.equal([1, 2]);
+		expect(stmt.safeIntegers(true).iterate().nextBatch(10)).to.deep.equal([1n, 2n]);
+	});
+	it('should return rows in the other modes', function () {
+		const sql = 'SELECT a, b FROM entries ORDER BY rowid LIMIT 3';
+		for (const mode of ['', 'pluck', 'expand']) {
+			const stmt = this.db.prepare(sql);
+			if (mode) stmt[mode]();
+			const want = stmt.all();
+			const it = stmt.iterate();
+			expect([...it.nextBatch(2), ...it.nextBatch(2)]).to.deep.equal(want);
+			expect(stmt.busy).to.be.false;
+		}
+	});
+	it('should mix with next() and stop at return()', function () {
+		const stmt = this.db.prepare('SELECT b FROM entries ORDER BY rowid').raw();
+		const it = stmt.iterate();
+		expect(it.next().value).to.deep.equal([1]);
+		expect(it.nextBatch(3)).to.deep.equal([2, 3, 4]);
+		expect(it.next().value).to.deep.equal([5]);
+		expect(stmt.busy).to.be.true;
+		expect(it.return()).to.deep.equal({value: undefined, done: true});
+		expect(stmt.busy).to.be.false;
+		expect(it.nextBatch(3)).to.deep.equal([]);
+		expect(it.next()).to.deep.equal({value: undefined, done: true});
+		expect(stmt.all().length).to.equal(10);
+	});
+	it('should keep the statement and database busy until the end', function () {
+		const stmt = this.db.prepare('SELECT b FROM entries WHERE b > ? ORDER BY rowid').raw();
+		const insert = this.db.prepare("INSERT INTO entries VALUES ('bar', 0, 0, NULL, NULL)");
+		const it = stmt.iterate(2);
+		expect(it.nextBatch(4)).to.deep.equal([3, 4, 5, 6]);
+		expect(() => stmt.all()).to.throw(TypeError);
+		expect(() => insert.run()).to.throw(TypeError);
+		expect(() => this.db.close()).to.throw(TypeError);
+		expect(it.nextBatch(4)).to.deep.equal([7, 8, 9, 10]);
+		expect(stmt.busy).to.be.true;
+		expect(it.nextBatch(4)).to.deep.equal([]);
+		expect(stmt.busy).to.be.false;
+		insert.run();
+		// The bind parameters were cleared at the end.
+		expect(() => stmt.all()).to.throw(RangeError);
+	});
+	it('should validate maxRows', function () {
+		const stmt = this.db.prepare('SELECT b FROM entries');
+		const it = stmt.iterate();
+		expect(() => it.nextBatch()).to.throw(TypeError);
+		expect(() => it.nextBatch(1.5)).to.throw(TypeError);
+		expect(() => it.nextBatch('2')).to.throw(TypeError);
+		expect(() => it.nextBatch(0)).to.throw(RangeError);
+		expect(() => it.nextBatch(-1)).to.throw(RangeError);
+		expect(stmt.busy).to.be.true;
+		it.return();
+	});
+	it('should close the iterator when SQLite fails mid-batch', function () {
+		const stmt = this.db.prepare('SELECT CASE WHEN b > 5 THEN json(a || b) ELSE b END FROM entries ORDER BY rowid').raw();
+		const it = stmt.iterate();
+		expect(it.nextBatch(3)).to.deep.equal([1, 2, 3]);
+		expect(() => it.nextBatch(3)).to.throw(Database.SqliteError);
+		expect(stmt.busy).to.be.false;
+		expect(it.nextBatch(3)).to.deep.equal([]);
+	});
+	it('should propagate an exception thrown by a user-defined function', function () {
+		const err = new Error('boom');
+		this.db.function('fail', x => { if (x > 4) throw err; return x; });
+		const stmt = this.db.prepare('SELECT fail(b) FROM entries ORDER BY rowid').raw();
+		const it = stmt.iterate();
+		expect(it.nextBatch(3)).to.deep.equal([1, 2, 3]);
+		expect(() => it.nextBatch(3)).to.throw(err);
+		expect(stmt.busy).to.be.false;
+	});
+	it('should log the statement once per iterator in verbose mode', function () {
+		const logged = [];
+		const db = new Database(util.next(), {verbose: sql => logged.push(sql)});
+		const stmt = db.prepare('SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3').raw();
+		logged.length = 0;
+		drain(stmt.iterate(), 1, 1);
+		const it = stmt.iterate();
+		it.next();
+		it.nextBatch(5);
+		expect(logged).to.deep.equal([stmt.source, stmt.source]);
+		db.close();
+	});
+
+	describe('restart()', function () {
+		it('should run the statement again with new bind parameters', function () {
+			const stmt = this.db.prepare('SELECT b FROM entries WHERE b > ? ORDER BY rowid').raw();
+			const it = stmt.iterate(7);
+			expect(it.nextBatch(10)).to.deep.equal([8, 9, 10]);
+			expect(it.restart(8)).to.equal(it);
+			expect(stmt.busy).to.be.true;
+			expect(it.nextBatch(10)).to.deep.equal([9, 10]);
+			expect(it.restart(1).nextBatch(2)).to.deep.equal([2, 3]);
+			it.return();
+			expect(it.restart(5).nextBatch(10)).to.deep.equal([6, 7, 8, 9, 10]);
+			expect(stmt.busy).to.be.false;
+		});
+		it('should only restart a finished iterator', function () {
+			const stmt = this.db.prepare('SELECT b FROM entries ORDER BY rowid').raw();
+			const it = stmt.iterate();
+			it.nextBatch(2);
+			expect(() => it.restart()).to.throw(TypeError);
+			expect(it.nextBatch(2)).to.deep.equal([3, 4]);
+			it.return();
+			const other = stmt.iterate();
+			expect(() => it.restart()).to.throw(TypeError);
+			other.return();
+			expect(it.restart().nextBatch(1)).to.deep.equal([1]);
+			it.return();
+		});
+		it('should leave the iterator finished on bad bind parameters', function () {
+			const stmt = this.db.prepare('SELECT b FROM entries WHERE b > ?').raw();
+			const it = stmt.iterate(9);
+			expect(it.nextBatch(5)).to.deep.equal([10]);
+			expect(() => it.restart()).to.throw(RangeError);
+			expect(() => it.restart(1, 2)).to.throw(RangeError);
+			expect(stmt.busy).to.be.false;
+			expect(it.nextBatch(5)).to.deep.equal([]);
+			expect(stmt.all(9)).to.deep.equal([[10]]);
+		});
+		it('should use the statement\'s current modes', function () {
+			const stmt = this.db.prepare('SELECT a, b FROM entries ORDER BY rowid LIMIT 1');
+			const it = stmt.iterate();
+			expect(it.nextBatch(5)).to.deep.equal([{a: 'foo', b: 1}]);
+			stmt.raw().safeIntegers();
+			expect(it.restart().nextBatch(5)).to.deep.equal(['foo', 1n]);
+		});
+		it('should log the statement once per run in verbose mode', function () {
+			const logged = [];
+			const db = new Database(util.next(), {verbose: sql => logged.push(sql)});
+			const stmt = db.prepare('SELECT 1 UNION ALL SELECT 2').raw();
+			logged.length = 0;
+			const it = stmt.iterate();
+			it.nextBatch(5);
+			it.restart().nextBatch(5);
+			expect(logged).to.deep.equal([stmt.source, stmt.source]);
+			db.close();
+		});
+	});
+});