diff --git a/js/package.json b/js/package.json index 58394dff3..64176c787 100644 --- a/js/package.json +++ b/js/package.json @@ -99,6 +99,7 @@ "async": "^3.2.5", "autoevals": "^0.0.131", "cross-env": "^7.0.3", + "jiti": "^2.6.1", "npm-run-all": "^4.1.5", "openapi-zod-client": "^1.18.3", "prettier": "^3.5.3", diff --git a/js/src/logger.ts b/js/src/logger.ts index a5c63f0dc..1e03f4cda 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -5477,6 +5477,11 @@ export class ObjectFetcher const state = await this.getState(); const objectId = await this.id; const limit = batchSize ?? DEFAULT_FETCH_BATCH_SIZE; + const internalBtqlWithoutCursor = Object.fromEntries( + Object.entries(this._internal_btql ?? {}).filter( + ([key]) => key !== "cursor", + ), + ); let cursor = undefined; let iterations = 0; while (true) { @@ -5484,7 +5489,6 @@ export class ObjectFetcher `btql`, { query: { - ...this._internal_btql, select: [ { op: "star", @@ -5505,6 +5509,7 @@ export class ObjectFetcher }, cursor, limit, + ...internalBtqlWithoutCursor, }, use_columnstore: false, brainstore_realtime: true, diff --git a/js/src/object-fetcher.test.ts b/js/src/object-fetcher.test.ts new file mode 100644 index 000000000..ac2c8bd66 --- /dev/null +++ b/js/src/object-fetcher.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, test, vi } from "vitest"; +import { + DEFAULT_FETCH_BATCH_SIZE, + ObjectFetcher, + type BraintrustState, +} from "./logger"; +import { configureNode } from "./node"; + +configureNode(); + +type TestRecord = { id: string }; + +type MockBtqlResponse = { + data: Array>; + cursor?: string | null; +}; + +function createPostResponse(response: MockBtqlResponse) { + return { + json: vi.fn().mockResolvedValue(response), + }; +} + +function createPostMock( + response: MockBtqlResponse = { data: [], cursor: null }, +) { + return vi.fn().mockResolvedValue(createPostResponse(response)); +} + +class TestObjectFetcher extends ObjectFetcher { + constructor( + private readonly postMock: ReturnType, + internalBtql?: Record, + ) { + super("dataset", undefined, undefined, internalBtql); + } + + public get id(): Promise { + return Promise.resolve("test-dataset-id"); + } + + protected async getState(): Promise { + return { + apiConn: () => ({ + post: this.postMock, + }), + } as unknown as BraintrustState; + } +} + +async function triggerFetch( + fetcher: TestObjectFetcher, + options?: { batchSize?: number }, +) { + await fetcher.fetchedData(options); +} + +function getBtqlQuery( + postMock: ReturnType, + callIndex = 0, +) { + const call = postMock.mock.calls[callIndex]; + expect(call).toBeDefined(); + const requestBody = call[1] as { query: Record }; + return requestBody.query; +} + +describe("ObjectFetcher internal BTQL limit handling", () => { + test("preserves custom _internal_btql limit instead of default batch size", async () => { + const postMock = createPostMock(); + const fetcher = new TestObjectFetcher(postMock, { + limit: 50, + where: { op: "eq", left: "foo", right: "bar" }, + }); + + await triggerFetch(fetcher); + + expect(postMock).toHaveBeenCalledTimes(1); + const query = getBtqlQuery(postMock); + expect(query.limit).toBe(50); + expect(query.where).toEqual({ op: "eq", left: "foo", right: "bar" }); + }); + + test("uses default batch size when no _internal_btql limit is provided", async () => { + const postMock = createPostMock(); + const fetcher = new TestObjectFetcher(postMock); + + await triggerFetch(fetcher); + + const query = getBtqlQuery(postMock); + expect(query.limit).toBe(DEFAULT_FETCH_BATCH_SIZE); + }); + + test("uses explicit fetch batchSize when no _internal_btql limit is provided", async () => { + const postMock = createPostMock(); + const fetcher = new TestObjectFetcher(postMock); + + await triggerFetch(fetcher, { batchSize: 17 }); + + const query = getBtqlQuery(postMock); + expect(query.limit).toBe(17); + }); + + test("does not allow _internal_btql cursor to override pagination cursor", 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, { + cursor: "stale-cursor", + limit: 1, + }); + + await triggerFetch(fetcher); + + expect(postMock).toHaveBeenCalledTimes(2); + const firstQuery = getBtqlQuery(postMock, 0); + const secondQuery = getBtqlQuery(postMock, 1); + expect(firstQuery.cursor).toBeUndefined(); + expect(secondQuery.cursor).toBe("next-page-cursor"); + }); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index f2a5502d1..75e86c94c 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -40,7 +40,7 @@ importers: version: 2.6.6(@types/node@20.10.5)(typescript@5.3.3) tsup: specifier: ^8.3.5 - version: 8.3.5(@swc/core@1.15.8)(postcss@8.5.6)(typescript@5.3.3)(yaml@2.8.2) + version: 8.3.5(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(typescript@5.3.3)(yaml@2.8.2) typedoc: specifier: ^0.25.13 version: 0.25.13(typescript@5.3.3) @@ -70,7 +70,7 @@ importers: version: link:../../js tsup: specifier: ^8.3.5 - version: 8.3.5(@swc/core@1.15.8)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2) + version: 8.3.5(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2) tsx: specifier: ^3.14.0 version: 3.14.0 @@ -110,7 +110,7 @@ importers: version: link:../../js tsup: specifier: ^8.5.0 - version: 8.5.1(@swc/core@1.15.8)(postcss@8.5.6)(typescript@5.5.4)(yaml@2.8.2) + version: 8.5.1(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(typescript@5.5.4)(yaml@2.8.2) typedoc: specifier: ^0.28.15 version: 0.28.15(typescript@5.5.4) @@ -143,7 +143,7 @@ importers: version: link:../../js tsup: specifier: ^8.5.0 - version: 8.5.1(@swc/core@1.15.8)(postcss@8.5.6)(typescript@5.5.4)(yaml@2.8.2) + version: 8.5.1(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(typescript@5.5.4)(yaml@2.8.2) typescript: specifier: 5.5.4 version: 5.5.4 @@ -168,7 +168,7 @@ importers: version: 20.10.5 tsup: specifier: ^8.3.5 - version: 8.3.5(@swc/core@1.15.8)(postcss@8.5.6)(typescript@5.3.3)(yaml@2.8.2) + version: 8.3.5(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(typescript@5.3.3)(yaml@2.8.2) typescript: specifier: ^5.3.3 version: 5.3.3 @@ -325,10 +325,10 @@ importers: version: 9.0.7 '@typescript-eslint/eslint-plugin': specifier: ^8.49.0 - version: 8.50.0(@typescript-eslint/parser@8.50.0(eslint@9.39.2)(typescript@5.4.4))(eslint@9.39.2)(typescript@5.4.4) + version: 8.50.0(@typescript-eslint/parser@8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4))(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4) '@typescript-eslint/parser': specifier: ^8.49.0 - version: 8.50.0(eslint@9.39.2)(typescript@5.4.4) + version: 8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4) ai: specifier: ^6.0.0 version: 6.0.37(zod@3.25.76) @@ -341,6 +341,9 @@ importers: cross-env: specifier: ^7.0.3 version: 7.0.3 + jiti: + specifier: ^2.6.1 + version: 2.6.1 npm-run-all: specifier: ^4.1.5 version: 4.1.5 @@ -361,7 +364,7 @@ importers: version: 29.1.4(@babel/core@7.28.0)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.28.0))(esbuild@0.27.0)(jest@29.7.0(@types/node@20.10.5)(ts-node@10.9.2(@swc/core@1.15.8)(@types/node@20.10.5)(typescript@5.4.4)))(typescript@5.4.4) tsup: specifier: ^8.5.1 - version: 8.5.1(@swc/core@1.15.8)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2) + version: 8.5.1(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2) tsx: specifier: ^3.14.0 version: 3.14.0 @@ -3611,11 +3614,12 @@ packages: glob@10.4.5: resolution: {integrity: sha512-7Bv8RF0k6xjo7d4A/PxYLbUCfb6c+Vpd2/mB2yRDlew7Jb5hEXiCD9ibfO7wpk8i4sevK6DFny9h7EYbM3/sHg==} + deprecated: Old versions of glob are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me hasBin: true glob@7.2.3: resolution: {integrity: sha512-nFR0zLpU2YCaRxwoCJvL6UvCH2JFyFVIvwTLsIf21AuHlMskA1hhTdk+LlYJtOlYt9v6dvszD2BGRqBL+iQK9Q==} - deprecated: Glob versions prior to v9 are no longer supported + deprecated: Old versions of glob are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me globals@11.12.0: resolution: {integrity: sha512-WOBp/EEGUiIsJSp7wcv/y6MO+lV9UoncWqxuFfm8eBwzWNgyfBd6Gz+IeKQ9jCmyhoH99g15M3T+QaVHFjizVA==} @@ -4041,6 +4045,10 @@ packages: node-notifier: optional: true + jiti@2.6.1: + resolution: {integrity: sha512-ekilCSN1jwRvIbgeg/57YFh8qQDNbwDb9xT/qu2DAHbFFZUicIl4ygVaAvzveMhMVr3LnpSKTNnwt8PoOfmKhQ==} + hasBin: true + joycon@3.1.1: resolution: {integrity: sha512-34wB/Y7MW7bzjKRjUKTa46I2Z7eV62Rkhva+KkopW7Qvv/OSWBqvkSY7vusOPrNuZcUG3tApvdVgNB8POj3SPw==} engines: {node: '>=10'} @@ -5221,6 +5229,7 @@ packages: tar@7.5.2: resolution: {integrity: sha512-7NyxrTE4Anh8km8iEy7o0QYPs+0JKBTj5ZaqHg6B39erLg0qYXN3BijtShwbsNSvQ+LN75+KV+C4QR/f6Gwnpg==} engines: {node: '>=18'} + deprecated: Old versions of tar are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me termi-link@1.1.0: resolution: {integrity: sha512-2qSN6TnomHgVLtk+htSWbaYs4Rd2MH/RU7VpHTy6MBstyNyWbM4yKd1DCYpE3fDg8dmGWojXCngNi/MHCzGuAA==} @@ -6888,9 +6897,9 @@ snapshots: '@esbuild/win32-x64@0.27.0': optional: true - '@eslint-community/eslint-utils@4.9.0(eslint@9.39.2)': + '@eslint-community/eslint-utils@4.9.0(eslint@9.39.2(jiti@2.6.1))': dependencies: - eslint: 9.39.2 + eslint: 9.39.2(jiti@2.6.1) eslint-visitor-keys: 3.4.3 '@eslint-community/regexpp@4.12.2': {} @@ -8133,15 +8142,15 @@ snapshots: dependencies: '@types/yargs-parser': 21.0.3 - '@typescript-eslint/eslint-plugin@8.50.0(@typescript-eslint/parser@8.50.0(eslint@9.39.2)(typescript@5.4.4))(eslint@9.39.2)(typescript@5.4.4)': + '@typescript-eslint/eslint-plugin@8.50.0(@typescript-eslint/parser@8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4))(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4)': dependencies: '@eslint-community/regexpp': 4.12.2 - '@typescript-eslint/parser': 8.50.0(eslint@9.39.2)(typescript@5.4.4) + '@typescript-eslint/parser': 8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4) '@typescript-eslint/scope-manager': 8.50.0 - '@typescript-eslint/type-utils': 8.50.0(eslint@9.39.2)(typescript@5.4.4) - '@typescript-eslint/utils': 8.50.0(eslint@9.39.2)(typescript@5.4.4) + '@typescript-eslint/type-utils': 8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4) + '@typescript-eslint/utils': 8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4) '@typescript-eslint/visitor-keys': 8.50.0 - eslint: 9.39.2 + eslint: 9.39.2(jiti@2.6.1) ignore: 7.0.5 natural-compare: 1.4.0 ts-api-utils: 2.1.0(typescript@5.4.4) @@ -8149,14 +8158,14 @@ snapshots: transitivePeerDependencies: - supports-color - '@typescript-eslint/parser@8.50.0(eslint@9.39.2)(typescript@5.4.4)': + '@typescript-eslint/parser@8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4)': dependencies: '@typescript-eslint/scope-manager': 8.50.0 '@typescript-eslint/types': 8.50.0 '@typescript-eslint/typescript-estree': 8.50.0(typescript@5.4.4) '@typescript-eslint/visitor-keys': 8.50.0 debug: 4.4.3 - eslint: 9.39.2 + eslint: 9.39.2(jiti@2.6.1) typescript: 5.4.4 transitivePeerDependencies: - supports-color @@ -8179,13 +8188,13 @@ snapshots: dependencies: typescript: 5.4.4 - '@typescript-eslint/type-utils@8.50.0(eslint@9.39.2)(typescript@5.4.4)': + '@typescript-eslint/type-utils@8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4)': dependencies: '@typescript-eslint/types': 8.50.0 '@typescript-eslint/typescript-estree': 8.50.0(typescript@5.4.4) - '@typescript-eslint/utils': 8.50.0(eslint@9.39.2)(typescript@5.4.4) + '@typescript-eslint/utils': 8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4) debug: 4.4.3 - eslint: 9.39.2 + eslint: 9.39.2(jiti@2.6.1) ts-api-utils: 2.1.0(typescript@5.4.4) typescript: 5.4.4 transitivePeerDependencies: @@ -8208,13 +8217,13 @@ snapshots: transitivePeerDependencies: - supports-color - '@typescript-eslint/utils@8.50.0(eslint@9.39.2)(typescript@5.4.4)': + '@typescript-eslint/utils@8.50.0(eslint@9.39.2(jiti@2.6.1))(typescript@5.4.4)': dependencies: - '@eslint-community/eslint-utils': 4.9.0(eslint@9.39.2) + '@eslint-community/eslint-utils': 4.9.0(eslint@9.39.2(jiti@2.6.1)) '@typescript-eslint/scope-manager': 8.50.0 '@typescript-eslint/types': 8.50.0 '@typescript-eslint/typescript-estree': 8.50.0(typescript@5.4.4) - eslint: 9.39.2 + eslint: 9.39.2(jiti@2.6.1) typescript: 5.4.4 transitivePeerDependencies: - supports-color @@ -9364,9 +9373,9 @@ snapshots: eslint-visitor-keys@4.2.1: {} - eslint@9.39.2: + eslint@9.39.2(jiti@2.6.1): dependencies: - '@eslint-community/eslint-utils': 4.9.0(eslint@9.39.2) + '@eslint-community/eslint-utils': 4.9.0(eslint@9.39.2(jiti@2.6.1)) '@eslint-community/regexpp': 4.12.2 '@eslint/config-array': 0.21.1 '@eslint/config-helpers': 0.4.2 @@ -9400,6 +9409,8 @@ snapshots: minimatch: 3.1.2 natural-compare: 1.4.0 optionator: 0.9.4 + optionalDependencies: + jiti: 2.6.1 transitivePeerDependencies: - supports-color @@ -10425,6 +10436,8 @@ snapshots: - supports-color - ts-node + jiti@2.6.1: {} + joycon@3.1.1: {} js-levenshtein@1.1.6: {} @@ -11167,10 +11180,11 @@ snapshots: pluralize@8.0.0: {} - postcss-load-config@6.0.1(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2): + postcss-load-config@6.0.1(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2): dependencies: lilconfig: 3.1.3 optionalDependencies: + jiti: 2.6.1 postcss: 8.5.6 tsx: 3.14.0 yaml: 2.8.2 @@ -11894,7 +11908,7 @@ snapshots: tslib@2.8.1: {} - tsup@8.3.5(@swc/core@1.15.8)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2): + tsup@8.3.5(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2): dependencies: bundle-require: 5.1.0(esbuild@0.24.2) cac: 6.7.14 @@ -11904,7 +11918,7 @@ snapshots: esbuild: 0.24.2 joycon: 3.1.1 picocolors: 1.1.1 - postcss-load-config: 6.0.1(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) + postcss-load-config: 6.0.1(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) resolve-from: 5.0.0 rollup: 4.35.0 source-map: 0.8.0-beta.0 @@ -11922,7 +11936,7 @@ snapshots: - tsx - yaml - tsup@8.3.5(@swc/core@1.15.8)(postcss@8.5.6)(typescript@5.3.3)(yaml@2.8.2): + tsup@8.3.5(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(typescript@5.3.3)(yaml@2.8.2): dependencies: bundle-require: 5.1.0(esbuild@0.24.2) cac: 6.7.14 @@ -11932,7 +11946,7 @@ snapshots: esbuild: 0.24.2 joycon: 3.1.1 picocolors: 1.1.1 - postcss-load-config: 6.0.1(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) + postcss-load-config: 6.0.1(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) resolve-from: 5.0.0 rollup: 4.35.0 source-map: 0.8.0-beta.0 @@ -11950,7 +11964,7 @@ snapshots: - tsx - yaml - tsup@8.5.1(@swc/core@1.15.8)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2): + tsup@8.5.1(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(typescript@5.4.4)(yaml@2.8.2): dependencies: bundle-require: 5.1.0(esbuild@0.27.0) cac: 6.7.14 @@ -11961,7 +11975,7 @@ snapshots: fix-dts-default-cjs-exports: 1.0.1 joycon: 3.1.1 picocolors: 1.1.1 - postcss-load-config: 6.0.1(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) + postcss-load-config: 6.0.1(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) resolve-from: 5.0.0 rollup: 4.35.0 source-map: 0.7.6 @@ -11979,7 +11993,7 @@ snapshots: - tsx - yaml - tsup@8.5.1(@swc/core@1.15.8)(postcss@8.5.6)(typescript@5.5.4)(yaml@2.8.2): + tsup@8.5.1(@swc/core@1.15.8)(jiti@2.6.1)(postcss@8.5.6)(typescript@5.5.4)(yaml@2.8.2): dependencies: bundle-require: 5.1.0(esbuild@0.27.0) cac: 6.7.14 @@ -11990,7 +12004,7 @@ snapshots: fix-dts-default-cjs-exports: 1.0.1 joycon: 3.1.1 picocolors: 1.1.1 - postcss-load-config: 6.0.1(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) + postcss-load-config: 6.0.1(jiti@2.6.1)(postcss@8.5.6)(tsx@3.14.0)(yaml@2.8.2) resolve-from: 5.0.0 rollup: 4.35.0 source-map: 0.7.6