From 1d2e8ffad412420d6aadb813b9a09977d46149e1 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 5 Aug 2026 11:25:09 -0400 Subject: [PATCH 1/2] fix(dataset): enforce cumulative btql fetch limits The object fetcher treated `_internal_btql.limit` as a per-request page size instead of a total result cap. It yielded every response and kept following cursors, so fully consuming a dataset or calling `fetchedData()` could return more records than requested. Before: limit=3 -> page(2, cursor) -> page(2) -> 4 rows After: limit=3 -> page(2, cursor) -> page(limit=1) -> 3 rows Track the remaining record budget across requests, cap each page by both the batch size and that budget, and stop yielding or paginating when it is exhausted. Add fake API coverage for cursor termination, oversized responses, and the interaction between `batchSize` and the cumulative limit. --- js/src/logger.ts | 19 +++++++++-- js/src/object-fetcher.test.ts | 64 +++++++++++++++++++++++++++++------ 2 files changed, 69 insertions(+), 14 deletions(-) diff --git a/js/src/logger.ts b/js/src/logger.ts index 8bdb8a67b..00d801a4a 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -6927,8 +6927,7 @@ export class ObjectFetcher implements AsyncIterable< const objectId = await this.id; const batchLimit = batchSize ?? DEFAULT_FETCH_BATCH_SIZE; const internalLimit = getInternalBtqlLimit(this._internal_btql); - const limit = - batchSize !== undefined ? batchSize : (internalLimit ?? batchLimit); + let remainingLimit = internalLimit; const internalBtqlWithoutReservedQueryKeys = Object.fromEntries( Object.entries(this._internal_btql ?? {}).filter( ([key]) => @@ -6941,6 +6940,13 @@ export class ObjectFetcher implements AsyncIterable< let cursor = undefined; let iterations = 0; while (true) { + if (remainingLimit !== undefined && remainingLimit <= 0) { + return; + } + const limit = + remainingLimit === undefined + ? batchLimit + : Math.min(batchLimit, remainingLimit); const resp = await state.apiConn().post( `btql`, { @@ -6982,9 +6988,16 @@ export class ObjectFetcher implements AsyncIterable< const respJson = await resp.json(); const mutate = this.mutateRecord; for (const record of respJson.data ?? []) { - yield mutate + if (remainingLimit !== undefined && remainingLimit <= 0) { + return; + } + const mutatedRecord = mutate ? mutate(record) : (record as WithTransactionId); + if (remainingLimit !== undefined) { + remainingLimit--; + } + yield mutatedRecord; } if (!respJson.cursor) { break; diff --git a/js/src/object-fetcher.test.ts b/js/src/object-fetcher.test.ts index 43d9d751b..5c6936cdd 100644 --- a/js/src/object-fetcher.test.ts +++ b/js/src/object-fetcher.test.ts @@ -119,16 +119,6 @@ describe("ObjectFetcher internal BTQL limit handling", () => { expect(query.limit).toBe(17); }); - test("explicit batchSize overrides _internal_btql.limit", async () => { - const postMock = createPostMock(); - const fetcher = new TestObjectFetcher(postMock, { limit: 100 }); - - await triggerFetch(fetcher, { batchSize: 25 }); - - const query = getBtqlQuery(postMock); - expect(query.limit).toBe(25); - }); - test("does not allow _internal_btql cursor to override pagination cursor", async () => { const postMock = vi .fn() @@ -146,7 +136,7 @@ describe("ObjectFetcher internal BTQL limit handling", () => { ); const fetcher = new TestObjectFetcher(postMock, { cursor: "stale-cursor", - limit: 1, + limit: 2, }); await triggerFetch(fetcher); @@ -158,6 +148,58 @@ describe("ObjectFetcher internal BTQL limit handling", () => { expect(secondQuery.cursor).toBe("next-page-cursor"); }); + test("stops pagination once the cumulative _internal_btql limit is reached", async () => { + const postMock = vi + .fn() + .mockResolvedValueOnce( + createPostResponse({ + data: [{ id: "record-1" }], + cursor: "next-page-cursor", + }), + ) + .mockResolvedValueOnce( + createPostResponse({ + data: [{ id: "record-2" }], + cursor: null, + }), + ); + const fetcher = new TestObjectFetcher(postMock, { limit: 1 }); + + const records = await fetcher.fetchedData(); + + expect(records).toEqual([{ id: "record-1" }]); + expect(postMock).toHaveBeenCalledTimes(1); + }); + + test("combines batchSize with the cumulative _internal_btql limit", async () => { + const postMock = vi + .fn() + .mockResolvedValueOnce( + createPostResponse({ + data: [{ id: "record-1" }, { id: "record-2" }], + cursor: "next-page-cursor", + }), + ) + .mockResolvedValueOnce( + createPostResponse({ + data: [{ id: "record-3" }, { id: "record-4" }], + cursor: null, + }), + ); + const fetcher = new TestObjectFetcher(postMock, { limit: 3 }); + + const records = await fetcher.fetchedData({ batchSize: 2 }); + + expect(records).toEqual([ + { id: "record-1" }, + { id: "record-2" }, + { id: "record-3" }, + ]); + expect(postMock).toHaveBeenCalledTimes(2); + expect(getBtqlQuery(postMock, 0).limit).toBe(2); + expect(getBtqlQuery(postMock, 1).limit).toBe(1); + }); + test("does not allow _internal_btql select/from to override base object query", async () => { const postMock = createPostMock(); const fetcher = new TestObjectFetcher(postMock, { From 2a73916b5a577978e55b4d4e17a70fe665b985fc Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 5 Aug 2026 11:30:51 -0400 Subject: [PATCH 2/2] chore: add dataset btql limit changeset --- .changeset/dataset-cumulative-btql-limit.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/dataset-cumulative-btql-limit.md diff --git a/.changeset/dataset-cumulative-btql-limit.md b/.changeset/dataset-cumulative-btql-limit.md new file mode 100644 index 000000000..7160d45bc --- /dev/null +++ b/.changeset/dataset-cumulative-btql-limit.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(dataset): Enforce `_internal_btql.limit` across paginated fetches