From 4fc80628ebe077d1d3cd728d2b77da86f50548e9 Mon Sep 17 00:00:00 2001 From: igoramf Date: Mon, 9 Nov 2026 17:23:57 -0300 Subject: [PATCH] fix(bigquery): carry job location into page 2+ (getQueryResults) The opaque pageToken kept only jobId + BQ token, and getQueryResults never sent location, so page 2 of any job outside the US/EU multi-regions (e.g. southamerica-east1) 404'd "Not found: Job". The token now carries jobReference.location (legacy 2-part tokens still decode), and the job-completion polling loops pass it too. --- .../server/lib/bigquery-client.ts | 10 ++++ .../server/lib/pagination-token.ts | 29 +++++++---- google-big-query/server/tools/bigquery.ts | 9 +++- .../tests/tools/bigquery-query.test.ts | 51 +++++++++++++++++++ 4 files changed, 88 insertions(+), 11 deletions(-) diff --git a/google-big-query/server/lib/bigquery-client.ts b/google-big-query/server/lib/bigquery-client.ts index 7b8606ad..0dad9aab 100644 --- a/google-big-query/server/lib/bigquery-client.ts +++ b/google-big-query/server/lib/bigquery-client.ts @@ -191,10 +191,16 @@ export class BigQueryClient { pageToken?: string; startIndex?: string; timeoutMs?: number; + /** Job location — required for jobs outside the US/EU multi-regions. */ + location?: string; }, ): Promise { const url = new URL(ENDPOINTS.QUERY_RESULTS(projectId, jobId)); + if (options?.location) { + url.searchParams.set("location", options.location); + } + if (options?.maxResults) { url.searchParams.set("maxResults", String(options.maxResults)); } @@ -258,6 +264,7 @@ export class BigQueryClient { projectId, response.jobReference.jobId, { + location: response.jobReference.location, timeoutMs: options.timeoutMs ?? DEFAULTS.QUERY_TIMEOUT_MS, }, ); @@ -308,6 +315,7 @@ export class BigQueryClient { totalBytesProcessed: string; jobId?: string; pageToken?: string; + location?: string; }> { let response = await this.query(projectId, options); @@ -325,6 +333,7 @@ export class BigQueryClient { projectId, response.jobReference.jobId, { + location: response.jobReference.location, timeoutMs: options.timeoutMs ?? DEFAULTS.QUERY_TIMEOUT_MS, maxResults: options.maxResults ?? DEFAULTS.MAX_RESULTS, }, @@ -351,6 +360,7 @@ export class BigQueryClient { totalBytesProcessed: response.totalBytesProcessed || "0", jobId: response.pageToken ? response.jobReference?.jobId : undefined, pageToken: response.pageToken, + location: response.jobReference?.location, }; } diff --git a/google-big-query/server/lib/pagination-token.ts b/google-big-query/server/lib/pagination-token.ts index bb08c0d9..d9b4bc30 100644 --- a/google-big-query/server/lib/pagination-token.ts +++ b/google-big-query/server/lib/pagination-token.ts @@ -1,20 +1,31 @@ // server/lib/pagination-token.ts /** - * Encodes a BigQuery jobId and API pageToken into a single opaque string. - * Format: base64(jobId + "\n" + apiPageToken) - * The newline is the separator — jobIds and pageTokens never contain newlines. + * Encodes a BigQuery jobId, API pageToken and job location into one opaque + * string. Format: base64(jobId + "\n" + apiPageToken [+ "\n" + location]) + * The newline is the separator — none of the three ever contains one. + * + * The location is required to fetch page 2+ of a job outside the US/EU + * multi-regions (e.g. southamerica-east1): jobs.getQueryResults without it + * looks the job up in US and 404s "Not found: Job". + * Tokens without a location (minted before this change) still decode. */ -export function encodePageToken(jobId: string, apiToken: string): string { - return btoa(`${jobId}\n${apiToken}`); +export function encodePageToken( + jobId: string, + apiToken: string, + location?: string, +): string { + return btoa( + location ? `${jobId}\n${apiToken}\n${location}` : `${jobId}\n${apiToken}`, + ); } export function decodePageToken(token: string): { jobId: string; apiToken: string; + location?: string; } { - const decoded = atob(token); - const nl = decoded.indexOf("\n"); - if (nl === -1) throw new Error("Invalid pageToken format"); - return { jobId: decoded.slice(0, nl), apiToken: decoded.slice(nl + 1) }; + const [jobId, apiToken, location] = atob(token).split("\n"); + if (apiToken === undefined) throw new Error("Invalid pageToken format"); + return { jobId, apiToken, ...(location ? { location } : {}) }; } diff --git a/google-big-query/server/tools/bigquery.ts b/google-big-query/server/tools/bigquery.ts index a4a55402..7877debe 100644 --- a/google-big-query/server/tools/bigquery.ts +++ b/google-big-query/server/tools/bigquery.ts @@ -148,14 +148,18 @@ export const createQueryTool = (env: Env) => totalBytesProcessed: string; jobId?: string; pageToken?: string; + location?: string; }; if (context.pageToken) { // Subsequent page: decode opaque token and call getQueryResults - const { jobId, apiToken } = decodePageToken(context.pageToken); + const { jobId, apiToken, location } = decodePageToken( + context.pageToken, + ); const page = await client.getQueryResults(context.projectId, jobId, { maxResults: context.maxResults, pageToken: apiToken, + location, }); if (!page.jobComplete) { throw new Error("Unexpected: job not complete on page fetch"); @@ -168,6 +172,7 @@ export const createQueryTool = (env: Env) => totalBytesProcessed: "0", jobId: page.pageToken ? jobId : undefined, pageToken: page.pageToken, + location, }; } else { // First page: execute query @@ -203,7 +208,7 @@ export const createQueryTool = (env: Env) => const nextPageToken = result.jobId && result.pageToken - ? encodePageToken(result.jobId, result.pageToken) + ? encodePageToken(result.jobId, result.pageToken, result.location) : undefined; return { diff --git a/google-big-query/tests/tools/bigquery-query.test.ts b/google-big-query/tests/tools/bigquery-query.test.ts index 950b4ccf..fd97a7f7 100644 --- a/google-big-query/tests/tools/bigquery-query.test.ts +++ b/google-big-query/tests/tools/bigquery-query.test.ts @@ -149,3 +149,54 @@ test("throws when neither query nor pageToken is provided", async () => { }), ).rejects.toThrow("query is required when pageToken is not provided"); }); + +test("job location round-trips through nextPageToken into page 2's getQueryResults", async () => { + // Non-US jobs (southamerica-east1) 404 on getQueryResults without location. + let capturedUrl = ""; + globalThis.fetch = mock((url: string | URL) => { + capturedUrl = url.toString(); + return Promise.resolve( + makeResponse({ + jobComplete: true, + jobReference: { + jobId: "jobSA", + projectId: "myproject", + location: "southamerica-east1", + }, + schema: SCHEMA, + rows: ROWS, + totalRows: "500", + totalBytesProcessed: "100", + cacheHit: false, + pageToken: "BQ_P2", + }), + ); + }) as unknown as typeof fetch; + + const { createQueryTool } = await import("../../server/tools/bigquery.ts"); + const tool = createQueryTool({ + MESH_REQUEST_CONTEXT: { authorization: "Bearer fake" }, + } as never); + + // @ts-ignore + const first = await tool.execute({ + context: { projectId: "myproject", query: "SELECT 1", maxResults: 1 }, + runtimeContext: {} as never, + }); + expect(decodePageToken(first.nextPageToken!).location).toBe( + "southamerica-east1", + ); + + // @ts-ignore + await tool.execute({ + context: { projectId: "myproject", pageToken: first.nextPageToken }, + runtimeContext: {} as never, + }); + expect(capturedUrl).toContain("/queries/jobSA"); + expect(capturedUrl).toContain("location=southamerica-east1"); +}); + +test("legacy token without location still decodes", () => { + const legacy = btoa("job1\ntok1"); + expect(decodePageToken(legacy)).toEqual({ jobId: "job1", apiToken: "tok1" }); +});