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 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, {