From 726e617c96d945cebcf9e1d9866e03db180ad4b2 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 7 Aug 2026 11:41:23 +0200 Subject: [PATCH] feat(appkit-ui): share in-flight useAnalyticsQuery requests Identical analytics requests (same query key, parameters, format, and dev mode) now share a single in-flight network request instead of one per hook instance. A module-singleton request store (mirroring the ResourceStatusStore idiom) owns the transport; useAnalyticsQuery becomes a useSyncExternalStore subscriber. Late subscribers read the current snapshot; the request is torn down a tick after the last subscriber unmounts, so a StrictMode unmount->remount reuses it rather than aborting. useChartData and all charts inherit the dedup for free. Dedup-only, no result cache. UseAnalyticsQueryResult is unchanged (non-breaking). Adds a /query-dedup playground route that counts analytics fetches in-page to make the behavior observable. Closes #496 Signed-off-by: MarioCadenas --- apps/dev-playground/client/src/lib/nav.ts | 8 + .../client/src/routeTree.gen.ts | 21 + .../client/src/routes/query-dedup.route.tsx | 209 +++++++++ .../__tests__/use-analytics-query.test.ts | 122 +++++ .../use-analytics-warehouse-status.test.tsx | 6 + .../react/hooks/analytics-request-store.ts | 421 ++++++++++++++++++ .../src/react/hooks/use-analytics-query.ts | 374 +++------------- 7 files changed, 848 insertions(+), 313 deletions(-) create mode 100644 apps/dev-playground/client/src/routes/query-dedup.route.tsx create mode 100644 packages/appkit-ui/src/react/hooks/analytics-request-store.ts diff --git a/apps/dev-playground/client/src/lib/nav.ts b/apps/dev-playground/client/src/lib/nav.ts index 00f70dfec..e3a32a611 100644 --- a/apps/dev-playground/client/src/lib/nav.ts +++ b/apps/dev-playground/client/src/lib/nav.ts @@ -5,6 +5,7 @@ import { FileCode2Icon, FolderIcon, GaugeIcon, + LayersIcon, LayoutDashboardIcon, LineChartIcon, type LucideIcon, @@ -78,6 +79,13 @@ export const NAV_GROUPS: ReadonlyArray = [ "Type-safe parameter builders and query generators for Databricks SQL.", icon: FileCode2Icon, }, + { + to: "/query-dedup", + label: "Query Dedup", + description: + "Many components, one request: identical analytics queries share a single in-flight fetch.", + icon: LayersIcon, + }, ], }, { diff --git a/apps/dev-playground/client/src/routeTree.gen.ts b/apps/dev-playground/client/src/routeTree.gen.ts index 450287592..48c72fca1 100644 --- a/apps/dev-playground/client/src/routeTree.gen.ts +++ b/apps/dev-playground/client/src/routeTree.gen.ts @@ -17,6 +17,7 @@ import { Route as SqlHelpersRouteRouteImport } from './routes/sql-helpers.route' import { Route as SmartDashboardRouteRouteImport } from './routes/smart-dashboard.route' import { Route as ServingRouteRouteImport } from './routes/serving.route' import { Route as ReconnectRouteRouteImport } from './routes/reconnect.route' +import { Route as QueryDedupRouteRouteImport } from './routes/query-dedup.route' import { Route as PolicyMatrixRouteRouteImport } from './routes/policy-matrix.route' import { Route as LakebaseRouteRouteImport } from './routes/lakebase.route' import { Route as JobsRouteRouteImport } from './routes/jobs.route' @@ -69,6 +70,11 @@ const ReconnectRouteRoute = ReconnectRouteRouteImport.update({ path: '/reconnect', getParentRoute: () => rootRouteImport, } as any) +const QueryDedupRouteRoute = QueryDedupRouteRouteImport.update({ + id: '/query-dedup', + path: '/query-dedup', + getParentRoute: () => rootRouteImport, +} as any) const PolicyMatrixRouteRoute = PolicyMatrixRouteRouteImport.update({ id: '/policy-matrix', path: '/policy-matrix', @@ -137,6 +143,7 @@ export interface FileRoutesByFullPath { '/jobs': typeof JobsRouteRoute '/lakebase': typeof LakebaseRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute + '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute '/serving': typeof ServingRouteRoute '/smart-dashboard': typeof SmartDashboardRouteRoute @@ -158,6 +165,7 @@ export interface FileRoutesByTo { '/jobs': typeof JobsRouteRoute '/lakebase': typeof LakebaseRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute + '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute '/serving': typeof ServingRouteRoute '/smart-dashboard': typeof SmartDashboardRouteRoute @@ -180,6 +188,7 @@ export interface FileRoutesById { '/jobs': typeof JobsRouteRoute '/lakebase': typeof LakebaseRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute + '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute '/serving': typeof ServingRouteRoute '/smart-dashboard': typeof SmartDashboardRouteRoute @@ -203,6 +212,7 @@ export interface FileRouteTypes { | '/jobs' | '/lakebase' | '/policy-matrix' + | '/query-dedup' | '/reconnect' | '/serving' | '/smart-dashboard' @@ -224,6 +234,7 @@ export interface FileRouteTypes { | '/jobs' | '/lakebase' | '/policy-matrix' + | '/query-dedup' | '/reconnect' | '/serving' | '/smart-dashboard' @@ -245,6 +256,7 @@ export interface FileRouteTypes { | '/jobs' | '/lakebase' | '/policy-matrix' + | '/query-dedup' | '/reconnect' | '/serving' | '/smart-dashboard' @@ -267,6 +279,7 @@ export interface RootRouteChildren { JobsRouteRoute: typeof JobsRouteRoute LakebaseRouteRoute: typeof LakebaseRouteRoute PolicyMatrixRouteRoute: typeof PolicyMatrixRouteRoute + QueryDedupRouteRoute: typeof QueryDedupRouteRoute ReconnectRouteRoute: typeof ReconnectRouteRoute ServingRouteRoute: typeof ServingRouteRoute SmartDashboardRouteRoute: typeof SmartDashboardRouteRoute @@ -335,6 +348,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ReconnectRouteRouteImport parentRoute: typeof rootRouteImport } + '/query-dedup': { + id: '/query-dedup' + path: '/query-dedup' + fullPath: '/query-dedup' + preLoaderRoute: typeof QueryDedupRouteRouteImport + parentRoute: typeof rootRouteImport + } '/policy-matrix': { id: '/policy-matrix' path: '/policy-matrix' @@ -427,6 +447,7 @@ const rootRouteChildren: RootRouteChildren = { JobsRouteRoute: JobsRouteRoute, LakebaseRouteRoute: LakebaseRouteRoute, PolicyMatrixRouteRoute: PolicyMatrixRouteRoute, + QueryDedupRouteRoute: QueryDedupRouteRoute, ReconnectRouteRoute: ReconnectRouteRoute, ServingRouteRoute: ServingRouteRoute, SmartDashboardRouteRoute: SmartDashboardRouteRoute, diff --git a/apps/dev-playground/client/src/routes/query-dedup.route.tsx b/apps/dev-playground/client/src/routes/query-dedup.route.tsx new file mode 100644 index 000000000..979188711 --- /dev/null +++ b/apps/dev-playground/client/src/routes/query-dedup.route.tsx @@ -0,0 +1,209 @@ +import { + Badge, + Button, + Card, + CardContent, + CardHeader, + CardTitle, + useAnalyticsQuery, +} from "@databricks/appkit-ui/react"; +import { createFileRoute, retainSearchParams } from "@tanstack/react-router"; +import { useEffect, useState, useSyncExternalStore } from "react"; +import { Header } from "@/components/layout/header"; + +export const Route = createFileRoute("/query-dedup")({ + component: QueryDedupRoute, + search: { + middlewares: [retainSearchParams(true)], + }, +}); + +// Two zero-parameter queries. Panels on the same key share one request; the +// key toggle demonstrates that a *different* key opens its own request. +const QUERY_KEYS = ["apps_list", "example"] as const; +type DemoQueryKey = (typeof QUERY_KEYS)[number]; +const ANALYTICS_PATH = "/api/analytics/query/"; + +/** + * Route-local counter for analytics network requests. Wraps `window.fetch` + * while this route is mounted (restored on unmount) and counts POSTs to the + * analytics query endpoint. This is what makes dedup observable in-page + * instead of only in the DevTools Network tab — it counts the real transport + * calls `useAnalyticsQuery` makes, without instrumenting the hook itself. + */ +const requestCounter = (() => { + let count = 0; + const listeners = new Set<() => void>(); + const emit = () => { + for (const l of listeners) l(); + }; + return { + increment() { + count += 1; + emit(); + }, + subscribe(listener: () => void) { + listeners.add(listener); + return () => listeners.delete(listener); + }, + get() { + return count; + }, + }; +})(); + +/** Install the fetch wrapper for the lifetime of the route. */ +function useAnalyticsRequestCounter(): number { + useEffect(() => { + const original = window.fetch; + window.fetch = (input, init) => { + const url = + typeof input === "string" + ? input + : input instanceof URL + ? input.toString() + : input.url; + if (url.includes(ANALYTICS_PATH) && init?.method === "POST") { + requestCounter.increment(); + } + return original(input, init); + }; + return () => { + window.fetch = original; + }; + }, []); + + return useSyncExternalStore( + requestCounter.subscribe, + requestCounter.get, + requestCounter.get, + ); +} + +/** + * A single independent consumer of a shared query. Each mounted panel is a + * separate `useAnalyticsQuery` hook instance — without dedup, each would fire + * its own request. + */ +function Panel({ label, queryKey }: { label: string; queryKey: DemoQueryKey }) { + const { data, loading, error } = useAnalyticsQuery(queryKey, {}); + const rows = Array.isArray(data) ? data.length : 0; + + return ( + + + + Panel {label} + {loading ? ( + loading… + ) : error ? ( + error + ) : ( + {rows} rows + )} + + + + useAnalyticsQuery("{queryKey}") + + + ); +} + +const PANEL_LABELS = ["A", "B", "C", "D", "E", "F", "G", "H"]; + +function QueryDedupRoute() { + const requestCount = useAnalyticsRequestCounter(); + const [panelCount, setPanelCount] = useState(4); + // When true, the last panel switches to a different query key, so it can no + // longer share the request — the counter ticks up to prove distinct keys + // still fan out independently. + const [splitLast, setSplitLast] = useState(false); + + const labels = PANEL_LABELS.slice(0, panelCount); + const distinctKeys = splitLast && panelCount > 1 ? 2 : 1; + + return ( +
+
+
+ + + +
+
+ {panelCount} +
+
+ components mounted +
+
+
→
+
+
+ {requestCount} +
+
+ network request{requestCount === 1 ? "" : "s"} fired +
+
+
+ {distinctKeys === 1 ? ( + <> + All {panelCount} panels share one key — without dedup this + would be{" "} + + {panelCount} + {" "} + requests. + + ) : ( + <> + Two distinct keys in use → two requests, no matter how many + panels share each. + + )} +
+
+ + +
+ + + +
+ {labels.map((label, i) => { + const isSplit = splitLast && i === labels.length - 1; + return ( + + ); + })} +
+
+
+ ); +} diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts index 4c5f1dd58..c9558288c 100644 --- a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts @@ -37,8 +37,21 @@ vi.mock("../use-query-hmr", () => ({ useQueryHMR: vi.fn(), })); +import { + getSnapshot, + resetAnalyticsRequestStore, + retain, + start, + subscribe, +} from "../analytics-request-store"; import { useAnalyticsQuery } from "../use-analytics-query"; +const JSON_OPTS = { + url: "/api/analytics/query/q", + payload: JSON.stringify({ parameters: null, format: "JSON_ARRAY" }), + format: "JSON_ARRAY", +}; + function markAborted() { const sig = capturedCallbacks.signal; if (!sig) throw new Error("signal not captured yet"); @@ -50,6 +63,9 @@ describe("useAnalyticsQuery", () => { vi.clearAllMocks(); lastConnectArgs = null; capturedCallbacks = {}; + // The request store is a module singleton; clear it between tests so + // entries (and their `connectSSE` call counts) don't leak across cases. + resetAnalyticsRequestStore(); }); afterEach(() => { @@ -459,4 +475,110 @@ describe("useAnalyticsQuery", () => { expect(result.current.data).toBeNull(); }); }); + + describe("shared in-flight requests (dedup)", () => { + test("two hook instances with the same key share one request", () => { + const { unmount: unmount1 } = renderHook(() => + useAnalyticsQuery("shared" as any, { a: 1 } as any), + ); + const { unmount: unmount2 } = renderHook(() => + useAnalyticsQuery("shared" as any, { a: 1 } as any), + ); + + // Both instances resolve to the same cache key → one network request. + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + unmount1(); + unmount2(); + }); + + test("different params do not share a request", () => { + renderHook(() => useAnalyticsQuery("shared" as any, { a: 1 } as any)); + renderHook(() => useAnalyticsQuery("shared" as any, { a: 2 } as any)); + + expect(mockConnectSSE).toHaveBeenCalledTimes(2); + }); + + test("a late instance sees the in-flight result of an existing request", async () => { + const { result: first } = renderHook(() => + useAnalyticsQuery("shared" as any, { a: 1 } as any), + ); + + // Resolve the shared request via the first instance's SSE stream. + await act(async () => { + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "result", data: [{ id: 7 }] }), + }); + }); + await waitFor(() => expect(first.current.data).toEqual([{ id: 7 }])); + + // A second instance mounting on the same key reads the resolved + // snapshot immediately without opening a new stream. + const { result: second } = renderHook(() => + useAnalyticsQuery("shared" as any, { a: 1 } as any), + ); + + expect(second.current.data).toEqual([{ id: 7 }]); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + }); + }); + + describe("request store lifecycle", () => { + test("retaining the same key twice starts the request once", () => { + const release1 = retain("k", JSON_OPTS); + const release2 = retain("k", JSON_OPTS); + + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + release1(); + release2(); + }); + + test("releasing to zero then re-retaining within a tick reuses the request", () => { + const release = retain("k", JSON_OPTS); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + // Synchronous unmount→remount (StrictMode): teardown is deferred, so the + // re-retain cancels it and keeps the same in-flight request. + release(); + retain("k", JSON_OPTS); + + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + }); + + test("re-retaining after the deferred teardown fires starts a fresh request", async () => { + const release = retain("k", JSON_OPTS); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + + release(); + // Let the deferred teardown run: the entry is dropped. + await new Promise((resolve) => setTimeout(resolve, 0)); + + retain("k", JSON_OPTS); + expect(mockConnectSSE).toHaveBeenCalledTimes(2); + }); + + test("start fans new state out to every subscriber of a key", async () => { + retain("k", JSON_OPTS); + const listener = vi.fn(); + subscribe("k", listener); + + await act(async () => { + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "result", data: [{ id: 1 }] }), + }); + }); + + expect(listener).toHaveBeenCalled(); + expect(getSnapshot("k").data).toEqual([{ id: 1 }]); + }); + + test("autoStart:false does not start the request until start() is called", () => { + retain("k", JSON_OPTS, false); + expect(mockConnectSSE).not.toHaveBeenCalled(); + + start("k"); + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + }); + }); }); diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-warehouse-status.test.tsx b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-warehouse-status.test.tsx index 103904423..fa90e2123 100644 --- a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-warehouse-status.test.tsx +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-warehouse-status.test.tsx @@ -36,6 +36,7 @@ vi.mock("../use-query-hmr", () => ({ })); import { ResourceStatusIndicator } from "../../resource-status-indicator"; +import { resetAnalyticsRequestStore } from "../analytics-request-store"; import { useAnalyticsQuery } from "../use-analytics-query"; import { ResourceStatusProvider, @@ -59,6 +60,11 @@ function queryIndicatorToast(): HTMLElement | null { describe("useAnalyticsQuery + ResourceStatusProvider integration", () => { afterEach(() => { cleanup(); + // `useAnalyticsQuery` is backed by a module-singleton request store; every + // Chart here shares the `chart_one` key, so clear it between tests (after + // unmount) to cancel deferred teardowns and avoid entry reuse leaking a + // captured `onMessage` across cases. + resetAnalyticsRequestStore(); vi.clearAllMocks(); }); diff --git a/packages/appkit-ui/src/react/hooks/analytics-request-store.ts b/packages/appkit-ui/src/react/hooks/analytics-request-store.ts new file mode 100644 index 000000000..988ad8cb3 --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/analytics-request-store.ts @@ -0,0 +1,421 @@ +import { ArrowClient, connectSSE } from "@/js"; +import type { WarehouseStatus } from "./types"; + +/** + * Shared in-flight request store for `useAnalyticsQuery`. + * + * Multiple hook instances that resolve to the same request (same query key, + * parameters, format, and dev mode) share a single network request keyed by a + * cache key. Each keyed {@link Entry} owns one transport (SSE or direct Arrow + * fetch) and fans both the final result and mid-flight `warehouse_status` + * updates out to every subscriber via `useSyncExternalStore`. + * + * Dedup-only: a keyed entry lives exactly as long as it has subscribers. When + * the last one releases, teardown is deferred a tick (so a StrictMode + * unmount→remount reuses the in-flight request instead of aborting it); if no + * one has re-subscribed by then, the request is aborted and the entry dropped. + * There is no cross-lifecycle result cache. + */ + +const GENERIC_LOAD_ERROR = "Unable to load data, please try again"; + +/** Map a fetch/SSE transport error to a user-facing message. */ +function userFacingFetchError(error: unknown): string { + if (error instanceof Error) { + if (error.name === "AbortError") { + return "Request timed out, please try again"; + } + if (error.message.includes("Failed to fetch")) { + return "Network error. Please check your connection."; + } + } + return GENERIC_LOAD_ERROR; +} + +function isWarehouseStatusPayload(value: unknown): value is WarehouseStatus { + return ( + typeof value === "object" && + value !== null && + typeof (value as WarehouseStatus).state === "string" + ); +} + +/** Options describing the request a keyed entry runs. */ +interface AnalyticsRequestOptions { + /** Full request URL (already includes the encoded query key and dev suffix). */ + url: string; + /** Serialized `{ parameters, format }` body. */ + payload: string; + /** Response format; selects the transport. */ + format: string; +} + +/** Immutable per-key request state; mirrors the hook's public result shape. */ +interface AnalyticsRequestSnapshot { + data: unknown; + loading: boolean; + error: string | null; + errorCode: string | null; + warehouseStatus: WarehouseStatus | null; +} + +/** Idle snapshot returned for keys with no live entry. Referentially stable. */ +export const EMPTY_SNAPSHOT: AnalyticsRequestSnapshot = { + data: null, + loading: false, + error: null, + errorCode: null, + warehouseStatus: null, +}; + +/** Snapshot a request resets to when it (re)starts. */ +const LOADING_SNAPSHOT: AnalyticsRequestSnapshot = { + data: null, + loading: true, + error: null, + errorCode: null, + warehouseStatus: null, +}; + +interface Entry { + snapshot: AnalyticsRequestSnapshot; + refCount: number; + abortController: AbortController | null; + teardownTimer: ReturnType | null; + /** True once `start` has run at least once; guards re-run on late `retain`. */ + started: boolean; + options: AnalyticsRequestOptions; +} + +const entries = new Map(); + +// Listeners are keyed independently of `entries` so a subscriber registered +// before its entry exists (React can call `useSyncExternalStore`'s subscribe +// before the `retain` effect runs) still receives notifications once the +// request starts. +const listenersByKey = new Map void>>(); + +function notify(key: string): void { + const listeners = listenersByKey.get(key); + if (!listeners) return; + for (const listener of listeners) listener(); +} + +/** Replace an entry's snapshot immutably and notify subscribers. */ +function patch( + key: string, + entry: Entry, + next: Partial, +): void { + entry.snapshot = { ...entry.snapshot, ...next }; + notify(key); +} + +async function handleSseMessage( + key: string, + entry: Entry, + parsed: Record, +): Promise { + if (parsed.type === "warehouse_status") { + if (!isWarehouseStatusPayload(parsed.status)) { + patch(key, entry, { loading: false, error: GENERIC_LOAD_ERROR }); + console.error( + "[useAnalyticsQuery] Malformed warehouse_status event", + parsed, + ); + return; + } + patch(key, entry, { warehouseStatus: parsed.status }); + return; + } + + // JSON result. The SSE wire schema is intentionally loose (`data` is an + // optional array of unknown values), so a structural check is enough here — + // no need to ship a schema validator (zod, ~60 KB gz) to the browser just + // to read our own same-origin server's messages. Missing or non-array + // `data` normalizes to [] so `undefined` never bleeds into the hook's + // `T | null` state. + if (parsed.type === "result") { + patch(key, entry, { + loading: false, + data: Array.isArray(parsed.data) ? parsed.data : [], + }); + return; + } + + // NOTE: ARROW_STREAM no longer flows over SSE — the server streams the + // raw Arrow IPC bytes back as the query response body, handled by + // `fetchArrowDirect` instead of this SSE handler. + + if (parsed.type === "error" || parsed.error || parsed.code) { + const errorMsg = + (parsed.error as string | undefined) || + (parsed.message as string | undefined) || + "Unable to execute query"; + // Propagate the upstream structured code so UI consumers can branch on + // a stable identifier (e.g. format-switch on + // RESULT_TOO_LARGE_FOR_JSON_FALLBACK or ARROW_DELIVERY_UNSUPPORTED) + // instead of parsing the human-readable message. + patch(key, entry, { + loading: false, + error: errorMsg, + ...(typeof parsed.errorCode === "string" + ? { errorCode: parsed.errorCode } + : {}), + }); + if (parsed.code) { + console.error( + `[useAnalyticsQuery] Code: ${parsed.code}, Message: ${errorMsg}`, + ); + } + return; + } + + // Not a warehouse-status, result, or error event — surface a generic error + // rather than silently dropping an unrecognized payload. + console.error("[useAnalyticsQuery] Unrecognized SSE payload", parsed); + patch(key, entry, { loading: false, error: GENERIC_LOAD_ERROR }); +} + +/** + * Fetch the real column names for a statement from the fallback endpoint, + * used when a very wide schema's names didn't fit in the response header. + * Returns undefined on any failure so decoding falls back to the raw Arrow + * schema names. + */ +async function fetchArrowColumns( + statementId: string, + signal: AbortSignal, +): Promise { + try { + const res = await fetch( + `/api/analytics/columns/${encodeURIComponent(statementId)}`, + { signal }, + ); + if (!res.ok) return undefined; + const body = (await res.json()) as { columns?: unknown }; + return Array.isArray(body.columns) ? (body.columns as string[]) : undefined; + } catch { + return undefined; + } +} + +/** + * Fetch an ARROW_STREAM query result as raw Arrow IPC bytes directly from + * the query endpoint (no SSE, no second /arrow-result request) and decode + * it into a Table. The server streams the bytes back as the POST response + * body; errors before the first byte arrive as a JSON `{ error, errorCode }`. + */ +async function fetchArrowDirect( + key: string, + entry: Entry, + signal: AbortSignal, +): Promise { + try { + const response = await fetch(entry.options.url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: entry.options.payload, + signal, + }); + if (signal.aborted) return; + + if (!response.ok) { + let message = GENERIC_LOAD_ERROR; + let code: string | null = null; + try { + const body = (await response.json()) as { + error?: string; + errorCode?: string; + }; + if (body.error) message = body.error; + if (typeof body.errorCode === "string") code = body.errorCode; + } catch { + // Non-JSON error body — keep the generic message. + } + patch(key, entry, { loading: false, error: message, errorCode: code }); + return; + } + + const buffer = await response.arrayBuffer(); + if (signal.aborted) return; + // Databricks encodes ARROW_STREAM columns positionally (col_0, …); the + // server sends the real manifest names so we can relabel the decoded + // Table (charts look columns up by name). Normally inline in the + // `X-Appkit-Arrow-Columns` header; for very wide schemas the header + // carries only a statement-id reference and we fetch the names. + let columnNames: string[] | undefined; + const header = response.headers.get("X-Appkit-Arrow-Columns"); + if (header) { + try { + columnNames = JSON.parse(decodeURIComponent(header)); + } catch { + // Malformed header — fall back to the raw Arrow schema names. + } + } else { + const ref = response.headers.get("X-Appkit-Arrow-Columns-Ref"); + if (ref) { + columnNames = await fetchArrowColumns(ref, signal); + } + } + const table = await ArrowClient.processArrowBuffer( + new Uint8Array(buffer), + columnNames, + ); + patch(key, entry, { loading: false, data: table }); + } catch (error) { + if (signal.aborted) return; + patch(key, entry, { loading: false, error: userFacingFetchError(error) }); + } +} + +/** + * (Re)start the request for a keyed entry: abort any in-flight transport, + * reset the snapshot to loading, and run the format-appropriate transport. + * The new state fans out to every current subscriber. + */ +export function start(key: string): void { + const entry = entries.get(key); + if (!entry) return; + + entry.abortController?.abort(); + + entry.started = true; + entry.snapshot = LOADING_SNAPSHOT; + notify(key); + + const abortController = new AbortController(); + entry.abortController = abortController; + const { signal } = abortController; + + // ARROW_STREAM: the server streams raw Arrow IPC bytes back on the query + // response body (no SSE). Fetch and decode directly. + if (entry.options.format === "ARROW_STREAM") { + void fetchArrowDirect(key, entry, signal); + return; + } + + connectSSE({ + url: entry.options.url, + payload: entry.options.payload, + signal, + onMessage: async (message) => { + // Drop late envelopes from a stream whose controller was already + // aborted (React StrictMode unmount→remount). Mirrors onError below. + if (signal.aborted) return; + try { + const parsed = JSON.parse(message.data) as Record; + await handleSseMessage(key, entry, parsed); + } catch (error) { + // A `JSON.parse` failure (or any other thrown error inside the + // SSE message handler) used to leave the hook permanently in + // `loading=true` with no error surfaced — the UI would just + // spin forever. Clear loading and report a user-facing error + // so the consumer can render a retry affordance. + // + // We also abort the SSE connection: if the upstream is + // emitting un-parseable frames, leaving the stream open just + // re-fires the same failure on the next message. Closing + // forces the consumer into a clean retry path. + console.warn("[useAnalyticsQuery] Malformed message received", error); + patch(key, entry, { loading: false, error: GENERIC_LOAD_ERROR }); + abortController.abort(); + } + }, + onError: (error) => { + if (signal.aborted) return; + + if (error instanceof Error) { + console.error("[useAnalyticsQuery] Error", { + url: entry.options.url, + error: error.message, + stack: error.stack, + }); + } + patch(key, entry, { loading: false, error: userFacingFetchError(error) }); + }, + }); +} + +/** + * Register a subscriber for `key`, creating and starting the shared request + * on first use. Returns a `release` function that must be called on unmount. + * + * @param key Cache key uniquely identifying the request. + * @param options Request options; only used when the entry is first created. + * @param autoStart Whether to start the request on creation. Default true. + */ +export function retain( + key: string, + options: AnalyticsRequestOptions, + autoStart = true, +): () => void { + let entry = entries.get(key); + if (!entry) { + entry = { + snapshot: EMPTY_SNAPSHOT, + refCount: 0, + abortController: null, + teardownTimer: null, + started: false, + options, + }; + entries.set(key, entry); + } + + // A late joiner cancels any pending teardown so it keeps the live request. + if (entry.teardownTimer !== null) { + clearTimeout(entry.teardownTimer); + entry.teardownTimer = null; + } + entry.refCount += 1; + + if (autoStart && !entry.started) { + start(key); + } + + return () => release(key); +} + +function release(key: string): void { + const entry = entries.get(key); + if (!entry) return; + entry.refCount -= 1; + if (entry.refCount > 0) return; + + // Defer teardown one tick: a StrictMode unmount→remount (or a fast + // route swap) re-`retain`s within the same tick and reuses the request. + entry.teardownTimer = setTimeout(() => { + const current = entries.get(key); + if (!current || current.refCount > 0) return; + current.abortController?.abort(); + entries.delete(key); + }, 0); +} + +export function subscribe(key: string, listener: () => void): () => void { + let listeners = listenersByKey.get(key); + if (!listeners) { + listeners = new Set(); + listenersByKey.set(key, listeners); + } + listeners.add(listener); + return () => { + listeners.delete(listener); + if (listeners.size === 0) listenersByKey.delete(key); + }; +} + +export function getSnapshot(key: string): AnalyticsRequestSnapshot { + return entries.get(key)?.snapshot ?? EMPTY_SNAPSHOT; +} + +/** Test-only: abort every in-flight request and clear the store. */ +export function resetAnalyticsRequestStore(): void { + for (const entry of entries.values()) { + if (entry.teardownTimer !== null) clearTimeout(entry.teardownTimer); + entry.abortController?.abort(); + } + entries.clear(); + listenersByKey.clear(); +} diff --git a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts index 1c63ffe14..bc1df83f0 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -4,9 +4,9 @@ import { useId, useMemo, useRef, - useState, + useSyncExternalStore, } from "react"; -import { ArrowClient, connectSSE } from "@/js"; +import * as store from "./analytics-request-store"; import type { AnalyticsFormat, InferParams, @@ -14,7 +14,6 @@ import type { QueryKey, UseAnalyticsQueryOptions, UseAnalyticsQueryResult, - WarehouseStatus, } from "./types"; import { useAnalyticsWarehousePublisher } from "./use-analytics-warehouse-status"; import { useQueryHMR } from "./use-query-hmr"; @@ -61,212 +60,7 @@ function getDevMode(): string { return dev ? `?dev=${dev}` : ""; } -const GENERIC_LOAD_ERROR = "Unable to load data, please try again"; - -/** Map a fetch/SSE transport error to a user-facing message. */ -function userFacingFetchError(error: unknown): string { - if (error instanceof Error) { - if (error.name === "AbortError") { - return "Request timed out, please try again"; - } - if (error.message.includes("Failed to fetch")) { - return "Network error. Please check your connection."; - } - } - return GENERIC_LOAD_ERROR; -} - -interface AnalyticsQuerySseContext { - setLoading: (loading: boolean) => void; - setError: (error: string | null) => void; - setErrorCode: (code: string | null) => void; - setData: (data: ResultType | null) => void; - setWarehouseStatus: (status: WarehouseStatus | null) => void; - publishWarehouseStatus: (status: WarehouseStatus | null) => void; - unpublishWarehouseStatus: () => void; -} - -function isWarehouseStatusPayload(value: unknown): value is WarehouseStatus { - return ( - typeof value === "object" && - value !== null && - typeof (value as WarehouseStatus).state === "string" - ); -} - -async function handleAnalyticsSseMessage( - parsed: Record, - ctx: AnalyticsQuerySseContext, -): Promise { - if (parsed.type === "warehouse_status") { - if (!isWarehouseStatusPayload(parsed.status)) { - ctx.setLoading(false); - ctx.setError(GENERIC_LOAD_ERROR); - ctx.unpublishWarehouseStatus(); - console.error( - "[useAnalyticsQuery] Malformed warehouse_status event", - parsed, - ); - return; - } - ctx.setWarehouseStatus(parsed.status); - ctx.publishWarehouseStatus(parsed.status); - return; - } - - // JSON result. The SSE wire schema is intentionally loose (`data` is an - // optional array of unknown values), so a structural check is enough here — - // no need to ship a schema validator (zod, ~60 KB gz) to the browser just - // to read our own same-origin server's messages. Missing or non-array - // `data` normalizes to [] so `undefined` never bleeds into the hook's - // `T | null` state. - if (parsed.type === "result") { - ctx.setLoading(false); - ctx.setData((Array.isArray(parsed.data) ? parsed.data : []) as ResultType); - ctx.unpublishWarehouseStatus(); - return; - } - - // NOTE: ARROW_STREAM no longer flows over SSE — the server streams the - // raw Arrow IPC bytes back as the query response body, handled by - // `fetchArrowDirect` instead of this SSE handler. - - if (parsed.type === "error" || parsed.error || parsed.code) { - const errorMsg = - (parsed.error as string | undefined) || - (parsed.message as string | undefined) || - "Unable to execute query"; - ctx.setLoading(false); - ctx.setError(errorMsg); - ctx.unpublishWarehouseStatus(); - // Propagate the upstream structured code so UI consumers can branch on - // a stable identifier (e.g. format-switch on - // RESULT_TOO_LARGE_FOR_JSON_FALLBACK or ARROW_DELIVERY_UNSUPPORTED) - // instead of parsing the human-readable message. - if (typeof parsed.errorCode === "string") { - ctx.setErrorCode(parsed.errorCode); - } - if (parsed.code) { - console.error( - `[useAnalyticsQuery] Code: ${parsed.code}, Message: ${errorMsg}`, - ); - } - return; - } - - // Not a warehouse-status, result, or error event — surface a generic error - // rather than silently dropping an unrecognized payload. - console.error("[useAnalyticsQuery] Unrecognized SSE payload", parsed); - ctx.setLoading(false); - ctx.setError(GENERIC_LOAD_ERROR); - ctx.unpublishWarehouseStatus(); -} - -interface ArrowDirectContext { - url: string; - payload: string; - signal: AbortSignal; - setLoading: (loading: boolean) => void; - setError: (error: string | null) => void; - setErrorCode: (code: string | null) => void; - setData: (data: unknown) => void; - unpublishWarehouseStatus: () => void; -} - -/** - * Fetch the real column names for a statement from the fallback endpoint, - * used when a very wide schema's names didn't fit in the response header. - * Returns undefined on any failure so decoding falls back to the raw Arrow - * schema names. - */ -async function fetchArrowColumns( - statementId: string, - signal: AbortSignal, -): Promise { - try { - const res = await fetch( - `/api/analytics/columns/${encodeURIComponent(statementId)}`, - { signal }, - ); - if (!res.ok) return undefined; - const body = (await res.json()) as { columns?: unknown }; - return Array.isArray(body.columns) ? (body.columns as string[]) : undefined; - } catch { - return undefined; - } -} - -/** - * Fetch an ARROW_STREAM query result as raw Arrow IPC bytes directly from - * the query endpoint (no SSE, no second /arrow-result request) and decode - * it into a Table. The server streams the bytes back as the POST response - * body; errors before the first byte arrive as a JSON `{ error, errorCode }`. - */ -async function fetchArrowDirect(ctx: ArrowDirectContext): Promise { - try { - const response = await fetch(ctx.url, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: ctx.payload, - signal: ctx.signal, - }); - if (ctx.signal.aborted) return; - - if (!response.ok) { - let message = GENERIC_LOAD_ERROR; - let code: string | null = null; - try { - const body = (await response.json()) as { - error?: string; - errorCode?: string; - }; - if (body.error) message = body.error; - if (typeof body.errorCode === "string") code = body.errorCode; - } catch { - // Non-JSON error body — keep the generic message. - } - ctx.setLoading(false); - ctx.setError(message); - if (code) ctx.setErrorCode(code); - ctx.unpublishWarehouseStatus(); - return; - } - - const buffer = await response.arrayBuffer(); - if (ctx.signal.aborted) return; - // Databricks encodes ARROW_STREAM columns positionally (col_0, …); the - // server sends the real manifest names so we can relabel the decoded - // Table (charts look columns up by name). Normally inline in the - // `X-Appkit-Arrow-Columns` header; for very wide schemas the header - // carries only a statement-id reference and we fetch the names. - let columnNames: string[] | undefined; - const header = response.headers.get("X-Appkit-Arrow-Columns"); - if (header) { - try { - columnNames = JSON.parse(decodeURIComponent(header)); - } catch { - // Malformed header — fall back to the raw Arrow schema names. - } - } else { - const ref = response.headers.get("X-Appkit-Arrow-Columns-Ref"); - if (ref) { - columnNames = await fetchArrowColumns(ref, ctx.signal); - } - } - const table = await ArrowClient.processArrowBuffer( - new Uint8Array(buffer), - columnNames, - ); - ctx.setData(table); - ctx.setLoading(false); - ctx.unpublishWarehouseStatus(); - } catch (error) { - if (ctx.signal.aborted) return; - ctx.setLoading(false); - ctx.unpublishWarehouseStatus(); - ctx.setError(userFacingFetchError(error)); - } -} +const NOOP_SUBSCRIBE: (listener: () => void) => () => void = () => () => {}; /** * Subscribe to an analytics query and return its latest result. JSON_ARRAY @@ -274,6 +68,12 @@ async function fetchArrowDirect(ctx: ArrowDirectContext): Promise { * results are fetched as raw Arrow bytes directly from the query endpoint. * Integration hook between client and analytics plugin. * + * Identical requests (same query key, parameters, format, and dev mode) share + * a single in-flight network request: the first mounting instance starts it, + * later instances subscribe to the same {@link store} entry and see the same + * result and warehouse-status updates. The request is torn down once its last + * subscriber unmounts. + * * The return type is automatically inferred based on the format: * - `format: "JSON_ARRAY"` (default): Returns typed array from QueryRegistry * - `format: "ARROW_STREAM"`: Returns TypedArrowTable with row type preserved @@ -316,13 +116,6 @@ export function useAnalyticsQuery< const urlSuffix = `/api/analytics/query/${encodeURIComponent(queryKey)}${devMode}`; type ResultType = InferResultByFormat; - const [data, setData] = useState(null); - const [loading, setLoading] = useState(false); - const [error, setError] = useState(null); - const [errorCode, setErrorCode] = useState(null); - const [warehouseStatus, setWarehouseStatus] = - useState(null); - const abortControllerRef = useRef(null); const publisherId = useId(); const { @@ -358,114 +151,69 @@ export function useAnalyticsQuery< } }, [stableParameters, format, maxParametersSize]); - const start = useCallback(() => { - if (payload === null) { - setError("Failed to serialize query parameters"); - return; - } + // Cache key shared across hook instances. `payload` already serializes + // `{ parameters, format }`, so identical requests collapse to one key. + const cacheKey = payload === null ? null : `${urlSuffix}::${payload}`; - abortControllerRef.current?.abort(); + const subscribe = useCallback( + (listener: () => void) => + cacheKey === null + ? NOOP_SUBSCRIBE(listener) + : store.subscribe(cacheKey, listener), + [cacheKey], + ); + const getSnapshot = useCallback( + () => + cacheKey === null ? store.EMPTY_SNAPSHOT : store.getSnapshot(cacheKey), + [cacheKey], + ); + const snapshot = useSyncExternalStore(subscribe, getSnapshot, getSnapshot); - setLoading(true); - setError(null); - setErrorCode(null); - setData(null); - setWarehouseStatus(null); - publishWarehouseStatus(null); + const start = useCallback(() => { + if (cacheKey !== null) store.start(cacheKey); + }, [cacheKey]); - const abortController = new AbortController(); - abortControllerRef.current = abortController; + // Register with the shared store on mount / key change; release on cleanup. + // The store starts the request on first retain of a key and reuses the + // in-flight request for later subscribers. + useEffect(() => { + if (cacheKey === null || payload === null) return; + return store.retain( + cacheKey, + { url: urlSuffix, payload, format }, + autoStart, + ); + }, [cacheKey, urlSuffix, payload, format, autoStart]); - // ARROW_STREAM: the server streams raw Arrow IPC bytes back on the query - // response body (no SSE). Fetch and decode directly. - if (format === "ARROW_STREAM") { - void fetchArrowDirect({ - url: urlSuffix, - payload, - signal: abortController.signal, - setLoading, - setError, - setErrorCode, - setData: (table) => setData(table as ResultType), - unpublishWarehouseStatus, - }); - return; + // Mirror this instance's warehouse status into the nearest resource-status + // provider while the request is in flight; clear the slot once it settles. + useEffect(() => { + if (snapshot.loading) { + publishWarehouseStatus(snapshot.warehouseStatus); + } else { + unpublishWarehouseStatus(); } - - const sseContext: AnalyticsQuerySseContext = { - setLoading, - setError, - setErrorCode, - setData, - setWarehouseStatus, - publishWarehouseStatus, - unpublishWarehouseStatus, - }; - - connectSSE({ - url: urlSuffix, - payload, - signal: abortController.signal, - onMessage: async (message) => { - // Drop late envelopes from a stream whose controller was already - // aborted (React StrictMode unmount→remount). Mirrors onError below. - if (abortController.signal.aborted) return; - try { - const parsed = JSON.parse(message.data) as Record; - await handleAnalyticsSseMessage(parsed, sseContext); - } catch (error) { - // A `JSON.parse` failure (or any other thrown error inside the - // SSE message handler) used to leave the hook permanently in - // `loading=true` with no error surfaced — the UI would just - // spin forever. Clear loading and report a user-facing error - // so the consumer can render a retry affordance. - // - // We also abort the SSE connection: if the upstream is - // emitting un-parseable frames, leaving the stream open just - // re-fires the same failure on the next message. Closing - // forces the consumer into a clean retry path. - console.warn("[useAnalyticsQuery] Malformed message received", error); - setLoading(false); - setError(GENERIC_LOAD_ERROR); - abortController.abort(); - } - }, - onError: (error) => { - if (abortController.signal.aborted) return; - setLoading(false); - unpublishWarehouseStatus(); - - if (error instanceof Error) { - console.error("[useAnalyticsQuery] Error", { - queryKey, - error: error.message, - stack: error.stack, - }); - } - setError(userFacingFetchError(error)); - }, - }); }, [ - queryKey, - payload, - urlSuffix, - format, + snapshot.loading, + snapshot.warehouseStatus, publishWarehouseStatus, unpublishWarehouseStatus, ]); - useEffect(() => { - if (autoStart) { - start(); - } - - return () => { - abortControllerRef.current?.abort(); - unpublishWarehouseStatus(); - }; - }, [start, autoStart, unpublishWarehouseStatus]); + useEffect(() => unpublishWarehouseStatus, [unpublishWarehouseStatus]); useQueryHMR(queryKey, start); - return { data, loading, error, errorCode, warehouseStatus }; + return { + data: snapshot.data as ResultType | null, + loading: snapshot.loading, + // A serialization failure never creates a store entry; surface the same + // error the pre-shared-request implementation did. + error: + cacheKey === null + ? "Failed to serialize query parameters" + : snapshot.error, + errorCode: snapshot.errorCode, + warehouseStatus: snapshot.warehouseStatus, + }; }