diff --git a/packages/appkit-ui/src/react/hooks/__tests__/analytics-request-store.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/analytics-request-store.test.ts new file mode 100644 index 000000000..601ae408a --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/__tests__/analytics-request-store.test.ts @@ -0,0 +1,276 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +let capturedCallbacks: { + onMessage?: (msg: { data: string }) => void; + onError?: (err: Error) => void; + signal?: AbortSignal; +} = {}; + +const mockFetchArrow = vi.fn(); + +// Mock connectSSE to capture calls and track uncached requests. +const mockConnectSSE = vi.fn((args: any): unknown => { + capturedCallbacks = { + onMessage: args?.onMessage, + onError: args?.onError, + signal: args?.signal, + }; + return () => {}; +}); + +const mockProcessArrowBuffer = vi.fn(); + +vi.mock("@/js", () => ({ + connectSSE: (...args: unknown[]) => mockConnectSSE(...(args as [any])), + ArrowClient: { + fetchArrow: (...args: unknown[]) => mockFetchArrow(...args), + processArrowBuffer: (...args: unknown[]) => mockProcessArrowBuffer(...args), + }, +})); + +import { + getSnapshot, + refetch, + resetAnalyticsRequestStore, + retain, + start, + subscribe, +} from "../analytics-request-store"; + +const JSON_OPTS = { + url: "/api/analytics/query/q", + payload: JSON.stringify({ parameters: null, format: "JSON_ARRAY" }), + format: "JSON_ARRAY", +}; + +describe("analytics-request-store", () => { + beforeEach(() => { + vi.clearAllMocks(); + capturedCallbacks = {}; + resetAnalyticsRequestStore(); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + describe("retain and start", () => { + test("retain with autoStart:false does not start the request", () => { + retain("k", JSON_OPTS, false); + expect(mockConnectSSE).not.toHaveBeenCalled(); + }); + + test("start on an autoStart:false entry begins the request", () => { + retain("k", JSON_OPTS, false); + expect(mockConnectSSE).not.toHaveBeenCalled(); + + start("k"); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + }); + + test("retain with default autoStart:true starts the request immediately", () => { + retain("k", JSON_OPTS); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + }); + }); + + describe("refetch", () => { + test("refetch aborts in-flight run and restarts with same options", () => { + retain("k", JSON_OPTS); + const firstSignal = capturedCallbacks.signal; + + // refetch should call connectSSE a second time. + refetch("k"); + + expect(mockConnectSSE).toHaveBeenCalledTimes(2); + + // The second call is a fresh start with a distinct signal. + const secondCall = mockConnectSSE.mock.calls[1]; + const secondSignal = secondCall[0].signal; + + expect(secondSignal).not.toBe(firstSignal); + }); + + test("refetch marks the key for uncached execution on next start", () => { + const release = retain("k", JSON_OPTS, false); + expect(mockConnectSSE).not.toHaveBeenCalled(); + + // Mark for uncached. + refetch("k"); + + // start() is called; the entry should be re-invoked with skipCache set. + // Check that the payload now includes skipCache:true. + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + const call = mockConnectSSE.mock.calls[0]; + const payload = call[0].payload; + const parsed = JSON.parse(payload); + expect(parsed.skipCache).toBe(true); + + release(); + }); + + test("skipCache flag is included in payload for SSE requests", () => { + const sseOpts = { + url: "/api/analytics/query/test", + payload: JSON.stringify({ parameters: { x: 1 }, format: "JSON_ARRAY" }), + format: "JSON_ARRAY", + skipCache: true, + }; + + retain("ssekey", sseOpts); + + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + const call = mockConnectSSE.mock.calls[0]; + const payload = call[0].payload; + const parsed = JSON.parse(payload); + + expect(parsed).toEqual({ + parameters: { x: 1 }, + format: "JSON_ARRAY", + skipCache: true, + }); + }); + + test("skipCache flag is included in payload for ARROW_STREAM requests", async () => { + const fakeTable = { numRows: 1, schema: { fields: [] } }; + const fakeBytes = new Uint8Array([1, 2, 3]); + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + arrayBuffer: async () => fakeBytes.buffer, + headers: { get: () => null }, + }); + vi.stubGlobal("fetch", fetchMock); + mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); + + const arrowOpts = { + url: "/api/analytics/query/arrow_test", + payload: JSON.stringify({ parameters: null, format: "ARROW_STREAM" }), + format: "ARROW_STREAM", + skipCache: true, + }; + + retain("arrowkey", arrowOpts); + + // Wait for the arrow fetch to be called. + await new Promise((r) => setTimeout(r, 10)); + + expect(fetchMock).toHaveBeenCalledTimes(1); + const [, init] = fetchMock.mock.calls[0]; + const payload = JSON.parse(init.body); + + expect(payload).toEqual({ + parameters: null, + format: "ARROW_STREAM", + skipCache: true, + }); + }); + + test("uncached mark is consumed on next start and does not persist", () => { + const release1 = retain("k", JSON_OPTS, false); + + // Mark for uncached. + refetch("k"); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + let payload = JSON.parse(mockConnectSSE.mock.calls[0][0].payload); + expect(payload.skipCache).toBe(true); + + release1(); + + // Re-retain the same key without refetch; the uncached mark should not persist. + resetAnalyticsRequestStore(); + vi.clearAllMocks(); + + const release2 = retain("k", JSON_OPTS); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + payload = JSON.parse(mockConnectSSE.mock.calls[0][0].payload); + expect(payload.skipCache).toBeUndefined(); + + release2(); + }); + }); + + describe("retain with autoStart:false and refetch composition", () => { + test("retain(autoStart:false) + refetch defers start but executes with uncached on refetch", () => { + const release = retain("deferred", JSON_OPTS, false); + + // No request yet. + expect(mockConnectSSE).not.toHaveBeenCalled(); + + // refetch marks for uncached and starts. + refetch("deferred"); + + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + const payload = JSON.parse(mockConnectSSE.mock.calls[0][0].payload); + expect(payload.skipCache).toBe(true); + + release(); + }); + + test("multiple refetch calls re-invoke the runner each time", () => { + retain("multi", JSON_OPTS, false); + + refetch("multi"); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + refetch("multi"); + expect(mockConnectSSE).toHaveBeenCalledTimes(2); + + refetch("multi"); + expect(mockConnectSSE).toHaveBeenCalledTimes(3); + + // All three should have skipCache. + for (let i = 0; i < 3; i++) { + const payload = JSON.parse(mockConnectSSE.mock.calls[i][0].payload); + expect(payload.skipCache).toBe(true); + } + }); + }); + + describe("snapshot subscribers", () => { + test("refetch triggers updates to subscribers via snapshot", () => { + const release = retain("snap", JSON_OPTS, false); + const listener = vi.fn(); + subscribe("snap", listener); + + refetch("snap"); + + // Refetch should have triggered a start → loading snapshot notification. + expect(listener).toHaveBeenCalled(); + expect(getSnapshot("snap").loading).toBe(true); + + release(); + }); + }); + + describe("uncached mark does not leak from a no-op refetch", () => { + test("refetch on a key with no live entry does not skipCache a later retain", () => { + // No entry exists for "ghost": store.start no-ops, so the runner never + // runs and never consumes the mark. It must not stick to the next retain. + refetch("ghost"); + expect(mockConnectSSE).not.toHaveBeenCalled(); + + const release = retain("ghost", JSON_OPTS); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + const payload = JSON.parse(mockConnectSSE.mock.calls[0][0].payload); + expect(payload.skipCache).toBeUndefined(); + + release(); + }); + + test("resetAnalyticsRequestStore clears any pending uncached mark", () => { + refetch("ghost"); + resetAnalyticsRequestStore(); + vi.clearAllMocks(); + + const release = retain("ghost", JSON_OPTS); + const payload = JSON.parse(mockConnectSSE.mock.calls[0][0].payload); + expect(payload.skipCache).toBeUndefined(); + + release(); + }); + }); +}); diff --git a/packages/appkit-ui/src/react/hooks/analytics-request-store.ts b/packages/appkit-ui/src/react/hooks/analytics-request-store.ts index 203bc0aea..d72f81ea1 100644 --- a/packages/appkit-ui/src/react/hooks/analytics-request-store.ts +++ b/packages/appkit-ui/src/react/hooks/analytics-request-store.ts @@ -34,6 +34,8 @@ interface AnalyticsRequestOptions { payload: string; /** Response format; selects the transport. */ format: string; + /** @internal Force skip cache on next execution; cleared after each run. */ + skipCache?: boolean; } /** Immutable per-key request state; mirrors the hook's public result shape. */ @@ -162,15 +164,34 @@ async function fetchArrowDirect( * format-appropriate transport, reporting state through `controls.patch`. */ function runAnalyticsRequest( + cacheKey: string, options: AnalyticsRequestOptions, ): RequestRunner { return (controls) => { controls.patch(LOADING_SNAPSHOT); + // Consume any one-shot uncached mark for this key. + const shouldSkipCache = options.skipCache || uncachedKeys.has(cacheKey); + if (uncachedKeys.has(cacheKey)) { + uncachedKeys.delete(cacheKey); + } + + // Build the request payload, injecting skipCache when needed. + let requestPayload = options.payload; + if (shouldSkipCache) { + try { + const parsed = JSON.parse(options.payload); + requestPayload = JSON.stringify({ ...parsed, skipCache: true }); + } catch { + // If payload parsing fails, use the original payload and let it fail downstream. + requestPayload = options.payload; + } + } + // ARROW_STREAM: the server streams raw Arrow IPC bytes back on the query // response body (no SSE). Fetch and decode directly. if (options.format === "ARROW_STREAM") { - void fetchArrowDirect(controls, options); + void fetchArrowDirect(controls, { ...options, payload: requestPayload }); return; } @@ -195,7 +216,7 @@ function runAnalyticsRequest( connectSSE({ url: options.url, - payload: options.payload, + payload: requestPayload, signal: controls.signal, onMessage: (message) => handleAnalyticsSseMessage(message.data, sseContext), @@ -206,6 +227,13 @@ function runAnalyticsRequest( const store = createRequestStore(EMPTY_SNAPSHOT); +/** + * Internal tracking of which cache keys should force uncached execution on their + * next run. Entries are removed after the run to avoid sticking to subsequent + * invocations with the same key. + */ +const uncachedKeys = new Set(); + /** * Register a subscriber for `key`, starting the shared request on first use. * Returns a `release` function that must be called on unmount. @@ -215,12 +243,35 @@ export function retain( options: AnalyticsRequestOptions, autoStart = true, ): () => void { - return store.retain(key, runAnalyticsRequest(options), autoStart); + return store.retain(key, runAnalyticsRequest(key, options), autoStart); } export const start = store.start; export const subscribe = store.subscribe; export const getSnapshot = store.getSnapshot; +/** + * Internal API: re-run a request with cache bypassed. + * + * Marks the key to force uncached execution on its next run, then calls + * store.start() to abort any in-flight run and restart. The mark is consumed + * immediately by the runner, so it doesn't persist across multiple refetch + * calls or independent hook instances. + * @internal + */ +export function refetch(cacheKey: string): void { + uncachedKeys.add(cacheKey); + store.start(cacheKey); + // `store.start` is a synchronous no-op when the key has no live entry (e.g. + // the last subscriber released and teardown ran). In that case the runner + // never ran and never consumed the mark, so clear it here to avoid leaking + // skipCache onto a later `retain` of the same key. When the runner did run + // it already deleted the mark synchronously, so this is a harmless no-op. + uncachedKeys.delete(cacheKey); +} + /** Test-only: abort every in-flight request and clear the store. */ -export const resetAnalyticsRequestStore = store.reset; +export function resetAnalyticsRequestStore(): void { + uncachedKeys.clear(); + store.reset(); +} diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index dc3543be4..a4b668515 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -28,7 +28,6 @@ import { AppKitError, ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; -import type { WorkspaceClient } from "../../workspace-client"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { @@ -245,8 +244,16 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { res: express.Response, ): Promise { const { query_key } = req.params; - const { parameters, format: rawFormat = "JSON_ARRAY" } = - req.body as IAnalyticsQueryRequest; + const { + parameters, + format: rawFormat = "JSON_ARRAY", + skipCache: rawSkipCache = false, + } = req.body as IAnalyticsQueryRequest; + + // Only an explicit boolean `true` triggers a cache-bypassing refresh. + // Coerce any other value (truthy strings, numbers) to false so an + // untrusted body can't enable it with a non-boolean. + const skipCache = rawSkipCache === true; if ( rawFormat !== "JSON_ARRAY" && @@ -305,6 +312,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { query, isAsUser, parameters, + skipCache, ); return; } @@ -327,6 +335,17 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { ], }; + // Refresh (polling/refetch): drop any cached entry so the execute below is + // a miss. It re-hits the warehouse, repopulates the shared entry + // (write-through, so other readers stop seeing the stale value), and still + // coalesces with concurrent refreshes via the cache's in-flight dedup — a + // plain cache-disable would do neither. + if (skipCache) { + await this.cache.delete( + this.cache.generateKey(cacheConfig.cacheKey, executorKey), + ); + } + // Cache/retry/timeout are scoped to the SQL execution itself (inner // `execute`) so the warehouse-readiness phase isn't subject to retries // and the generator value never leaks into the cache. @@ -821,6 +840,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { query: string, isAsUser: boolean, parameters: IAnalyticsQueryRequest["parameters"], + skipCache: boolean = false, ): Promise { const executor = isAsUser ? this.asUser(req) : this; const executorKey = isAsUser ? this.resolveUserId(req) : "global"; @@ -871,6 +891,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { query, parameters, executorKey, + skipCache, ), this.SQLClient, query, @@ -1002,12 +1023,13 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { query: string, parameters: IAnalyticsQueryRequest["parameters"], executorKey: string, + skipCache: boolean = false, ): QueryExecutor { const hashedQuery = this.queryProcessor.hashQuery(query); const cache = this.cache; const ttl = queryDefaults.cache?.ttl; return { - query: (q, params, formatParameters, signal) => { + query: async (q, params, formatParameters, signal) => { // Only the inline-arrow attempt is cacheable — EXTERNAL_LINKS carry // short-lived pre-signed URLs, so those pass straight through. if ( @@ -1016,19 +1038,30 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { ) { return executor.query(q, params, formatParameters, signal); } + + const arrowKey = [ + "analytics:query:arrow", + query_key, + JSON.stringify(parameters), + hashedQuery, + executorKey, + ]; + + // Refresh (polling/refetch): drop the cached entry so getOrExecute is a + // miss. It re-hits the warehouse, repopulates the shared entry + // (write-through), and still coalesces concurrent refreshes — a plain + // bypass would leave other readers stale to TTL and lose that dedup. + if (skipCache) { + await cache.delete(cache.generateKey(arrowKey, executorKey)); + } + // On a standard warehouse this throws a capability rejection — the // cache never stores a rejection, so the fallback still sees the // structured error. On Reyden it returns a bounded (<=25 MiB) // attachment that caches like the JSON path's rows. The shared signal // dedupes concurrent renders (e.g. React StrictMode double-mount). return cache.getOrExecute( - [ - "analytics:query:arrow", - query_key, - JSON.stringify(parameters), - hashedQuery, - executorKey, - ], + arrowKey, (sharedSignal) => executor.query(q, params, formatParameters, sharedSignal ?? signal), executorKey, diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 5101f9424..92bf08420 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -39,7 +39,11 @@ const { mockCacheStore, mockCacheInstance } = vi.hoisted(() => { const instance = { get: vi.fn(), set: vi.fn(), - delete: vi.fn(), + // Model the real CacheManager.delete: actually evict the entry so a + // subsequent getOrExecute misses (used by the skipCache refresh path). + delete: vi.fn(async (key: string) => { + store.delete(key); + }), getOrExecute: vi.fn( async (key: unknown[], fn: () => Promise, userKey: string) => { const cacheKey = generateKey(key, userKey); @@ -342,6 +346,106 @@ describe("Analytics Plugin", () => { expect(mockRes2.write).toHaveBeenCalledWith("event: result\n"); }); + test("skipCache refreshes the shared entry (re-hits warehouse + write-through)", async () => { + mockCacheInstance.delete.mockClear(); + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test WHERE foo = :foo", + isAsUser: false, + }); + + // Distinct results per execution so we can tell a fresh run from a + // cache hit: first the "stale" value, then the "fresh" value. + const executeMock = vi + .fn() + .mockResolvedValueOnce({ result: { data: [{ id: 1, v: "stale" }] } }) + .mockResolvedValueOnce({ result: { data: [{ id: 1, v: "fresh" }] } }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const body = { parameters: { foo: sql.string("bar") } }; + + // 1) Prime the shared cache with the "stale" value. + const res1 = createMockResponse(); + await handler( + createMockRequest({ params: { query_key: "q" }, body }), + res1, + ); + + // 2) Same params, but skipCache: must invalidate + re-execute (not served + // the stale hit), then repopulate the entry with the fresh value. + const res2 = createMockResponse(); + await handler( + createMockRequest({ + params: { query_key: "q" }, + body: { ...body, skipCache: true }, + }), + res2, + ); + + // 3) Same params, no skipCache: served from the REFRESHED entry — no new + // execution, and it returns the fresh value (write-through), not "stale". + const res3 = createMockResponse(); + await handler( + createMockRequest({ params: { query_key: "q" }, body }), + res3, + ); + + // Two executions total: initial + the skipCache refresh. Step 3 is a hit. + expect(executeMock).toHaveBeenCalledTimes(2); + expect(mockCacheInstance.delete).toHaveBeenCalledTimes(1); + + expect(res1.write).toHaveBeenCalledWith( + expect.stringContaining('"v":"stale"'), + ); + expect(res2.write).toHaveBeenCalledWith( + expect.stringContaining('"v":"fresh"'), + ); + expect(res3.write).toHaveBeenCalledWith( + expect.stringContaining('"v":"fresh"'), + ); + }); + + test("skipCache is ignored unless the body value is boolean true", async () => { + mockCacheInstance.delete.mockClear(); + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test WHERE foo = :foo", + isAsUser: false, + }); + const executeMock = vi.fn().mockResolvedValue({ + result: { data: [{ id: 1, name: "cached" }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const body = { parameters: { foo: sql.string("bar") } }; + + // Prime the cache, then repeat with a truthy-but-non-boolean skipCache. + await handler( + createMockRequest({ params: { query_key: "q" }, body }), + createMockResponse(), + ); + await handler( + createMockRequest({ + params: { query_key: "q" }, + body: { ...body, skipCache: "true" as unknown as boolean }, + }), + createMockResponse(), + ); + + // The non-boolean did not bypass: second request was a cache hit and the + // entry was never invalidated. + expect(executeMock).toHaveBeenCalledTimes(1); + expect(mockCacheInstance.delete).not.toHaveBeenCalled(); + }); + test("should share cache across users for .sql files (global cache)", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); diff --git a/packages/appkit/src/plugins/analytics/types.ts b/packages/appkit/src/plugins/analytics/types.ts index 924373f07..8fcadb1f6 100644 --- a/packages/appkit/src/plugins/analytics/types.ts +++ b/packages/appkit/src/plugins/analytics/types.ts @@ -125,6 +125,13 @@ export function normalizeAnalyticsFormat( export interface IAnalyticsQueryRequest { parameters?: Record; format?: AnalyticsFormat; + /** + * Internal flag to bypass server-side TTL cache for this request. + * When true, the request re-hits the warehouse even with identical params. + * Used by polling/refetch operations to force fresh execution. + * @internal + */ + skipCache?: boolean; } export interface AnalyticsQueryResponse {