diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 86b60986..a9c688da 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -7,11 +7,7 @@ import type { StreamExecutionSettings, } from "shared"; import { SQLWarehouseConnector } from "../../connectors"; -import { - getCurrentUserId, - getWarehouseId, - getWorkspaceClient, -} from "../../context"; +import { getWarehouseId, getWorkspaceClient } from "../../context"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; @@ -152,8 +148,7 @@ export class AnalyticsPlugin extends Plugin { // get execution context - user-scoped if .obo.sql, otherwise service principal const executor = isAsUser ? this.asUser(req) : this; - const userKey = getCurrentUserId(); - const executorKey = isAsUser ? userKey : "global"; + const executorKey = isAsUser ? this.resolveUserId(req) : "global"; const queryParameters = format === "ARROW" diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 71a8fd98..9a30440e 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -483,6 +483,65 @@ describe("Analytics Plugin", () => { ); }); + test("OBO cache key must use the end user's ID, not the service principal's", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM my_data", + isAsUser: true, + }); + + const executeMock = vi + .fn() + .mockResolvedValueOnce({ + result: { data: [{ owner: "alice-data" }] }, + }) + .mockResolvedValueOnce({ + result: { data: [{ owner: "bob-data" }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + + // User Alice makes an OBO query + const aliceReq = createMockRequest({ + params: { query_key: "my_data" }, + body: { parameters: {} }, + headers: { + "x-forwarded-access-token": "alice-token", + "x-forwarded-user": "alice", + }, + }); + const aliceRes = createMockResponse(); + await handler(aliceReq, aliceRes); + + // User Bob makes the SAME OBO query with the SAME parameters + const bobReq = createMockRequest({ + params: { query_key: "my_data" }, + body: { parameters: {} }, + headers: { + "x-forwarded-access-token": "bob-token", + "x-forwarded-user": "bob", + }, + }); + const bobRes = createMockResponse(); + await handler(bobReq, bobRes); + + // Both queries must execute — different users must not share OBO cache + expect(executeMock).toHaveBeenCalledTimes(2); + + // Alice sees her own data + expect(aliceRes.write).toHaveBeenCalledWith( + expect.stringContaining('"owner":"alice-data"'), + ); + // Bob sees his own data, NOT Alice's cached result + expect(bobRes.write).toHaveBeenCalledWith( + expect.stringContaining('"owner":"bob-data"'), + ); + }); + test("should handle AbortSignal cancellation", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter();