fix: wait for asynchronous row batch delivery

docs/content/docs/guides/sync-and-async.mdx

@@ -49,6 +49,8 @@ The [`NitroSQLite` export](/api/react-native-nitro-sqlite/variables/NitroSQLite)
 
 `NitroSQLite.native` exposes the underlying [`NitroSQLiteNative` hybrid object](/api/react-native-nitro-sqlite/hybrid-objects/NitroSQLiteNative). Its methods take a database name and return raw results without the connection helper's `rows` container. They also throw native errors without converting them to `NitroSQLiteError`. Its `executeAsync()` accepts an optional `onRows` callback that receives leading rows in batches, in order; the resolved `results` then holds only the rows after the last batch.
 
+The native promise can resolve before those callbacks run. Its `rowBatchCount` reports how many batches to expect. Wait for all of them before appending `results` to the collected rows. Without `onRows`, the count is zero and `results` contains every row. The managed `executeAsync()` API handles this waiting and assembly for you.
+
 ```ts
 import { NitroSQLite } from 'react-native-nitro-sqlite'
 

example/tests/unit/specs/operations/execute.spec.ts

@@ -168,6 +168,7 @@ export default function registerExecuteUnitTests() {
           )
 
           expectQueryResultRows(result, size)
+          expect(result.rowBatchCount).toBe(0)
         }
       } finally {
         db.close()
@@ -186,6 +187,7 @@ export default function registerExecuteUnitTests() {
           )
 
           expectQueryResultRows(result, size)
+          expect(result.rowBatchCount).toBe(Math.floor(size / 256))
           expect(result.results).toBe(result.rows._array)
           expect(
             result.rows._array.every((row, index) => row.id === index + 1),
@@ -198,6 +200,7 @@ export default function registerExecuteUnitTests() {
             [1_000],
           )
           expectQueryResultRows(result, 1_000)
+          expect(result.rowBatchCount).toBe(3)
           expect(
             result.rows._array.every((row, index) => row.id === index + 1),
           ).toBe(true)

packages/react-native-nitro-sqlite/cpp/NitroSQLiteOperations.cpp

@@ -221,6 +221,7 @@ namespace {
     int columnCount = 0;
     std::vector<std::string> columnNames;
     std::vector<SQLiteQueryResultRow> rows;
+    size_t rowBatchCount = 0;
     bool columnsCaptured = false;
 
     const auto captureColumns = [&] {
@@ -283,6 +284,7 @@ namespace {
       rows.push_back(std::move(row));
       if (onRows && rows.size() == kRowBatchSize) {
         onRows(SQLiteQueryResults(columnNames, std::move(rows)));
+        rowBatchCount++;
         rows = {};
       }
     });
@@ -309,6 +311,7 @@ namespace {
     result.insertId = static_cast<double>(sqlite3_last_insert_rowid(db));
     result.results = SQLiteQueryResults(std::move(columnNames), std::move(rows));
     result.metadata = std::move(metadata);
+    result.rowBatchCount = static_cast<double>(rowBatchCount);
     return result;
   }
 

packages/react-native-nitro-sqlite/cpp/NitroSQLiteOperations.hpp

@@ -54,7 +54,8 @@ void sqliteAttachDb(const SQLiteConnectionPtr& connection, const std::string& do
 void sqliteDetachDb(const SQLiteConnectionPtr& connection, const std::string& alias);
 
 /** Execute one statement. With @p onRows, full batches of rows go to it as they are read and the
- * returned result holds only the rows after the last batch.
+ * returned result holds only the rows after the last batch. Its rowBatchCount reports how many
+ * batches were sent; JavaScript delivery may still be pending when this function returns.
  */
 NitroSQLiteQueryResult sqliteExecute(const SQLiteConnectionPtr& connection, const std::string& query,
                                      const std::optional<SQLiteQueryParams>& params, const SQLiteRowBatchHandler& onRows = nullptr);

packages/react-native-nitro-sqlite/src/__tests__/execute.test.ts

@@ -10,7 +10,7 @@ import {
   executeAsyncManaged,
   executeManaged,
 } from '../operations/execute'
-import { nativeResult } from './testUtils'
+import { deferred, nativeResult } from './testUtils'
 
 const dbName = 'execute-test'
 const query = 'SELECT ? AS value'
@@ -127,6 +127,65 @@ describe('execute', () => {
     expect(result.rows.item(3)).toEqual({ value: 4 })
   })
 
+  it.each([
+    { size: 256, earlyBatches: 0, pendingNative: false },
+    { size: 257, earlyBatches: 0, pendingNative: true },
+    { size: 512, earlyBatches: 0, pendingNative: false },
+    { size: 1_000, earlyBatches: 0, pendingNative: false },
+    { size: 1_000, earlyBatches: 1, pendingNative: false },
+    { size: 1_000, earlyBatches: 3, pendingNative: false },
+  ])(
+    'waits for all batches with $size rows, $earlyBatches early batches, and pendingNative=$pendingNative',
+    async ({ size, earlyBatches, pendingNative }) => {
+      const rows = Array.from({ length: size }, (_, value) => ({ value }))
+      const rowBatchCount = Math.floor(size / 256)
+      const native = {
+        ...nativeResult(rows.slice(rowBatchCount * 256)),
+        rowBatchCount,
+      }
+      const completion = deferred<typeof native>()
+      const deliveries: (() => void)[] = []
+      jest
+        .mocked(HybridNitroSQLite.executeAsync)
+        .mockImplementationOnce((_dbName, _query, _params, onRows) => {
+          for (let batch = 0; batch < rowBatchCount; batch++) {
+            deliveries.push(() =>
+              onRows?.(rows.slice(batch * 256, (batch + 1) * 256)),
+            )
+          }
+          deliveries.splice(0, earlyBatches).forEach((deliver) => deliver())
+          return pendingNative ? completion.promise : Promise.resolve(native)
+        })
+
+      const completed = jest.fn()
+      const pending = executeAsync<{ value: number }>(
+        dbName,
+        query,
+        params,
+      ).then((result) => {
+        completed()
+        return result
+      })
+      completion.resolve(native)
+      await new Promise<void>((resolve) => setImmediate(resolve))
+      expect(completed).toHaveBeenCalledTimes(deliveries.length === 0 ? 1 : 0)
+
+      while (deliveries.length > 0) {
+        deliveries.shift()?.()
+        if (deliveries.length > 0) {
+          await new Promise<void>((resolve) => setImmediate(resolve))
+          expect(completed).not.toHaveBeenCalled()
+        }
+      }
+
+      const result = await pending
+      expect(result.results).toEqual(rows)
+      expect(result.rows._array).toBe(result.results)
+      expect(result.rows.length).toBe(size)
+      expect(result.rows.item(size - 1)).toEqual(rows[size - 1])
+    },
+  )
+
   it('keeps the native rows when no batch was delivered', async () => {
     const native = nativeResult([{ value: 7 }])
     jest.mocked(HybridNitroSQLite.executeAsync).mockResolvedValueOnce(native)
@@ -136,6 +195,26 @@ describe('execute', () => {
     expect(result.rows._array).toBe(native.results)
   })
 
+  it('keeps a managed connection busy until row delivery completes', async () => {
+    openDatabaseQueue(dbName)
+    const delivery = deferred<() => void>()
+    jest
+      .mocked(HybridNitroSQLite.executeAsync)
+      .mockImplementationOnce(async (_dbName, _query, _params, onRows) => {
+        delivery.resolve(() => onRows?.([{ value: 1 }]))
+        return { ...nativeResult([{ value: 2 }]), rowBatchCount: 1 }
+      })
+
+    const pending = executeAsync(dbName, query)
+    const deliver = await delivery.promise
+    await new Promise<void>((resolve) => setImmediate(resolve))
+    expect(() => execute(dbName, query)).toThrow('busy')
+
+    deliver()
+    await expect(pending).resolves.toMatchObject({ rows: { length: 2 } })
+    expect(() => closeDatabaseQueue(dbName)).not.toThrow()
+  })
+
   it('converts native asynchronous errors and releases a managed queue', async () => {
     openDatabaseQueue(dbName)
     jest

packages/react-native-nitro-sqlite/src/operations/execute.ts

@@ -57,7 +57,7 @@ export function executeNative<Row extends QueryResultRow = never>(
  * @param dbName Name of an open native database.
  * @param query SQL statement with optional positional placeholders.
  * @param params Values bound to the placeholders.
- * @returns A promise of the query result with typed rows.
+ * @returns A promise of the complete query result after all row batches arrive.
  */
 export async function executeAsync<Row extends QueryResultRow = never>(
   dbName: string,
@@ -90,6 +90,8 @@ export async function executeAsyncNative<Row extends QueryResultRow = never>(
   try {
     // Large results arrive in batches, so rows are converted while SQLite reads the rest.
     let leadingRows: Row[] | undefined
+    let receivedBatchCount = 0
+    let onBatchReceived: (() => void) | undefined
     const nativeResult = await HybridNitroSQLite.executeAsync(
       dbName,
       query,
@@ -100,8 +102,19 @@ export async function executeAsyncNative<Row extends QueryResultRow = never>(
         } else {
           Array.prototype.push.apply(leadingRows, rows as Row[])
         }
+        receivedBatchCount++
+        onBatchReceived?.()
       },
     )
+    // Native computation can finish before its queued JavaScript callbacks run.
+    const expectedBatchCount = nativeResult.rowBatchCount ?? 0
+    if (receivedBatchCount < expectedBatchCount) {
+      await new Promise<void>((resolve) => {
+        onBatchReceived = () => {
+          if (receivedBatchCount >= expectedBatchCount) resolve()
+        }
+      })
+    }
     return buildJSQueryResult<Row>(nativeResult, leadingRows)
   } catch (error) {
     throw NitroSQLiteError.fromError(error)

packages/react-native-nitro-sqlite/src/specs/NitroSQLite.nitro.ts

@@ -100,7 +100,8 @@ export interface NitroSQLite
    * @param onRows Receives leading rows in batches, in order, while SQLite still reads the rest, so
    * JavaScript can convert a batch while the next one is read. Batched rows are not repeated in the
    * result's `results`, which then holds only the rows after the last batch.
-   * @returns A promise of the native query result.
+   * @returns A promise of the native query result. Row callbacks may still be queued when it
+   * resolves. Wait for `rowBatchCount` callbacks before combining their rows with `results`.
    */
   executeAsync(
     dbName: string,

packages/react-native-nitro-sqlite/src/specs/NitroSQLiteQueryResult.nitro.ts

@@ -22,6 +22,13 @@ export interface NitroSQLiteQueryResult {
 
   /** Column metadata keyed by result column name, when available. */
   readonly metadata?: Record<string, NitroSQLiteQueryColumnMetadata>
+
+  /** Number of row batches sent through the native `onRows` callback, or zero when none were sent.
+   * Callbacks may arrive after the native promise resolves. Wait for this many batches before
+   * combining them with `results`. The managed async API waits for them automatically.
+   * An omitted count is treated as zero. Batch tracking requires an updated native build.
+   */
+  readonly rowBatchCount?: number
 }
 
 // TODO: Investigate why this cannot be represented in Nitro