From 4fffad81c17c83a86aa684b0ca2fd8052ee106ac Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 8 Oct 2026 16:05:15 +0200 Subject: [PATCH 1/4] feat(shared): migrate genie to modular @databricks/sdk-genie Route the facade's genie accessor through a lazily built modular GenieClient (sdk-genie 0.54.0, compatible with sdk-core/auth/options 0.51.0) and port the genie connector to its camelCase model. - Poll getConversationMessage ourselves instead of using the SDK waiter: the modular wait() has no onProgress, which the SSE status events need. Cadence, timeout and the "failed to reach COMPLETED state" errors match the legacy waiter. - Convert the query statement response back to the snake_case wire shape with bigint counts as numbers, so the SSE contract and JSON.stringify keep working. - Patch sdk-genie: its model types data_array as protobuf ListValue[], so the unmarshal rejected every real JSON_ARRAY query result. - Match RESOURCE_DOES_NOT_EXIST on ApiError .code / .errorCode too. The CLI's listSpaces/getSpace calls stay on the legacy client. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../server/testing-kit.integration.test.ts | 2 +- .../api/appkit/Interface.WorkspaceClient.md | 4 +- docs/docs/plugins/testing.md | 4 +- knip.json | 1 + packages/appkit/package.json | 1 + .../appkit/src/connectors/genie/client.ts | 212 ++++++--- .../src/connectors/genie/tests/client.test.ts | 352 ++++++++------- .../src/plugins/genie/tests/genie.test.ts | 417 ++++++++++-------- .../tests/mock-workspace-client.test.ts | 2 +- packages/appkit/src/workspace-client/index.ts | 9 +- packages/shared/package.json | 1 + .../shared/src/workspace-client/client.ts | 11 +- packages/shared/src/workspace-client/index.ts | 2 - .../shared/src/workspace-client/legacy.ts | 8 +- .../shared/src/workspace-client/modular.ts | 50 +++ .../workspace-client/tests/modular.test.ts | 57 ++- packages/shared/src/workspace-client/types.ts | 7 +- patches/@databricks__sdk-genie@0.54.0.patch | 26 ++ pnpm-lock.yaml | 20 + pnpm-workspace.yaml | 1 + 20 files changed, 743 insertions(+), 444 deletions(-) create mode 100644 patches/@databricks__sdk-genie@0.54.0.patch diff --git a/apps/dev-playground/server/testing-kit.integration.test.ts b/apps/dev-playground/server/testing-kit.integration.test.ts index 5d0443eca..fa644ba4f 100644 --- a/apps/dev-playground/server/testing-kit.integration.test.ts +++ b/apps/dev-playground/server/testing-kit.integration.test.ts @@ -68,7 +68,7 @@ describe("createMockWorkspaceClient — build the fake client yourself", () => { }); expect(getMock(client, "jobs.getRun")).toHaveBeenCalledWith({ run_id: 1 }); await expect( - client.genie.getMessage({ id: "m-1" }), + client.genie.genieGetConversationMessage({ messageId: "m-1" }), ).resolves.toBeUndefined(); }); }); diff --git a/docs/docs/api/appkit/Interface.WorkspaceClient.md b/docs/docs/api/appkit/Interface.WorkspaceClient.md index 26a680581..9d3fec157 100644 --- a/docs/docs/api/appkit/Interface.WorkspaceClient.md +++ b/docs/docs/api/appkit/Interface.WorkspaceClient.md @@ -56,10 +56,10 @@ UC Volumes / Files API. ### genie ```ts -readonly genie: GenieService; +readonly genie: GenieClient; ``` -Genie / dashboards. +Genie (modular SDK). *** diff --git a/docs/docs/plugins/testing.md b/docs/docs/plugins/testing.md index 6d06d355b..7a7797c01 100644 --- a/docs/docs/plugins/testing.md +++ b/docs/docs/plugins/testing.md @@ -343,7 +343,7 @@ await mock.attach(plugin); ``` `options` is: -- `responses` — seed the mock workspace client with responses keyed by dotted path (`"jobs.getRun"`, `"genie.getMessage"`). A value can be static or a function of call arguments and the abort signal. +- `responses` — seed the mock workspace client with responses keyed by dotted path (`"jobs.getRun"`, `"genie.genieGetConversationMessage"`). A value can be static or a function of call arguments and the abort signal. - `env` — set environment variables scoped to the test; they are restored on plugin detach. - `strict` — throw if a handler calls an undeclared workspace-client path (instead of silently resolving `undefined`). The built-in defaults still count as declared. @@ -513,7 +513,7 @@ const client = createMockWorkspaceClient({ }); await client.jobs.getRun({ run_id: 1 }); // → { state: "TERMINATED" } -await client.genie.getMessage({ id: "m-1" }); // → undefined, does not throw +await client.genie.genieGetConversationMessage({ messageId: "m-1" }); // → undefined, does not throw ``` `createTestApp` installs one of these for you, so reach for it directly only when you're driving a plugin through `createTestPluginContext` or `mockServiceContext`. diff --git a/knip.json b/knip.json index b46121129..c57c8be24 100644 --- a/knip.json +++ b/knip.json @@ -13,6 +13,7 @@ "@databricks/sdk-auth", "@databricks/sdk-core", "@databricks/sdk-experimental", + "@databricks/sdk-genie", "@databricks/sdk-options", "@databricks/sdk-scim", "@databricks/sdk-statementexecution", diff --git a/packages/appkit/package.json b/packages/appkit/package.json index a00cc0c2e..3eda04dba 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -74,6 +74,7 @@ "@databricks/sdk-auth": "0.51.0", "@databricks/sdk-core": "0.51.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-options": "0.51.0", "@databricks/sdk-scim": "0.51.0", "@databricks/sdk-statementexecution": "0.52.0", diff --git a/packages/appkit/src/connectors/genie/client.ts b/packages/appkit/src/connectors/genie/client.ts index ec8a987d1..ed77a5998 100644 --- a/packages/appkit/src/connectors/genie/client.ts +++ b/packages/appkit/src/connectors/genie/client.ts @@ -1,13 +1,7 @@ import { createLogger } from "../../logging"; -import { - type GenieMessage, - Time, - TimeUnits, - type Waiter, - type WorkspaceClient, -} from "../../workspace-client"; +import type { GenieMessage, WorkspaceClient } from "../../workspace-client"; import { genieConnectorDefaults } from "./defaults"; -import { pollWaiter } from "./poll-waiter"; +import { type Pollable, pollWaiter } from "./poll-waiter"; import type { GenieAttachmentResponse, GenieConversationHistoryResponse, @@ -26,7 +20,11 @@ const GenieErrors = { QUERY_RESULT_FAILED: "Failed to fetch query result", } as const; -type CreateMessageWaiter = Waiter; +type CreateMessageWaiter = Pollable; + +// Legacy SDK waiter defaults, kept so polling cadence is unchanged. +const DEFAULT_WAIT_TIMEOUT_MS = 10 * 60_000; +const MAX_POLL_INTERVAL_MS = 10_000; interface GenieConnectorConfig { timeout?: number; @@ -35,38 +33,72 @@ interface GenieConnectorConfig { function mapAttachments(message: GenieMessage): GenieAttachmentResponse[] { return ( - message.attachments?.map((att) => ({ - attachmentId: att.attachment_id, - query: att.query - ? { - title: att.query.title, - description: att.query.description, - query: att.query.query, - statementId: att.query.statement_id, - } - : undefined, - text: att.text ? { content: att.text.content } : undefined, - suggestedQuestions: att.suggested_questions?.questions, + message.attachments?.map(({ attachmentId, attachment }) => ({ + attachmentId, + query: + attachment?.$case === "query" + ? { + title: attachment.query.title, + description: attachment.query.description, + query: attachment.query.query, + statementId: attachment.query.statementId, + } + : undefined, + text: + attachment?.$case === "text" + ? { content: attachment.text.content } + : undefined, + suggestedQuestions: + attachment?.$case === "suggestedQuestions" + ? attachment.suggestedQuestions.questions + : undefined, })) ?? [] ); } function toMessageResponse(message: GenieMessage): GenieMessageResponse { return { - messageId: message.message_id, - conversationId: message.conversation_id, - spaceId: message.space_id, + messageId: message.messageId ?? "", + conversationId: message.conversationId ?? "", + spaceId: message.spaceId ?? "", status: message.status ?? "COMPLETED", - content: message.content, + content: message.content ?? "", attachments: mapAttachments(message), error: message.error?.error, }; } +/** + * The modular SDK returns the statement response camelCased with int64 fields as + * `bigint`. The SSE contract (`GenieStatementResponse`, read by appkit-ui) is the + * raw snake_case API shape, and `JSON.stringify` throws on `bigint`, so convert + * back here. Recurses into objects only; `data_array` rows pass through as-is. + */ +function toWireShape(value: unknown): unknown { + if (typeof value === "bigint") return Number(value); + if (Array.isArray(value)) return value.map(toWireShape); + if (value === null || typeof value !== "object") return value; + return Object.fromEntries( + Object.entries(value).map(([key, v]) => [ + key.replace(/[A-Z]/g, (c) => `_${c.toLowerCase()}`), + toWireShape(v), + ]), + ); +} + function classifyGenieError(error: unknown): string { const message = error instanceof Error ? error.message : String(error); + // Modular ApiError carries the code on `.code`, legacy on `.errorCode`. + const { code, errorCode } = (error ?? {}) as { + code?: unknown; + errorCode?: unknown; + }; - if (message.includes("RESOURCE_DOES_NOT_EXIST")) { + if ( + code === "RESOURCE_DOES_NOT_EXIST" || + errorCode === "RESOURCE_DOES_NOT_EXIST" || + message.includes("RESOURCE_DOES_NOT_EXIST") + ) { return GenieErrors.SPACE_ACCESS_DENIED; } @@ -100,26 +132,74 @@ export class GenieConnector { conversationId: string; messageId: string; }> { - if (conversationId) { - const waiter = await workspaceClient.genie.createMessage({ - space_id: spaceId, - conversation_id: conversationId, - content, - }); - return { - messageWaiter: waiter, - conversationId, - messageId: waiter.message_id ?? "", - }; - } - const start = await workspaceClient.genie.startConversation({ - space_id: spaceId, - content, - }); + const started = conversationId + ? await workspaceClient.genie.genieCreateConversationMessage({ + spaceId, + conversationId, + content, + }) + : await workspaceClient.genie.genieStartConversation({ + spaceId, + content, + }); return { - messageWaiter: start as unknown as CreateMessageWaiter, - conversationId: start.conversation_id, - messageId: start.message_id, + messageWaiter: this.messagePoller( + workspaceClient, + spaceId, + started.conversationId, + started.messageId, + ), + conversationId: started.conversationId, + messageId: started.messageId, + }; + } + + /** + * Polls `getConversationMessage` until COMPLETED. Replaces the SDK waiter: the + * modular `wait()` has no `onProgress`, which the SSE `status` events need. It + * mirrors the legacy waiter: progress on every poll, backoff of `attempt` + * seconds + 50-750ms jitter capped at 10s, a 10 minute default timeout, and the + * same `failed to reach COMPLETED state` errors `classifyGenieError` matches. + */ + private messagePoller( + workspaceClient: WorkspaceClient, + spaceId: string, + conversationId: string, + messageId: string, + ): CreateMessageWaiter { + return { + async wait(options) { + const timeout = + typeof options?.timeout === "number" + ? options.timeout + : DEFAULT_WAIT_TIMEOUT_MS; + const deadline = Date.now() + timeout; + let lastStatus: string | undefined; + for (let attempt = 1; Date.now() < deadline; attempt++) { + const message = + await workspaceClient.genie.genieGetConversationMessage({ + spaceId, + conversationId, + messageId, + }); + await options?.onProgress?.(message); + lastStatus = message.status; + if (lastStatus === "COMPLETED") return message; + if (lastStatus === "FAILED") { + throw new Error("failed to reach COMPLETED state, got FAILED"); + } + const jitter = 50 + Math.random() * 700; + await new Promise((resolve) => + setTimeout( + resolve, + Math.min(attempt * 1000 + jitter, MAX_POLL_INTERVAL_MS), + ), + ); + } + throw new Error( + `timed out: failed to reach COMPLETED state, got ${lastStatus}`, + ); + }, }; } @@ -128,9 +208,7 @@ export class GenieConnector { options?: { timeout?: number }, ): Promise { const timeout = options?.timeout ?? this.config.timeout; - const waitOptions = - timeout > 0 ? { timeout: new Time(timeout, TimeUnits.milliseconds) } : {}; - return messageWaiter.wait(waitOptions); + return messageWaiter.wait(timeout > 0 ? { timeout } : {}); } async listConversationMessages( @@ -145,18 +223,18 @@ export class GenieConnector { const pageSize = options?.pageSize ?? genieConnectorDefaults.initialPageSize; - const response = await workspaceClient.genie.listConversationMessages({ - space_id: spaceId, - conversation_id: conversationId, - page_size: pageSize, - ...(options?.pageToken ? { page_token: options.pageToken } : {}), + const response = await workspaceClient.genie.genieListConversationMessages({ + spaceId, + conversationId, + pageSize, + ...(options?.pageToken ? { pageToken: options.pageToken } : {}), }); const messages = (response.messages ?? []).reverse().map(toMessageResponse); return { messages, - nextPageToken: response.next_page_token ?? null, + nextPageToken: response.nextPageToken ?? null, }; } @@ -169,13 +247,13 @@ export class GenieConnector { _signal?: AbortSignal, ): Promise { const response = - await workspaceClient.genie.getMessageAttachmentQueryResult({ - space_id: spaceId, - conversation_id: conversationId, - message_id: messageId, - attachment_id: attachmentId, + await workspaceClient.genie.genieGetMessageAttachmentQueryResult({ + spaceId, + conversationId, + messageId, + attachmentId, }); - return response.statement_response as GenieStatementResponse; + return toWireShape(response.statementResponse) as GenieStatementResponse; } async *streamSendMessage( @@ -206,10 +284,7 @@ export class GenieConnector { const timeout = options?.timeout != null ? options.timeout : this.config.timeout; - const waitOptions = - timeout > 0 - ? { timeout: new Time(timeout, TimeUnits.milliseconds) } - : {}; + const waitOptions = timeout > 0 ? { timeout } : {}; let completedMessage!: GenieMessage; for await (const event of pollWaiter(messageWaiter, waitOptions)) { @@ -402,11 +477,10 @@ export class GenieConnector { while (true) { if (signal?.aborted) return; - const message = await workspaceClient.genie.getMessage({ - space_id: spaceId, - conversation_id: conversationId, - message_id: messageId, - }); + const message = await workspaceClient.genie.genieGetConversationMessage( + { spaceId, conversationId, messageId }, + { signal }, + ); if (message.status && message.status !== lastStatus) { lastStatus = message.status; diff --git a/packages/appkit/src/connectors/genie/tests/client.test.ts b/packages/appkit/src/connectors/genie/tests/client.test.ts index 50ae69f2b..287bef9d1 100644 --- a/packages/appkit/src/connectors/genie/tests/client.test.ts +++ b/packages/appkit/src/connectors/genie/tests/client.test.ts @@ -20,9 +20,9 @@ async function collect( function makeGenieMessage(overrides: Partial = {}): GenieMessage { return { - message_id: "msg-1", - conversation_id: "conv-1", - space_id: "space-1", + messageId: "msg-1", + conversationId: "conv-1", + spaceId: "space-1", status: "COMPLETED", content: "Hello from Genie", attachments: [], @@ -36,12 +36,15 @@ function makeGenieMessageWithQuery( return makeGenieMessage({ attachments: [ { - attachment_id: "att-1", - query: { - title: "Sales Query", - description: "Total sales", - query: "SELECT sum(amount) FROM sales", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Sales Query", + description: "Total sales", + query: "SELECT sum(amount) FROM sales", + statementId: "stmt-1", + }, }, }, ], @@ -53,37 +56,26 @@ function makeGenieMessageWithQuery( function createMockWorkspaceClient() { return { genie: { - startConversation: vi.fn(), - createMessage: vi.fn(), - getMessage: vi.fn(), - listConversationMessages: vi.fn(), - getMessageAttachmentQueryResult: vi.fn(), + genieStartConversation: vi.fn(), + genieCreateConversationMessage: vi.fn(), + genieGetConversationMessage: vi.fn(), + genieListConversationMessages: vi.fn(), + genieGetMessageAttachmentQueryResult: vi.fn(), }, } as any; } -/** - * Builds a mock waiter whose `.wait()` invokes `onProgress` for each - * progress value, then resolves with the final result. - */ -function createMockWaiter(opts: { - progressValues?: Partial[]; - result: GenieMessage; -}) { - return { - wait: vi.fn().mockImplementation(async (options: any = {}) => { - if (opts.progressValues) { - for (const value of opts.progressValues) { - if (options.onProgress) { - await options.onProgress(value); - } - } - } - return opts.result; - }), - message_id: opts.result.message_id, - conversation_id: opts.result.conversation_id, - }; +/** Stubs the start call; every poll then returns `result`. */ +function mockStart( + ws: ReturnType, + result: GenieMessage, + method = "genieStartConversation", +) { + ws.genie[method].mockResolvedValue({ + messageId: result.messageId, + conversationId: result.conversationId, + }); + ws.genie.genieGetConversationMessage.mockResolvedValue(result); } // --------------------------------------------------------------------------- @@ -105,17 +97,14 @@ describe("GenieConnector", () => { describe("streamSendMessage", () => { test("yields message_start, status updates, then message_result", async () => { + vi.useFakeTimers(); const completedMsg = makeGenieMessage(); - const waiter = createMockWaiter({ - progressValues: [ - { status: "EXECUTING_QUERY" }, - { status: "COMPLETED" }, - ], - result: completedMsg, - }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); + ws.genie.genieGetConversationMessage.mockResolvedValueOnce( + makeGenieMessage({ status: "EXECUTING_QUERY" }), + ); - const events = await collect( + const pending = collect( connector.streamSendMessage( ws, "space-1", @@ -123,6 +112,9 @@ describe("GenieConnector", () => { undefined, ), ); + await vi.runAllTimersAsync(); + const events = await pending; + vi.useRealTimers(); expect(events[0]).toEqual({ type: "message_start", @@ -144,51 +136,59 @@ describe("GenieConnector", () => { test("new conversation calls startConversation", async () => { const completedMsg = makeGenieMessage(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); await collect( connector.streamSendMessage(ws, "space-1", "hello", undefined), ); - expect(ws.genie.startConversation).toHaveBeenCalledWith({ - space_id: "space-1", + expect(ws.genie.genieStartConversation).toHaveBeenCalledWith({ + spaceId: "space-1", content: "hello", }); - expect(ws.genie.createMessage).not.toHaveBeenCalled(); + expect(ws.genie.genieCreateConversationMessage).not.toHaveBeenCalled(); }); test("existing conversation calls createMessage", async () => { const completedMsg = makeGenieMessage(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.createMessage.mockResolvedValue(waiter); + mockStart(ws, completedMsg, "genieCreateConversationMessage"); await collect( connector.streamSendMessage(ws, "space-1", "hello", "conv-existing"), ); - expect(ws.genie.createMessage).toHaveBeenCalledWith({ - space_id: "space-1", - conversation_id: "conv-existing", + expect(ws.genie.genieCreateConversationMessage).toHaveBeenCalledWith({ + spaceId: "space-1", + conversationId: "conv-existing", content: "hello", }); - expect(ws.genie.startConversation).not.toHaveBeenCalled(); + expect(ws.genie.genieStartConversation).not.toHaveBeenCalled(); }); test("emits query_result for attachments with statementIds", async () => { const completedMsg = makeGenieMessageWithQuery(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); + // The modular SDK hands back camelCase with int64 counts as bigint; the + // SSE payload must be the snake_case wire shape and JSON-serializable. + ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: { + statementId: "stmt-1", + manifest: { + schema: { columns: [{ name: "total", typeName: "DOUBLE" }] }, + totalRowCount: 1n, + }, + result: { dataArray: [["1234.56", null]], rowCount: 1n }, + }, + }); const statementResponse = { + statement_id: "stmt-1", manifest: { schema: { columns: [{ name: "total", type_name: "DOUBLE" }] }, + total_row_count: 1, }, - result: { data_array: [["1234.56"]] }, + result: { data_array: [["1234.56", null]], row_count: 1 }, }; - ws.genie.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: statementResponse, - }); const events = await collect( connector.streamSendMessage(ws, "space-1", "query", undefined), @@ -201,10 +201,25 @@ describe("GenieConnector", () => { statementId: "stmt-1", data: statementResponse, }); + expect(() => JSON.stringify(queryResult)).not.toThrow(); + }); + + test("maps a FAILED poll to the table permissions error", async () => { + mockStart(ws, makeGenieMessage({ status: "FAILED" })); + + const events = await collect( + connector.streamSendMessage(ws, "space-1", "q", undefined), + ); + + expect(events.at(-1)).toEqual({ + type: "error", + error: + "You may not have access to the data tables. Please verify your table permissions.", + }); }); test("yields error event on SDK failure", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("Network timeout"), ); @@ -216,7 +231,7 @@ describe("GenieConnector", () => { }); test("classifies RESOURCE_DOES_NOT_EXIST as access denied", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("RESOURCE_DOES_NOT_EXIST: space not found"), ); @@ -234,9 +249,8 @@ describe("GenieConnector", () => { test("emits error event when query result fetch fails", async () => { const completedMsg = makeGenieMessageWithQuery(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); - ws.genie.getMessageAttachmentQueryResult.mockRejectedValue( + mockStart(ws, completedMsg); + ws.genie.genieGetMessageAttachmentQueryResult.mockRejectedValue( new Error("statement expired"), ); @@ -258,12 +272,12 @@ describe("GenieConnector", () => { describe("streamConversation", () => { test("yields message_result for each message, then history_info", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [ - makeGenieMessage({ message_id: "m1", content: "first" }), - makeGenieMessage({ message_id: "m2", content: "second" }), + makeGenieMessage({ messageId: "m1", content: "first" }), + makeGenieMessage({ messageId: "m2", content: "second" }), ], - next_page_token: null, + nextPageToken: null, }); const events = await collect( @@ -286,39 +300,45 @@ describe("GenieConnector", () => { }); test("fetches query results in parallel when includeQueryResults=true", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [ makeGenieMessageWithQuery({ - message_id: "m1", + messageId: "m1", attachments: [ { - attachment_id: "att-a", - query: { - title: "Q1", - query: "SELECT 1", - statement_id: "stmt-a", + attachmentId: "att-a", + attachment: { + $case: "query", + query: { + title: "Q1", + query: "SELECT 1", + statementId: "stmt-a", + }, }, }, { - attachment_id: "att-b", - query: { - title: "Q2", - query: "SELECT 2", - statement_id: "stmt-b", + attachmentId: "att-b", + attachment: { + $case: "query", + query: { + title: "Q2", + query: "SELECT 2", + statementId: "stmt-b", + }, }, }, ], }), ], - next_page_token: null, + nextPageToken: null, }); const stmtResponse = { manifest: { schema: { columns: [] } }, - result: { data_array: [] }, + result: { dataArray: [] }, }; - ws.genie.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: stmtResponse, + ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: stmtResponse, }); const events = await collect( @@ -329,13 +349,15 @@ describe("GenieConnector", () => { const queryResults = events.filter((e) => e.type === "query_result"); expect(queryResults).toHaveLength(2); - expect(ws.genie.getMessageAttachmentQueryResult).toHaveBeenCalledTimes(2); + expect( + ws.genie.genieGetMessageAttachmentQueryResult, + ).toHaveBeenCalledTimes(2); }); test("skips query results when includeQueryResults=false", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [makeGenieMessageWithQuery()], - next_page_token: null, + nextPageToken: null, }); const events = await collect( @@ -345,44 +367,52 @@ describe("GenieConnector", () => { ); expect(events.filter((e) => e.type === "query_result")).toHaveLength(0); - expect(ws.genie.getMessageAttachmentQueryResult).not.toHaveBeenCalled(); + expect( + ws.genie.genieGetMessageAttachmentQueryResult, + ).not.toHaveBeenCalled(); }); test("handles partial query result failures via Promise.allSettled", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [ makeGenieMessage({ - message_id: "m1", + messageId: "m1", attachments: [ { - attachment_id: "att-ok", - query: { - title: "OK", - query: "SELECT 1", - statement_id: "stmt-ok", + attachmentId: "att-ok", + attachment: { + $case: "query", + query: { + title: "OK", + query: "SELECT 1", + statementId: "stmt-ok", + }, }, }, { - attachment_id: "att-fail", - query: { - title: "Fail", - query: "SELECT 2", - statement_id: "stmt-fail", + attachmentId: "att-fail", + attachment: { + $case: "query", + query: { + title: "Fail", + query: "SELECT 2", + statementId: "stmt-fail", + }, }, }, ], }), ], - next_page_token: null, + nextPageToken: null, }); const stmtResponse = { manifest: { schema: { columns: [] } }, - result: { data_array: [] }, + result: { dataArray: [] }, }; - ws.genie.getMessageAttachmentQueryResult - .mockResolvedValueOnce({ statement_response: stmtResponse }) + ws.genie.genieGetMessageAttachmentQueryResult + .mockResolvedValueOnce({ statementResponse: stmtResponse }) .mockRejectedValueOnce(new Error("statement expired")); const events = await collect( @@ -400,7 +430,7 @@ describe("GenieConnector", () => { }); test("yields error when listConversationMessages fails", async () => { - ws.genie.listConversationMessages.mockRejectedValue( + ws.genie.genieListConversationMessages.mockRejectedValue( new Error("RESOURCE_DOES_NOT_EXIST: conv not found"), ); @@ -423,7 +453,7 @@ describe("GenieConnector", () => { describe("streamGetMessage", () => { test("polls until COMPLETED, yields status + message_result", async () => { - ws.genie.getMessage + ws.genie.genieGetConversationMessage .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce(makeGenieMessage({ status: "COMPLETED" })); @@ -439,11 +469,11 @@ describe("GenieConnector", () => { }); expect(events[1]).toEqual({ type: "status", status: "COMPLETED" }); expect(events[2]).toMatchObject({ type: "message_result" }); - expect(ws.genie.getMessage).toHaveBeenCalledTimes(2); + expect(ws.genie.genieGetConversationMessage).toHaveBeenCalledTimes(2); }); test("polls until FAILED, yields status + message_result", async () => { - ws.genie.getMessage + ws.genie.genieGetConversationMessage .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce( makeGenieMessage({ @@ -472,7 +502,7 @@ describe("GenieConnector", () => { test("respects abort signal", async () => { const controller = new AbortController(); - ws.genie.getMessage.mockResolvedValue( + ws.genie.genieGetConversationMessage.mockResolvedValue( makeGenieMessage({ status: "EXECUTING_QUERY" }), ); @@ -500,7 +530,9 @@ describe("GenieConnector", () => { }); test("yields error when getMessage throws", async () => { - ws.genie.getMessage.mockRejectedValue(new Error("service unavailable")); + ws.genie.genieGetConversationMessage.mockRejectedValue( + new Error("service unavailable"), + ); const events = await collect( connector.streamGetMessage(ws, "space-1", "conv-1", "msg-1", { @@ -512,7 +544,7 @@ describe("GenieConnector", () => { }); test("does not duplicate status events for same status", async () => { - ws.genie.getMessage + ws.genie.genieGetConversationMessage .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce(makeGenieMessage({ status: "COMPLETED" })); @@ -538,11 +570,10 @@ describe("GenieConnector", () => { describe("sendMessage", () => { test("returns completed message response", async () => { const completedMsg = makeGenieMessage({ - message_id: "msg-42", - conversation_id: "conv-new", + messageId: "msg-42", + conversationId: "conv-new", }); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); const result = await connector.sendMessage( ws, @@ -566,17 +597,17 @@ describe("GenieConnector", () => { // listConversationMessages reverses the SDK response, so mock data // is ordered newest-first (as the SDK returns) and results are // oldest-first after reversal. - ws.genie.listConversationMessages + ws.genie.genieListConversationMessages .mockResolvedValueOnce({ messages: [ - makeGenieMessage({ message_id: "m2" }), - makeGenieMessage({ message_id: "m1" }), + makeGenieMessage({ messageId: "m2" }), + makeGenieMessage({ messageId: "m1" }), ], - next_page_token: "page2", + nextPageToken: "page2", }) .mockResolvedValueOnce({ - messages: [makeGenieMessage({ message_id: "m3" })], - next_page_token: null, + messages: [makeGenieMessage({ messageId: "m3" })], + nextPageToken: null, }); const result = await connector.getConversation(ws, "space-1", "conv-1"); @@ -587,7 +618,7 @@ describe("GenieConnector", () => { "m2", "m3", ]); - expect(ws.genie.listConversationMessages).toHaveBeenCalledTimes(2); + expect(ws.genie.genieListConversationMessages).toHaveBeenCalledTimes(2); }); test("respects maxMessages limit", async () => { @@ -596,13 +627,13 @@ describe("GenieConnector", () => { maxMessages: 2, }); - ws.genie.listConversationMessages.mockResolvedValueOnce({ + ws.genie.genieListConversationMessages.mockResolvedValueOnce({ messages: [ - makeGenieMessage({ message_id: "m1" }), - makeGenieMessage({ message_id: "m2" }), - makeGenieMessage({ message_id: "m3" }), + makeGenieMessage({ messageId: "m1" }), + makeGenieMessage({ messageId: "m2" }), + makeGenieMessage({ messageId: "m3" }), ], - next_page_token: "page2", + nextPageToken: "page2", }); const result = await smallConnector.getConversation( @@ -614,7 +645,7 @@ describe("GenieConnector", () => { // Should be sliced to maxMessages expect(result.messages).toHaveLength(2); // Should NOT fetch a second page since length already >= maxMessages - expect(ws.genie.listConversationMessages).toHaveBeenCalledTimes(1); + expect(ws.genie.genieListConversationMessages).toHaveBeenCalledTimes(1); }); }); @@ -625,14 +656,13 @@ describe("GenieConnector", () => { describe("mapAttachments", () => { test("handles query attachments", async () => { const msg = makeGenieMessageWithQuery(); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); // We drive through streamSendMessage to exercise mapAttachments - ws.genie.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: { + ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: { manifest: { schema: { columns: [] } }, - result: { data_array: [] }, + result: { dataArray: [] }, }, }); @@ -658,13 +688,15 @@ describe("GenieConnector", () => { const msg = makeGenieMessage({ attachments: [ { - attachment_id: "att-text", - text: { content: "Here is the explanation" }, + attachmentId: "att-text", + attachment: { + $case: "text", + text: { content: "Here is the explanation" }, + }, }, ], }); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); const events = await collect( connector.streamSendMessage(ws, "space-1", "q", undefined), @@ -683,15 +715,17 @@ describe("GenieConnector", () => { const msg = makeGenieMessage({ attachments: [ { - attachment_id: "att-sq", - suggested_questions: { - questions: ["What is X?", "Show me Y"], + attachmentId: "att-sq", + attachment: { + $case: "suggestedQuestions", + suggestedQuestions: { + questions: ["What is X?", "Show me Y"], + }, }, }, ], }); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); const events = await collect( connector.streamSendMessage(ws, "space-1", "q", undefined), @@ -708,8 +742,7 @@ describe("GenieConnector", () => { test("returns empty array when message has no attachments", async () => { const msg = makeGenieMessage({ attachments: undefined }); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); const events = await collect( connector.streamSendMessage(ws, "space-1", "q", undefined), @@ -726,7 +759,7 @@ describe("GenieConnector", () => { describe("classifyGenieError", () => { test("maps RESOURCE_DOES_NOT_EXIST to space access denied", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("RESOURCE_DOES_NOT_EXIST: space xyz"), ); @@ -740,8 +773,25 @@ describe("GenieConnector", () => { }); }); + test("maps an ApiError .code of RESOURCE_DOES_NOT_EXIST to space access denied", async () => { + ws.genie.genieStartConversation.mockRejectedValue( + Object.assign(new Error("Space xyz not found"), { + code: "RESOURCE_DOES_NOT_EXIST", + }), + ); + + const events = await collect( + connector.streamSendMessage(ws, "space-1", "hi", undefined), + ); + + expect(events[0]).toEqual({ + type: "error", + error: "You don't have access to this Genie Space.", + }); + }); + test("maps failed-to-reach-COMPLETED + FAILED to table permissions", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("failed to reach COMPLETED state, got FAILED"), ); @@ -757,7 +807,7 @@ describe("GenieConnector", () => { }); test("passes through unknown error messages", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("something unexpected"), ); @@ -772,7 +822,7 @@ describe("GenieConnector", () => { }); test("handles non-Error throwable", async () => { - ws.genie.startConversation.mockRejectedValue("string error"); + ws.genie.genieStartConversation.mockRejectedValue("string error"); const events = await collect( connector.streamSendMessage(ws, "space-1", "hi", undefined), diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index 5cd03a348..09b26d4c1 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.test.ts @@ -19,68 +19,73 @@ import type { IGenieConfig } from "../types"; useTestCache(); function createMockGenieService() { - const getMessageAttachmentQueryResult = vi.fn(); - - const createWaiter = ( + const genieGetMessageAttachmentQueryResult = vi.fn(); + // Messages the start/create stubs "created"; the poll serves them back. + const messages = new Map(); + const register = ( conversationId: string, messageId: string, attachments: any[] = [], status = "COMPLETED", - ) => ({ - wait: vi.fn().mockImplementation(async ({ onProgress }: any) => { - if (onProgress) { - await onProgress({ status: "ASKING_AI" }); - await onProgress({ status: "EXECUTING_QUERY" }); - } - return { - message_id: messageId, - conversation_id: conversationId, - space_id: "test-space-id", - content: "Here are your results", - status, - attachments, - error: undefined, - }; - }), - }); + ) => { + messages.set(messageId, { + messageId, + conversationId, + spaceId: "test-space-id", + content: "Here are your results", + status, + attachments, + error: undefined, + }); + return { conversationId, messageId }; + }; + + const genieGetConversationMessage = vi + .fn() + .mockImplementation(async ({ messageId }: any) => messages.get(messageId)); - const startConversation = vi.fn().mockImplementation(async () => ({ - conversation_id: "new-conv-id", - message_id: "new-msg-id", - ...createWaiter("new-conv-id", "new-msg-id", [ + const genieStartConversation = vi.fn().mockImplementation(async () => + register("new-conv-id", "new-msg-id", [ { - attachment_id: "att-1", - query: { - title: "Top Customers", - description: "Query for top customers", - query: "SELECT * FROM customers", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Top Customers", + description: "Query for top customers", + query: "SELECT * FROM customers", + statementId: "stmt-1", + }, }, }, ]), - })); + ); - const createMessage = vi.fn().mockImplementation(async () => - createWaiter("existing-conv-id", "followup-msg-id", [ + const genieCreateConversationMessage = vi.fn().mockImplementation(async () => + register("existing-conv-id", "followup-msg-id", [ { - attachment_id: "att-2", - query: { - title: "Follow-up Query", - query: "SELECT * FROM orders", - statement_id: "stmt-2", + attachmentId: "att-2", + attachment: { + $case: "query", + query: { + title: "Follow-up Query", + query: "SELECT * FROM orders", + statementId: "stmt-2", + }, }, }, ]), ); - const listConversationMessages = vi.fn(); + const genieListConversationMessages = vi.fn(); return { - startConversation, - createMessage, - getMessageAttachmentQueryResult, - listConversationMessages, - createWaiter, + genieStartConversation, + genieCreateConversationMessage, + genieGetConversationMessage, + genieGetMessageAttachmentQueryResult, + genieListConversationMessages, + register, }; } @@ -102,11 +107,11 @@ describe("Genie Plugin", () => { mockGenieService = createMockGenieService(); - mockGenieService.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: { + mockGenieService.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: { status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ ["Acme Corp", "1000000"], ["Globex", "500000"], ], @@ -114,8 +119,8 @@ describe("Genie Plugin", () => { manifest: { schema: { columns: [ - { name: "customer", type_name: "STRING" }, - { name: "revenue", type_name: "DECIMAL" }, + { name: "customer", typeName: "STRING" }, + { name: "revenue", typeName: "DECIMAL" }, ], }, }, @@ -238,9 +243,9 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect(mockRes.status).not.toHaveBeenCalledWith(404); - expect(mockGenieService.startConversation).toHaveBeenCalledWith( + expect(mockGenieService.genieStartConversation).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", + spaceId: "test-space-id", content: "What are my top customers?", }), ); @@ -291,12 +296,16 @@ describe("Genie Plugin", () => { }, }); const mockRes = createMockResponse(); + // One non-terminal poll before COMPLETED (costs one ~1s backoff sleep). + mockGenieService.genieGetConversationMessage.mockResolvedValueOnce({ + status: "ASKING_AI", + }); await handler(mockReq, mockRes); - expect(mockGenieService.startConversation).toHaveBeenCalledWith( + expect(mockGenieService.genieStartConversation).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", + spaceId: "test-space-id", content: "What are my top customers?", }), ); @@ -357,15 +366,17 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); - expect(mockGenieService.createMessage).toHaveBeenCalledWith( + expect( + mockGenieService.genieCreateConversationMessage, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "existing-conv-id", + spaceId: "test-space-id", + conversationId: "existing-conv-id", content: "Show me more details", }), ); - expect(mockGenieService.startConversation).not.toHaveBeenCalled(); + expect(mockGenieService.genieStartConversation).not.toHaveBeenCalled(); const writeCalls = mockRes.write.mock.calls.map((call: any[]) => call[0]); const allWritten = writeCalls.join(""); @@ -379,51 +390,46 @@ describe("Genie Plugin", () => { describe("multiple attachments", () => { test("should yield query_result for each query attachment", async () => { // Override startConversation to return multiple query attachments - mockGenieService.startConversation.mockImplementation(async () => ({ - conversation_id: "multi-conv-id", - message_id: "multi-msg-id", - wait: vi.fn().mockImplementation(async ({ onProgress }: any) => { - if (onProgress) { - await onProgress({ status: "ASKING_AI" }); - } - return { - message_id: "multi-msg-id", - conversation_id: "multi-conv-id", - space_id: "test-space-id", - content: "Here are two queries", - status: "COMPLETED", - attachments: [ - { - attachment_id: "att-q1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-q1", - }, - }, - { - attachment_id: "att-q2", - query: { - title: "Query 2", - query: "SELECT 2", - statement_id: "stmt-q2", - }, + mockGenieService.genieStartConversation.mockImplementation(async () => + mockGenieService.register("multi-conv-id", "multi-msg-id", [ + { + attachmentId: "att-q1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-q1", }, - { - attachment_id: "att-text", - text: { content: "Some explanation" }, + }, + }, + { + attachmentId: "att-q2", + attachment: { + $case: "query", + query: { + title: "Query 2", + query: "SELECT 2", + statementId: "stmt-q2", }, - ], - }; - }), - })); + }, + }, + { + attachmentId: "att-text", + attachment: { + $case: "text", + text: { content: "Some explanation" }, + }, + }, + ]), + ); - mockGenieService.getMessageAttachmentQueryResult + mockGenieService.genieGetMessageAttachmentQueryResult .mockResolvedValueOnce({ - statement_response: { result: { data: [["row1"]] } }, + statementResponse: { result: { data: [["row1"]] } }, }) .mockResolvedValueOnce({ - statement_response: { result: { data: [["row2"]] } }, + statementResponse: { result: { data: [["row2"]] } }, }); const plugin = new GeniePlugin(config); @@ -446,18 +452,18 @@ describe("Genie Plugin", () => { // getMessageAttachmentQueryResult should be called twice (once per query attachment) expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledTimes(2); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( - expect.objectContaining({ attachment_id: "att-q1" }), + expect.objectContaining({ attachmentId: "att-q1" }), ); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( - expect.objectContaining({ attachment_id: "att-q2" }), + expect.objectContaining({ attachmentId: "att-q2" }), ); const writeCalls = mockRes.write.mock.calls.map((call: any[]) => call[0]); @@ -473,7 +479,7 @@ describe("Genie Plugin", () => { describe("error handling", () => { test("should yield error event on SDK failure", async () => { - mockGenieService.startConversation.mockRejectedValue( + mockGenieService.genieStartConversation.mockRejectedValue( new Error("Genie service unavailable"), ); @@ -519,9 +525,9 @@ describe("Genie Plugin", () => { } function mockMessages(messages: any[]) { - mockGenieService.listConversationMessages.mockResolvedValue({ + mockGenieService.genieListConversationMessages.mockResolvedValue({ messages, - next_page_token: undefined, + nextPageToken: undefined, }); } @@ -551,26 +557,29 @@ describe("Genie Plugin", () => { test("should stream message_result events for each message", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "What are the top customers?", status: "COMPLETED", attachments: [], }, { - message_id: "msg-2", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-2", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Here are the results", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Top Customers", - query: "SELECT * FROM customers", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Top Customers", + query: "SELECT * FROM customers", + statementId: "stmt-1", + }, }, }, ], @@ -591,11 +600,13 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); - expect(mockGenieService.listConversationMessages).toHaveBeenCalledWith( + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - page_size: genieConnectorDefaults.initialPageSize, + spaceId: "test-space-id", + conversationId: "conv-123", + pageSize: genieConnectorDefaults.initialPageSize, }), ); @@ -621,18 +632,21 @@ describe("Genie Plugin", () => { test("should stream query_result events when includeQueryResults is true (default)", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Results", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-1", + }, }, }, ], @@ -654,13 +668,13 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - message_id: "msg-1", - attachment_id: "att-1", + spaceId: "test-space-id", + conversationId: "conv-123", + messageId: "msg-1", + attachmentId: "att-1", }), ); @@ -675,18 +689,21 @@ describe("Genie Plugin", () => { test("should NOT stream query_result events when includeQueryResults is false", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Results", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-1", + }, }, }, ], @@ -710,7 +727,7 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).not.toHaveBeenCalled(); const writeCalls = mockRes.write.mock.calls.map((call: any[]) => call[0]); @@ -722,18 +739,18 @@ describe("Genie Plugin", () => { }); test("should fetch only one page and emit history_info with nextPageToken", async () => { - mockGenieService.listConversationMessages.mockResolvedValueOnce({ + mockGenieService.genieListConversationMessages.mockResolvedValueOnce({ messages: [ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Most recent message", status: "COMPLETED", attachments: [], }, ], - next_page_token: "page-2-token", + nextPageToken: "page-2-token", }); const plugin = new GeniePlugin(config); @@ -753,15 +770,17 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); // Should only fetch one page (lazy loading) - expect(mockGenieService.listConversationMessages).toHaveBeenCalledTimes( - 1, - ); + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledTimes(1); - expect(mockGenieService.listConversationMessages).toHaveBeenCalledWith( + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - page_size: genieConnectorDefaults.initialPageSize, + spaceId: "test-space-id", + conversationId: "conv-123", + pageSize: genieConnectorDefaults.initialPageSize, }), ); @@ -778,9 +797,9 @@ describe("Genie Plugin", () => { test("should emit history_info with null nextPageToken when no more pages", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Only message", status: "COMPLETED", attachments: [], @@ -838,7 +857,7 @@ describe("Genie Plugin", () => { }); test("should yield error event on SDK failure", async () => { - mockGenieService.listConversationMessages.mockRejectedValue( + mockGenieService.genieListConversationMessages.mockRejectedValue( new Error("Conversation not found"), ); @@ -867,35 +886,41 @@ describe("Genie Plugin", () => { test("should fetch query results in parallel for multiple attachments across messages", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "First query", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-1", + }, }, }, ], }, { - message_id: "msg-2", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-2", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Second query", status: "COMPLETED", attachments: [ { - attachment_id: "att-2", - query: { - title: "Query 2", - query: "SELECT 2", - statement_id: "stmt-2", + attachmentId: "att-2", + attachment: { + $case: "query", + query: { + title: "Query 2", + query: "SELECT 2", + statementId: "stmt-2", + }, }, }, ], @@ -917,23 +942,23 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledTimes(2); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( expect.objectContaining({ - message_id: "msg-1", - attachment_id: "att-1", + messageId: "msg-1", + attachmentId: "att-1", }), ); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( expect.objectContaining({ - message_id: "msg-2", - attachment_id: "att-2", + messageId: "msg-2", + attachmentId: "att-2", }), ); @@ -950,7 +975,7 @@ describe("Genie Plugin", () => { describe("error classification", () => { test("should return user-friendly message for RESOURCE_DOES_NOT_EXIST error", async () => { - mockGenieService.startConversation.mockRejectedValue( + mockGenieService.genieStartConversation.mockRejectedValue( new Error( "RESOURCE_DOES_NOT_EXIST: No Genie space found with id test-space-id", ), @@ -984,7 +1009,7 @@ describe("Genie Plugin", () => { }); test("should return user-friendly message for FAILED state error (table access denied)", async () => { - mockGenieService.startConversation.mockRejectedValue( + mockGenieService.genieStartConversation.mockRejectedValue( new Error( "failed to reach COMPLETED state, got FAILED: [object Object]", ), @@ -1018,7 +1043,7 @@ describe("Genie Plugin", () => { }); test("should return user-friendly message for RESOURCE_DOES_NOT_EXIST on getConversation", async () => { - mockGenieService.listConversationMessages.mockRejectedValue( + mockGenieService.genieListConversationMessages.mockRejectedValue( new Error( "RESOURCE_DOES_NOT_EXIST: No Genie space found with id test-space-id", ), @@ -1078,9 +1103,9 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect(mockRes.status).not.toHaveBeenCalledWith(404); - expect(mockGenieService.startConversation).toHaveBeenCalledWith( + expect(mockGenieService.genieStartConversation).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "env-space-id", + spaceId: "env-space-id", content: "test question", }), ); @@ -1162,18 +1187,18 @@ describe("Genie Plugin", () => { describe("getConversation with pageToken", () => { test("should pass pageToken through to streamConversation", async () => { - mockGenieService.listConversationMessages.mockResolvedValueOnce({ + mockGenieService.genieListConversationMessages.mockResolvedValueOnce({ messages: [ { - message_id: "msg-old-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-old-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Older message", status: "COMPLETED", attachments: [], }, ], - next_page_token: "next-token-abc", + nextPageToken: "next-token-abc", }); const plugin = new GeniePlugin(config); @@ -1197,11 +1222,13 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); - expect(mockGenieService.listConversationMessages).toHaveBeenCalledWith( + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - page_token: "some-page-token", + spaceId: "test-space-id", + conversationId: "conv-123", + pageToken: "some-page-token", }), ); @@ -1215,7 +1242,7 @@ describe("Genie Plugin", () => { }); test("should yield error event when paginated request fails", async () => { - mockGenieService.listConversationMessages.mockRejectedValue( + mockGenieService.genieListConversationMessages.mockRejectedValue( new Error("Page token expired"), ); diff --git a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts index 5ec90a0d9..1dada4f5c 100644 --- a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts +++ b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts @@ -21,7 +21,7 @@ describe("createMockWorkspaceClient", () => { // Never-crash is the headline claim, so all nine are asserted, not sampled. test.each([ ["files", "listDirectory", undefined], - ["genie", "getMessage", undefined], + ["genie", "genieGetConversationMessage", undefined], ["jobs", "getRun", undefined], ["servingEndpoints", "get", undefined], ["warehouses", "getWarehouse", { state: "RUNNING" }], diff --git a/packages/appkit/src/workspace-client/index.ts b/packages/appkit/src/workspace-client/index.ts index a7d7e1c34..39e400c6a 100644 --- a/packages/appkit/src/workspace-client/index.ts +++ b/packages/appkit/src/workspace-client/index.ts @@ -5,14 +5,7 @@ * behind a stable facade. */ -export { - ApiError, - ConfigError, - Context, - createWorkspaceClient, - Time, - TimeUnits, -} from "shared"; +export { ApiError, ConfigError, Context, createWorkspaceClient } from "shared"; // Forwards every wrapper type — legacy service namespaces (files/jobs/serving), // the client option/waiter types, and the modular SDK client + model types // (warehouses, statementExecution). `sql` is gone: its statement + warehouse diff --git a/packages/shared/package.json b/packages/shared/package.json index 867a9f2b4..6002e4693 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -51,6 +51,7 @@ "@databricks/sdk-auth": "0.51.0", "@databricks/sdk-core": "0.51.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-options": "0.51.0", "@databricks/sdk-scim": "0.51.0", "@databricks/sdk-statementexecution": "0.52.0", diff --git a/packages/shared/src/workspace-client/client.ts b/packages/shared/src/workspace-client/client.ts index 29494bb3c..1552392e5 100644 --- a/packages/shared/src/workspace-client/client.ts +++ b/packages/shared/src/workspace-client/client.ts @@ -18,6 +18,8 @@ import { buildWarehousesClient, buildWorkspaceAuth, type ScimClient, + buildGenieClient, + type GenieClient, type StatementExecutionClient, type WarehousesClient, type WorkspaceAuth, @@ -32,6 +34,7 @@ export class AppKitWorkspaceClient implements WorkspaceClient { #statementExecution?: StatementExecutionClient; #auth?: WorkspaceAuth; #currentUser?: ScimClient; + #genie?: GenieClient; constructor(opts: WorkspaceClientOptions) { this.#opts = opts; @@ -49,8 +52,12 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#warehouses; } - get genie() { - return this.#getLegacy().genie; + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get genie(): GenieClient { + if (!this.#genie) { + this.#genie = buildGenieClient(this.#opts); + } + return this.#genie; } get jobs() { diff --git a/packages/shared/src/workspace-client/index.ts b/packages/shared/src/workspace-client/index.ts index 585bc46c5..f037e8ef4 100644 --- a/packages/shared/src/workspace-client/index.ts +++ b/packages/shared/src/workspace-client/index.ts @@ -10,9 +10,7 @@ export { createWorkspaceClient } from "./factory"; export type { CancellationToken, ClientOptions, - GenieMessage, LegacyWorkspaceClient, - Waiter, WorkspaceClientOptions, } from "./legacy"; // SDK value + type re-exports so AppKit modules import them from the wrapper. diff --git a/packages/shared/src/workspace-client/legacy.ts b/packages/shared/src/workspace-client/legacy.ts index 8784fbe57..745cf0976 100644 --- a/packages/shared/src/workspace-client/legacy.ts +++ b/packages/shared/src/workspace-client/legacy.ts @@ -77,7 +77,7 @@ export type { // // AppKit modules import these from the wrapper instead of the SDK so the // boundary rule holds. `Context` bridges AbortSignal → CancellationToken -// (serving/jobs/sql-warehouse); `Time`/`TimeUnits` drive genie polling; +// (serving/jobs/sql-warehouse); `Time`/`TimeUnits` are re-exported by `shared`'s index; // `ConfigError` is matched in service-context's auth-failure handling. // // These are sourced off the namespace import rather than `export { ... } from` @@ -90,9 +90,3 @@ export type { export const { ConfigError, Context, TimeUnits, loadConfigFile } = SDK; export const Time = SDK.Time ?? (SDK as unknown as { default: typeof SDK }).default.Time; - -// Deep-import types used by the genie connector's waiter idiom. Not exposed -// on the SDK's top-level index, so re-exported here to keep the genie -// connector off a direct `@databricks/sdk-experimental/dist/**` import. -export type { GenieMessage } from "@databricks/sdk-experimental/dist/apis/dashboards"; -export type { Waiter } from "@databricks/sdk-experimental/dist/wait"; diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 058701ddd..4e02de2f8 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -16,6 +16,12 @@ * (`patches/@databricks__sdk-statementexecution@0.46.0.patch`) that restores the * undocumented Reyden `attachment` response field, which the SDK's generated * unmarshal transform would otherwise strip. + * + * NOTE: genie relies on a pinned pnpm patch + * (`patches/@databricks__sdk-genie@0.54.0.patch`): the generated model types the + * query result's `data_array` as protobuf `ListValue[]` (`{ values: [...] }`), but + * the API returns plain `JSON_ARRAY` rows (`[["a", null], ...]`), so the unmarshal + * schema rejected every real query result. The patch restores `(string | null)[][]`. */ import { execFile } from "node:child_process"; import { promisify } from "node:util"; @@ -39,6 +45,7 @@ import { newFetchHttpClient, } from "@databricks/sdk-core/http"; import { type ProfileOptions, resolve } from "@databricks/sdk-core/profiles"; +import type { GenieClient } from "@databricks/sdk-genie/v1"; import type { ClientOptions } from "@databricks/sdk-options/client"; import { ScimClient } from "@databricks/sdk-scim/v1"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; @@ -499,7 +506,46 @@ export function buildScimClient(opts: WorkspaceClientOptions): ScimClient { return new ScimClient(mapToClientOptions(opts)); } +/** + * Defer loading an SDK client's module until its first method call, so app + * boot doesn't pay for clients it never uses. Generated methods return either + * a Promise or (`*Iter`) an AsyncGenerator, so each proxied call returns the + * pending result as a Promise that is also async-iterable. + */ +function lazyClient( + load: () => Promise T>, + options: ClientOptions, +): T { + let client: Promise | undefined; + return new Proxy({} as T, { + get(_target, prop) { + // Not thenable, or `await facade.someClient` would hang. + if (typeof prop === "symbol" || prop === "then") return undefined; + return (...args: unknown[]) => { + client ??= load().then((Client) => new Client(options)); + const result = client.then((c) => + (c as Record unknown>)[prop](...args), + ); + return Object.assign(result, { + async *[Symbol.asyncIterator]() { + yield* (await result) as AsyncIterable; + }, + }); + }; + }, + }); +} + +/** Build a modular Genie client from wrapper options. */ +export function buildGenieClient(opts: WorkspaceClientOptions): GenieClient { + return lazyClient( + () => import("@databricks/sdk-genie/v1").then((m) => m.GenieClient), + mapToClientOptions(opts), + ); +} + // ── Client type re-exports (for the facade accessor types) ─────────────── +export type { GenieClient } from "@databricks/sdk-genie/v1"; export type { ScimClient } from "@databricks/sdk-scim/v1"; export type { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; export type { WarehousesClient } from "@databricks/sdk-warehouses/v1"; @@ -524,6 +570,10 @@ export type { StatementStatus, StatementStatus_State, } from "@databricks/sdk-statementexecution/v1"; +export type { + GenieGetMessageQueryResultResponse, + GenieMessage, +} from "@databricks/sdk-genie/v1"; export type { EndpointHealth, EndpointInfo, diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 37c7711d2..c633eeda4 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -42,6 +42,19 @@ vi.mock("@databricks/sdk-warehouses/v1", () => ({ vi.mock("@databricks/sdk-statementexecution/v1", () => ({ StatementExecutionClient: vi.fn().mockImplementation((opts) => ({ opts })), })); +vi.mock("@databricks/sdk-genie/v1", () => ({ + GenieClient: vi.fn().mockImplementation((opts) => { + ctorOpts.push(opts); + return { + opts, + genieGetConversationMessage: async () => ({ id: "m" }), + async *listItemsIter() { + yield 1; + yield 2; + }, + }; + }), +})); vi.mock("@databricks/sdk-auth/credentials", () => ({ newPatCredentials: vi.fn((token: string) => { patTokens.push(token); @@ -104,7 +117,11 @@ vi.mock("node:child_process", () => ({ })); import { ApiError } from "../errors"; -import { buildWarehousesClient, buildWorkspaceAuth } from "../modular"; +import { + buildGenieClient, + buildWarehousesClient, + buildWorkspaceAuth, +} from "../modular"; /** Drive the wrapped httpClient with one request and return the UA it set. */ async function sentUserAgent( @@ -290,6 +307,44 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(ua).toBe("@databricks/appkit/0.64.0"); }); + test("genie (asUser) uses the OBO token as PAT and keeps the AppKit User-Agent", async () => { + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + await buildGenieClient({ + token: "user-token", + host: "https://x", + clientOptions: { + product: "@databricks/appkit", + productVersion: "0.64.0", + }, + } as never).genieGetConversationMessage({} as never); + expect(ctorOpts[0].credentials).toMatchObject({ + kind: "pat", + token: "user-token", + }); + expect(m2mOpts).toEqual([]); + expect(await sentUserAgent(ctorOpts[0].httpClient)).toBe( + "@databricks/appkit/0.64.0", + ); + }); + + test("lazy clients load the SDK module on first call and build once", async () => { + const client = buildGenieClient({ host: "https://x" }); + expect(ctorOpts).toEqual([]); + await expect( + client.genieGetConversationMessage({} as never), + ).resolves.toEqual({ id: "m" }); + const items: unknown[] = []; + for await (const item of ( + client as unknown as { listItemsIter(): AsyncIterable } + ).listItemsIter()) { + items.push(item); + } + expect(items).toEqual([1, 2]); + expect(ctorOpts).toHaveLength(1); + expect(Reflect.get(client, "then")).toBeUndefined(); + }); + test("no product configured (build-time) → transport adds no AppKit User-Agent", async () => { buildWarehousesClient({ host: "https://x" }); expect(await sentUserAgent(ctorOpts[0].httpClient)).toBeNull(); diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index 1eb59394a..7dd98c30f 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -19,6 +19,7 @@ import type { WarehousesClient, WorkspaceAuth, ScimClient, + GenieClient, } from "./modular"; // Legacy SDK type namespaces for un-migrated services, re-exported so AppKit @@ -30,7 +31,7 @@ import type { // that request has no `skip_cannot_use` filter, so it could pick a warehouse // the caller can't use. Statement + warehouse service types come from `./modular`. export type { files, jobs, serving, sql } from "@databricks/sdk-experimental"; -// Modular SDK client + model types (warehouses, statementExecution). +// Modular SDK client + model types (warehouses, statementExecution, genie). export type * from "./modular"; /** @@ -48,8 +49,8 @@ export interface WorkspaceClient extends WorkspaceAuth { /** SQL Warehouses (modular SDK). */ readonly warehouses: WarehousesClient; - /** Genie / dashboards. */ - readonly genie: LegacyWorkspaceClient["genie"]; + /** Genie (modular SDK). */ + readonly genie: GenieClient; /** Jobs. */ readonly jobs: LegacyWorkspaceClient["jobs"]; diff --git a/patches/@databricks__sdk-genie@0.54.0.patch b/patches/@databricks__sdk-genie@0.54.0.patch new file mode 100644 index 000000000..09e3a03ba --- /dev/null +++ b/patches/@databricks__sdk-genie@0.54.0.patch @@ -0,0 +1,26 @@ +diff --git a/dist/v1/model.d.ts b/dist/v1/model.d.ts +index ceb8ef56c1d7e4912851dbdce891623f4e8fc93c..4ff6fa02396a17485630132cf1a5bae4a3ada512 100644 +--- a/dist/v1/model.d.ts ++++ b/dist/v1/model.d.ts +@@ -1581,7 +1581,7 @@ interface ResultData { + * The `JSON_ARRAY` format is an array of arrays of values, where each non-null value is + * formatted as a string. Null values are encoded as JSON `null`. + */ +- dataArray?: ListValue[] | undefined; ++ dataArray?: (string | null)[][] | undefined; + /** The position within the sequence of result set chunks. */ + chunkIndex?: number | undefined; + /** The starting row offset within the result set. */ +diff --git a/dist/v1/model.js b/dist/v1/model.js +index 8770051f3102f70625db5758166d7fc1636afcb7..d20c134274fde0ffbf7d67a8346765fa7cba7f76 100644 +--- a/dist/v1/model.js ++++ b/dist/v1/model.js +@@ -846,7 +846,7 @@ const unmarshalResultSchema = z.object({ + })); + const unmarshalResultDataSchema = z.object({ + external_links: z.array(z.lazy(() => unmarshalExternalLinkSchema)).optional(), +- data_array: z.array(z.lazy(() => unmarshalListValueSchema)).optional(), ++ data_array: z.array(z.array(z.string().nullable())).optional(), + chunk_index: z.number().optional(), + row_offset: z.union([ + z.number(), diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index ba0ae702c..fd69b375f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -12,6 +12,7 @@ overrides: size-sensor: 1.0.3 patchedDependencies: + '@databricks/sdk-genie@0.54.0': 664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54 '@databricks/sdk-statementexecution@0.52.0': 9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c importers: @@ -288,6 +289,9 @@ importers: '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-genie': + specifier: 0.54.0 + version: 0.54.0(patch_hash=664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54) '@databricks/sdk-options': specifier: 0.51.0 version: 0.51.0 @@ -621,6 +625,9 @@ importers: '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-genie': + specifier: 0.54.0 + version: 0.54.0(patch_hash=664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54) '@databricks/sdk-options': specifier: 0.51.0 version: 0.51.0 @@ -2014,6 +2021,10 @@ packages: resolution: {integrity: sha512-dOJIt4F2nBk6HKObnv7Xbmy/qLYTy2835qhXSuW0Qw1QAXui9plmCet1KqG3yeQcMTyncWGbnhjGdQi8GEGQSA==} engines: {node: '>=22.0', npm: '>=10.0.0'} + '@databricks/sdk-genie@0.54.0': + resolution: {integrity: sha512-xT+dyuXocgwPJ5VA1Kr5/dSdKzzR9hRAkEOybeO2i+v891kuuFzVd/M9OM31PtO2/YZp95vdFH5Sb+jVGtm4nw==} + engines: {node: '>=22.0.0'} + '@databricks/sdk-options@0.51.0': resolution: {integrity: sha512-p5uBh64Y1onnvwlEwfRsmKmGRI23Y4Ly8fZiCgNUwVDe4Z9d0ctI0mo1v4gnDqlmBaBml3TZfqn5BJBROCXsEw==} engines: {node: '>=22.0.0'} @@ -14326,6 +14337,15 @@ snapshots: transitivePeerDependencies: - supports-color + '@databricks/sdk-genie@0.54.0(patch_hash=664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54)': + dependencies: + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 + '@databricks/sdk-options': 0.51.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + '@databricks/sdk-options@0.51.0': dependencies: '@databricks/sdk-auth': 0.51.0 diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index bc863eab6..3f48fc9af 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -19,4 +19,5 @@ allowBuilds: esbuild: true protobufjs: false patchedDependencies: + '@databricks/sdk-genie@0.54.0': patches/@databricks__sdk-genie@0.54.0.patch '@databricks/sdk-statementexecution@0.52.0': patches/@databricks__sdk-statementexecution@0.52.0.patch From e297a74f712402d496a8c3d4eeb30cd441794148 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 9 Oct 2026 12:47:18 +0200 Subject: [PATCH 2/4] fix(appkit): serialize bigint in SSE events The modular SDK returns int64 fields as bigint, which JSON.stringify throws on. Add a shared sseJsonReplacer used by both SSE serialization sites (stream-manager size check and SSEWriter.writeEvent) that emits them as numbers, matching the legacy SDK's raw JSON. The genie connector's toWireShape now only renames keys back to the snake_case contract, with a TODO to move appkit-ui to the camelCase shape and drop the conversion. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../api/appkit/Interface.WorkspaceClient.md | 94 +++++++++++++++++-- .../appkit/src/connectors/genie/client.ts | 14 ++- .../src/connectors/genie/tests/client.test.ts | 9 +- packages/appkit/src/stream/sse-writer.ts | 11 ++- packages/appkit/src/stream/stream-manager.ts | 4 +- .../appkit/src/stream/tests/stream.test.ts | 12 +++ 6 files changed, 123 insertions(+), 21 deletions(-) diff --git a/docs/docs/api/appkit/Interface.WorkspaceClient.md b/docs/docs/api/appkit/Interface.WorkspaceClient.md index 9d3fec157..c926698c1 100644 --- a/docs/docs/api/appkit/Interface.WorkspaceClient.md +++ b/docs/docs/api/appkit/Interface.WorkspaceClient.md @@ -7,6 +7,10 @@ migrated one at a time behind this stable interface. Accessors are legacy-typed for now (delegated to the underlying legacy SDK client); see the module docblock. +## Extends + +- `WorkspaceAuth` + ## Properties ### apiClient @@ -15,9 +19,9 @@ client); see the module docblock. readonly apiClient: ApiClient; ``` -Low-level HTTP transport (`apiClient.request(...)`). Used for endpoints -without a typed service method: SCIM header probe, warehouse listing, -serving SSE streaming, vector search, internal telemetry. +Legacy low-level HTTP transport. Prefer `request(...)` (modular, inherited +from `WorkspaceAuth`); still used by serving SSE streaming, vector search, +and the agents adapters, which take a structural `apiClient` shape. *** @@ -27,19 +31,19 @@ serving SSE streaming, vector search, internal telemetry. readonly config: Config; ``` -SDK `Config` — exposes `host` and `authenticate(headers)`. Used by the -files-upload path and agents auth-header stamping, which bypass the typed -services. +Legacy SDK `Config`. Prefer `getHost()` / `authenticate(headers)` (modular, +inherited from `WorkspaceAuth`); kept for structural `WorkspaceClientLike` +callers (supervisor adapter) until they migrate. *** ### currentUser ```ts -readonly currentUser: CurrentUserService; +readonly currentUser: ScimClient; ``` -Current user. +Current user (modular SDK SCIM client; `me({})` returns the caller). *** @@ -103,6 +107,80 @@ SQL Warehouses (modular SDK). ## Methods +### authenticate() + +```ts +authenticate(headers: Headers): Promise; +``` + +Set the auth header(s) (e.g. `Authorization`) on `headers`. + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `headers` | `Headers` | + +#### Returns + +`Promise`\<`void`\> + +#### Inherited from + +```ts +WorkspaceAuth.authenticate +``` + +*** + +### getHost() + +```ts +getHost(): Promise; +``` + +Scheme-normalized workspace host, without a trailing slash. + +#### Returns + +`Promise`\<`string`\> + +#### Inherited from + +```ts +WorkspaceAuth.getHost +``` + +*** + +### request() + +```ts +request(req: WorkspaceRequest): Promise; +``` + +Send a request through the modular transport (AppKit User-Agent + auth). +Returns the raw `Response` (body unread, so it can stream); throws +ApiError on a non-2xx status. + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `req` | `WorkspaceRequest` | + +#### Returns + +`Promise`\<`Response`\> + +#### Inherited from + +```ts +WorkspaceAuth.request +``` + +*** + ### toLegacyWorkspaceClient() ```ts diff --git a/packages/appkit/src/connectors/genie/client.ts b/packages/appkit/src/connectors/genie/client.ts index ed77a5998..84b3b291f 100644 --- a/packages/appkit/src/connectors/genie/client.ts +++ b/packages/appkit/src/connectors/genie/client.ts @@ -69,13 +69,17 @@ function toMessageResponse(message: GenieMessage): GenieMessageResponse { } /** - * The modular SDK returns the statement response camelCased with int64 fields as - * `bigint`. The SSE contract (`GenieStatementResponse`, read by appkit-ui) is the - * raw snake_case API shape, and `JSON.stringify` throws on `bigint`, so convert - * back here. Recurses into objects only; `data_array` rows pass through as-is. + * The modular SDK returns the statement response camelCased, but the SSE + * contract (`GenieStatementResponse`, read by appkit-ui) is the raw snake_case + * API shape, so rename keys back here. `bigint` counts are handled by the SSE + * serializer (`sseJsonReplacer`). Recurses into objects only; `data_array` rows + * pass through as-is. + * + * TODO: update appkit-ui (`genie-query-transform.ts`) and `GenieStatementResponse` + * to accept the camelCase shape, then drop this conversion (non-breaking path: + * read both shapes first, switch the server later). */ function toWireShape(value: unknown): unknown { - if (typeof value === "bigint") return Number(value); if (Array.isArray(value)) return value.map(toWireShape); if (value === null || typeof value !== "object") return value; return Object.fromEntries( diff --git a/packages/appkit/src/connectors/genie/tests/client.test.ts b/packages/appkit/src/connectors/genie/tests/client.test.ts index 287bef9d1..a9a9c5350 100644 --- a/packages/appkit/src/connectors/genie/tests/client.test.ts +++ b/packages/appkit/src/connectors/genie/tests/client.test.ts @@ -169,8 +169,8 @@ describe("GenieConnector", () => { const completedMsg = makeGenieMessageWithQuery(); mockStart(ws, completedMsg); - // The modular SDK hands back camelCase with int64 counts as bigint; the - // SSE payload must be the snake_case wire shape and JSON-serializable. + // The modular SDK hands back camelCase; the SSE payload keeps the + // snake_case wire shape (bigints are serialized by the SSE writer). ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ statementResponse: { statementId: "stmt-1", @@ -185,9 +185,9 @@ describe("GenieConnector", () => { statement_id: "stmt-1", manifest: { schema: { columns: [{ name: "total", type_name: "DOUBLE" }] }, - total_row_count: 1, + total_row_count: 1n, }, - result: { data_array: [["1234.56", null]], row_count: 1 }, + result: { data_array: [["1234.56", null]], row_count: 1n }, }; const events = await collect( @@ -201,7 +201,6 @@ describe("GenieConnector", () => { statementId: "stmt-1", data: statementResponse, }); - expect(() => JSON.stringify(queryResult)).not.toThrow(); }); test("maps a FAILED poll to the table permissions error", async () => { diff --git a/packages/appkit/src/stream/sse-writer.ts b/packages/appkit/src/stream/sse-writer.ts index 202f99536..b2c3372dc 100644 --- a/packages/appkit/src/stream/sse-writer.ts +++ b/packages/appkit/src/stream/sse-writer.ts @@ -9,6 +9,15 @@ import { } from "./types"; import { StreamValidator } from "./validator"; +/** + * `JSON.stringify` replacer for SSE events: the modular SDK returns int64 fields + * as `bigint`, which `JSON.stringify` throws on. Emit them as numbers, matching + * what the legacy SDK's raw JSON produced. + */ +export function sseJsonReplacer(_key: string, value: unknown): unknown { + return typeof value === "bigint" ? Number(value) : value; +} + export class SSEWriter { // setup SSE headers setupHeaders(res: IAppResponse): void { @@ -29,7 +38,7 @@ export class SSEWriter { if (res.writableEnded) return; const eventType = StreamValidator.sanitizeEventType(event.type); - const eventData = JSON.stringify(event); + const eventData = JSON.stringify(event, sseJsonReplacer); res.write(`id: ${eventId}\n`); res.write(`event: ${eventType}\n`); diff --git a/packages/appkit/src/stream/stream-manager.ts b/packages/appkit/src/stream/stream-manager.ts index eade24ec2..a483ef131 100644 --- a/packages/appkit/src/stream/stream-manager.ts +++ b/packages/appkit/src/stream/stream-manager.ts @@ -10,7 +10,7 @@ import { IdentityExpiredError } from "../errors/identity-expired"; import { createLogger } from "../logging/logger"; import { EventRingBuffer } from "./buffers"; import { streamDefaults } from "./defaults"; -import { SSEWriter } from "./sse-writer"; +import { SSEWriter, sseJsonReplacer } from "./sse-writer"; import { StreamRegistry } from "./stream-registry"; import { clearGraceTimer, clearRemovalTimer } from "./timers"; import { SSEErrorCode, type StreamEntry, type StreamOperation } from "./types"; @@ -273,7 +273,7 @@ export class StreamManager { for await (const event of streamEntry.generator) { if (streamEntry.abortController.signal.aborted) break; const eventId = randomUUID(); - const eventData = JSON.stringify(event); + const eventData = JSON.stringify(event, sseJsonReplacer); const { maxEventSize } = streamEntry; // UTF-8 bytes, not `String.length`: non-ASCII payloads used to slip diff --git a/packages/appkit/src/stream/tests/stream.test.ts b/packages/appkit/src/stream/tests/stream.test.ts index 15b3b63d6..0e4b021f6 100644 --- a/packages/appkit/src/stream/tests/stream.test.ts +++ b/packages/appkit/src/stream/tests/stream.test.ts @@ -97,6 +97,18 @@ describe("StreamManager", () => { expect(streamManager.getActiveCount()).toBe(0); }); + test("serializes bigint fields (modular SDK int64) as numbers", async () => { + const { mockRes, events } = createMockResponse(); + + async function* generator() { + yield { type: "result", rowCount: 2n }; + } + + await streamManager.stream(mockRes as any, generator); + + expect(events).toContain('data: {"type":"result","rowCount":2}\n\n'); + }); + test("should use default message type when not provided", async () => { const { mockRes, events } = createMockResponse(); From 14cfd4824f75911bc564ac5ab557a0a7af53f0b3 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 9 Oct 2026 15:19:37 +0200 Subject: [PATCH 3/4] feat(appkit-ui): read camelCase genie statement responses Genie query results come from the modular SDK in camelCase. Teach transformGenieData to read result.dataArray and column typeName, and keep accepting the legacy snake_case shape so it still works against servers that predate the migration. GenieStatementResponse gains the camelCase fields; the snake_case ones are kept as @deprecated optional fields. The playground's chart-inference fixtures move to the new shape. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../src/routes/chart-inference.route.tsx | 48 +++++++++---------- .../__tests__/genie-query-transform.test.ts | 39 +++++++++------ .../src/react/genie/genie-chart-inference.ts | 2 +- .../src/react/genie/genie-query-transform.ts | 22 +++++---- packages/shared/src/genie.ts | 17 +++++-- 5 files changed, 75 insertions(+), 53 deletions(-) diff --git a/apps/dev-playground/client/src/routes/chart-inference.route.tsx b/apps/dev-playground/client/src/routes/chart-inference.route.tsx index f799c4d61..49ec302bd 100644 --- a/apps/dev-playground/client/src/routes/chart-inference.route.tsx +++ b/apps/dev-playground/client/src/routes/chart-inference.route.tsx @@ -12,12 +12,12 @@ export const Route = createFileRoute("/chart-inference")({ }); // --------------------------------------------------------------------------- -// Helper to build a Genie-shaped statement_response from simple definitions +// Helper to build a Genie-shaped statement response from simple definitions // --------------------------------------------------------------------------- interface SampleColumn { name: string; - type_name: string; + typeName: string; } function makeStatementResponse( @@ -26,7 +26,7 @@ function makeStatementResponse( ) { return { manifest: { schema: { columns } }, - result: { data_array: rows }, + result: { dataArray: rows }, }; } @@ -46,8 +46,8 @@ const SAMPLES: { expected: "line", data: makeStatementResponse( [ - { name: "date", type_name: "DATE" }, - { name: "revenue", type_name: "DECIMAL" }, + { name: "date", typeName: "DATE" }, + { name: "revenue", typeName: "DECIMAL" }, ], [ ["2024-01-01", "12000"], @@ -71,8 +71,8 @@ const SAMPLES: { expected: "pie", data: makeStatementResponse( [ - { name: "region", type_name: "STRING" }, - { name: "sales", type_name: "DECIMAL" }, + { name: "region", typeName: "STRING" }, + { name: "sales", typeName: "DECIMAL" }, ], [ ["North America", "45000"], @@ -87,8 +87,8 @@ const SAMPLES: { expected: "bar", data: makeStatementResponse( [ - { name: "product", type_name: "STRING" }, - { name: "revenue", type_name: "DECIMAL" }, + { name: "product", typeName: "STRING" }, + { name: "revenue", typeName: "DECIMAL" }, ], Array.from({ length: 15 }, (_, i) => [ `Product ${String.fromCharCode(65 + i)}`, @@ -102,8 +102,8 @@ const SAMPLES: { expected: "line", data: makeStatementResponse( [ - { name: "city", type_name: "STRING" }, - { name: "population", type_name: "INT" }, + { name: "city", typeName: "STRING" }, + { name: "population", typeName: "INT" }, ], Array.from({ length: 150 }, (_, i) => [ `City ${i + 1}`, @@ -117,9 +117,9 @@ const SAMPLES: { expected: "line", data: makeStatementResponse( [ - { name: "month", type_name: "DATE" }, - { name: "revenue", type_name: "DECIMAL" }, - { name: "cost", type_name: "DECIMAL" }, + { name: "month", typeName: "DATE" }, + { name: "revenue", typeName: "DECIMAL" }, + { name: "cost", typeName: "DECIMAL" }, ], [ ["2024-01-01", "12000", "8000"], @@ -137,9 +137,9 @@ const SAMPLES: { expected: "bar", data: makeStatementResponse( [ - { name: "department", type_name: "STRING" }, - { name: "budget", type_name: "DECIMAL" }, - { name: "actual", type_name: "DECIMAL" }, + { name: "department", typeName: "STRING" }, + { name: "budget", typeName: "DECIMAL" }, + { name: "actual", typeName: "DECIMAL" }, ], [ ["Engineering", "500000", "480000"], @@ -159,8 +159,8 @@ const SAMPLES: { expected: "scatter", data: makeStatementResponse( [ - { name: "height_cm", type_name: "DOUBLE" }, - { name: "weight_kg", type_name: "DOUBLE" }, + { name: "height_cm", typeName: "DOUBLE" }, + { name: "weight_kg", typeName: "DOUBLE" }, ], Array.from({ length: 30 }, (_, i) => [ String(150 + i * 1.2), @@ -174,8 +174,8 @@ const SAMPLES: { expected: "none (table only)", data: makeStatementResponse( [ - { name: "metric", type_name: "STRING" }, - { name: "value", type_name: "DECIMAL" }, + { name: "metric", typeName: "STRING" }, + { name: "value", typeName: "DECIMAL" }, ], [["Total Revenue", "125000"]], ), @@ -186,9 +186,9 @@ const SAMPLES: { expected: "none (table only)", data: makeStatementResponse( [ - { name: "first_name", type_name: "STRING" }, - { name: "last_name", type_name: "STRING" }, - { name: "city", type_name: "STRING" }, + { name: "first_name", typeName: "STRING" }, + { name: "last_name", typeName: "STRING" }, + { name: "city", typeName: "STRING" }, ], [ ["Alice", "Smith", "New York"], diff --git a/packages/appkit-ui/src/react/genie/__tests__/genie-query-transform.test.ts b/packages/appkit-ui/src/react/genie/__tests__/genie-query-transform.test.ts index b7ab70134..697dfc184 100644 --- a/packages/appkit-ui/src/react/genie/__tests__/genie-query-transform.test.ts +++ b/packages/appkit-ui/src/react/genie/__tests__/genie-query-transform.test.ts @@ -56,20 +56,20 @@ describe("classifySqlType", () => { describe("transformGenieData", () => { function makeResponse( - columns: Array<{ name: string; type_name: string }>, + columns: Array<{ name: string; typeName: string }>, dataArray: (string | null)[][], ) { return { manifest: { schema: { columns } }, - result: { data_array: dataArray }, + result: { dataArray: dataArray }, }; } test("transforms basic numeric and string data", () => { const data = makeResponse( [ - { name: "region", type_name: "STRING" }, - { name: "sales", type_name: "DECIMAL" }, + { name: "region", typeName: "STRING" }, + { name: "sales", typeName: "DECIMAL" }, ], [ ["North", "1000.50"], @@ -99,8 +99,8 @@ describe("transformGenieData", () => { test("handles date columns as strings", () => { const data = makeResponse( [ - { name: "day", type_name: "DATE" }, - { name: "revenue", type_name: "INT" }, + { name: "day", typeName: "DATE" }, + { name: "revenue", typeName: "INT" }, ], [["2024-01-15", "500"]], ); @@ -113,8 +113,8 @@ describe("transformGenieData", () => { test("handles null values", () => { const data = makeResponse( [ - { name: "name", type_name: "STRING" }, - { name: "value", type_name: "INT" }, + { name: "name", typeName: "STRING" }, + { name: "value", typeName: "INT" }, ], [ [null, "10"], @@ -132,8 +132,8 @@ describe("transformGenieData", () => { test("handles non-numeric strings in numeric columns", () => { const data = makeResponse( [ - { name: "name", type_name: "STRING" }, - { name: "value", type_name: "INT" }, + { name: "name", typeName: "STRING" }, + { name: "value", typeName: "INT" }, ], [["a", "not_a_number"]], ); @@ -142,8 +142,8 @@ describe("transformGenieData", () => { expect(result?.rows[0].value).toBeNull(); }); - test("returns null for empty data_array", () => { - const data = makeResponse([{ name: "a", type_name: "STRING" }], []); + test("returns null for empty dataArray", () => { + const data = makeResponse([{ name: "a", typeName: "STRING" }], []); expect(transformGenieData(data)).toBeNull(); }); @@ -151,11 +151,20 @@ describe("transformGenieData", () => { expect( transformGenieData({ manifest: { schema: { columns: [] } }, - result: { data_array: [["x"]] }, + result: { dataArray: [["x"]] }, }), ).toBeNull(); }); + test("accepts the legacy snake_case shape", () => { + const result = transformGenieData({ + manifest: { schema: { columns: [{ name: "n", type_name: "INT" }] } }, + result: { data_array: [["42"]] }, + }); + expect(result?.rows).toEqual([{ n: 42 }]); + expect(result?.columns[0]).toMatchObject({ typeName: "INT" }); + }); + test("returns null for null/undefined input", () => { expect(transformGenieData(null)).toBeNull(); expect(transformGenieData(undefined)).toBeNull(); @@ -164,8 +173,8 @@ describe("transformGenieData", () => { test("handles rows shorter than columns (missing cells)", () => { const data = makeResponse( [ - { name: "a", type_name: "STRING" }, - { name: "b", type_name: "INT" }, + { name: "a", typeName: "STRING" }, + { name: "b", typeName: "INT" }, ], [["hello"]], ); diff --git a/packages/appkit-ui/src/react/genie/genie-chart-inference.ts b/packages/appkit-ui/src/react/genie/genie-chart-inference.ts index a7f9bc08f..5a8e954d4 100644 --- a/packages/appkit-ui/src/react/genie/genie-chart-inference.ts +++ b/packages/appkit-ui/src/react/genie/genie-chart-inference.ts @@ -5,7 +5,7 @@ * │ These rules determine what chart type is shown for Genie query │ * │ results. Modify thresholds and chart type mappings here. │ * │ │ - * │ Column types are classified from SQL type_name: │ + * │ Column types are classified from SQL typeName: │ * │ DATE: DATE, TIMESTAMP, TIMESTAMP_NTZ │ * │ NUMERIC: DECIMAL, INT, DOUBLE, FLOAT, LONG, etc. │ * │ STRING: STRING, VARCHAR, CHAR │ diff --git a/packages/appkit-ui/src/react/genie/genie-query-transform.ts b/packages/appkit-ui/src/react/genie/genie-query-transform.ts index d763db41c..db1b9cf89 100644 --- a/packages/appkit-ui/src/react/genie/genie-query-transform.ts +++ b/packages/appkit-ui/src/react/genie/genie-query-transform.ts @@ -2,8 +2,8 @@ * Converts Genie's statement_response data into a flat record array * suitable for charting. * - * The Genie API returns `{ manifest.schema.columns, result.data_array }` - * where each column carries a SQL `type_name`. This module parses values + * The Genie API returns `{ manifest.schema.columns, result.dataArray }` + * where each column carries a SQL `typeName`. This module parses values * according to those types so downstream chart code receives proper * numbers and strings. */ @@ -70,11 +70,14 @@ function parseValue(raw: string | null, category: ColumnCategory): unknown { * Expects `data` to have the shape: * ``` * { - * manifest: { schema: { columns: [{ name, type_name }, ...] } }, - * result: { data_array: [["val", ...], ...] } + * manifest: { schema: { columns: [{ name, typeName }, ...] } }, + * result: { dataArray: [["val", ...], ...] } * } * ``` * + * The legacy snake_case shape (`type_name`, `data_array`) is also accepted, so + * this keeps working against servers that predate the modular SDK migration. + * * Returns `null` when the data is empty or malformed. */ export function transformGenieData( @@ -87,16 +90,15 @@ export function transformGenieData( return null; } - const dataArray = data.result?.data_array; + const dataArray = data.result?.dataArray ?? data.result?.data_array; if (!dataArray || dataArray.length === 0) { return null; } - const columns: GenieColumnMeta[] = rawColumns.map((col) => ({ - name: col.name, - typeName: col.type_name, - category: classifySqlType(col.type_name), - })); + const columns: GenieColumnMeta[] = rawColumns.map((col) => { + const typeName = col.typeName ?? col.type_name ?? ""; + return { name: col.name, typeName, category: classifySqlType(typeName) }; + }); const rows: Record[] = dataArray.map((row) => { const record: Record = {}; diff --git a/packages/shared/src/genie.ts b/packages/shared/src/genie.ts index 9b568fa57..9ea830a55 100644 --- a/packages/shared/src/genie.ts +++ b/packages/shared/src/genie.ts @@ -25,15 +25,26 @@ export type GenieStreamEvent = loadedCount: number; }; -/** Shape of the Databricks SQL statement_response returned by Genie queries */ +/** + * Shape of the Databricks SQL statement response returned by Genie queries. + * The server sends the modular SDK's camelCase shape; the snake_case fields are + * the legacy raw-API shape, still accepted by appkit-ui for older servers. + */ export interface GenieStatementResponse { manifest: { schema: { - columns: Array<{ name: string; type_name: string }>; + columns: Array<{ + name: string; + typeName?: string; + /** @deprecated Legacy snake_case shape; use `typeName`. */ + type_name?: string; + }>; }; }; result: { - data_array: (string | null)[][]; + dataArray?: (string | null)[][]; + /** @deprecated Legacy snake_case shape; use `dataArray`. */ + data_array?: (string | null)[][]; }; } From 991979a8a9cd236eb8cda805d7f42e7a4d84202b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 9 Oct 2026 15:20:54 +0200 Subject: [PATCH 4/4] feat(appkit): stream genie query results in the modular camelCase shape Drop the connector's toWireShape conversion and send the modular SDK's statement response as-is (dataArray, typeName, ...). bigint counts are already handled by the SSE serializer, and appkit-ui reads both shapes. Custom clients reading the genie query_result SSE payload directly need to switch to the camelCase fields. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../appkit/src/connectors/genie/client.ts | 24 +----------------- .../src/connectors/genie/tests/client.test.ts | 25 +++++++------------ 2 files changed, 10 insertions(+), 39 deletions(-) diff --git a/packages/appkit/src/connectors/genie/client.ts b/packages/appkit/src/connectors/genie/client.ts index 84b3b291f..73d56d5f6 100644 --- a/packages/appkit/src/connectors/genie/client.ts +++ b/packages/appkit/src/connectors/genie/client.ts @@ -68,28 +68,6 @@ function toMessageResponse(message: GenieMessage): GenieMessageResponse { }; } -/** - * The modular SDK returns the statement response camelCased, but the SSE - * contract (`GenieStatementResponse`, read by appkit-ui) is the raw snake_case - * API shape, so rename keys back here. `bigint` counts are handled by the SSE - * serializer (`sseJsonReplacer`). Recurses into objects only; `data_array` rows - * pass through as-is. - * - * TODO: update appkit-ui (`genie-query-transform.ts`) and `GenieStatementResponse` - * to accept the camelCase shape, then drop this conversion (non-breaking path: - * read both shapes first, switch the server later). - */ -function toWireShape(value: unknown): unknown { - if (Array.isArray(value)) return value.map(toWireShape); - if (value === null || typeof value !== "object") return value; - return Object.fromEntries( - Object.entries(value).map(([key, v]) => [ - key.replace(/[A-Z]/g, (c) => `_${c.toLowerCase()}`), - toWireShape(v), - ]), - ); -} - function classifyGenieError(error: unknown): string { const message = error instanceof Error ? error.message : String(error); // Modular ApiError carries the code on `.code`, legacy on `.errorCode`. @@ -257,7 +235,7 @@ export class GenieConnector { messageId, attachmentId, }); - return toWireShape(response.statementResponse) as GenieStatementResponse; + return response.statementResponse as GenieStatementResponse; } async *streamSendMessage( diff --git a/packages/appkit/src/connectors/genie/tests/client.test.ts b/packages/appkit/src/connectors/genie/tests/client.test.ts index a9a9c5350..00deb2998 100644 --- a/packages/appkit/src/connectors/genie/tests/client.test.ts +++ b/packages/appkit/src/connectors/genie/tests/client.test.ts @@ -169,26 +169,19 @@ describe("GenieConnector", () => { const completedMsg = makeGenieMessageWithQuery(); mockStart(ws, completedMsg); - // The modular SDK hands back camelCase; the SSE payload keeps the - // snake_case wire shape (bigints are serialized by the SSE writer). - ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ - statementResponse: { - statementId: "stmt-1", - manifest: { - schema: { columns: [{ name: "total", typeName: "DOUBLE" }] }, - totalRowCount: 1n, - }, - result: { dataArray: [["1234.56", null]], rowCount: 1n }, - }, - }); + // The SDK's camelCase statement response is streamed as-is; bigints are + // serialized by the SSE writer. const statementResponse = { - statement_id: "stmt-1", + statementId: "stmt-1", manifest: { - schema: { columns: [{ name: "total", type_name: "DOUBLE" }] }, - total_row_count: 1n, + schema: { columns: [{ name: "total", typeName: "DOUBLE" }] }, + totalRowCount: 1n, }, - result: { data_array: [["1234.56", null]], row_count: 1n }, + result: { dataArray: [["1234.56", null]], rowCount: 1n }, }; + ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse, + }); const events = await collect( connector.streamSendMessage(ws, "space-1", "query", undefined),