diff --git a/docs/api.md b/docs/api.md index 25de1a0..e76d9a6 100644 --- a/docs/api.md +++ b/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)* diff --git a/lib/index.d.ts b/lib/index.d.ts index d48eba8..5240e7d 100644 --- a/lib/index.d.ts +++ b/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; + iterate(...params: BindParameters): StatementIterator; 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 extends IterableIterator { + /** + * 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 : BetterSqlite3.Statement<[BindParameters], Result>; type ColumnDefinition = BetterSqlite3.ColumnDefinition; + type StatementIterator = BetterSqlite3.StatementIterator; type Transaction = BetterSqlite3.Transaction; type Database = BetterSqlite3.Database; } diff --git a/src/objects/statement-iterator.cpp b/src/objects/statement-iterator.cpp index ff69021..133c418 100644 --- a/src/objects/statement-iterator.cpp +++ b/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 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 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 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(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(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(info.This()); REQUIRE_DATABASE_NOT_BUSY(iter->db_state); diff --git a/src/objects/statement-iterator.hpp b/src/objects/statement-iterator.hpp index 55caad7..ac7a754 100644 --- a/src/objects/statement-iterator.hpp +++ b/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; }; diff --git a/test/29.statement.iterate.next-batch.js b/test/29.statement.iterate.next-batch.js new file mode 100644 index 0000000..02c772a --- /dev/null +++ b/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(); + }); + }); +});