diff --git a/apps/dev-playground/client/src/lib/nav.ts b/apps/dev-playground/client/src/lib/nav.ts index feed19476..dacb87b44 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, @@ -86,6 +87,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 a57845549..5d9e2009f 100644 --- a/apps/dev-playground/client/src/routeTree.gen.ts +++ b/apps/dev-playground/client/src/routeTree.gen.ts @@ -16,6 +16,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 MetricViewsRouteRouteImport } from './routes/metric-views.route' import { Route as LakebaseRouteRouteImport } from './routes/lakebase.route' @@ -65,6 +66,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', @@ -145,6 +151,7 @@ export interface FileRoutesByFullPath { '/lakebase': typeof LakebaseRouteRoute '/metric-views': typeof MetricViewsRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute + '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute '/serving': typeof ServingRouteRoute '/smart-dashboard': typeof SmartDashboardRouteRoute @@ -167,6 +174,7 @@ export interface FileRoutesByTo { '/lakebase': typeof LakebaseRouteRoute '/metric-views': typeof MetricViewsRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute + '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute '/serving': typeof ServingRouteRoute '/smart-dashboard': typeof SmartDashboardRouteRoute @@ -190,6 +198,7 @@ export interface FileRoutesById { '/lakebase': typeof LakebaseRouteRoute '/metric-views': typeof MetricViewsRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute + '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute '/serving': typeof ServingRouteRoute '/smart-dashboard': typeof SmartDashboardRouteRoute @@ -214,6 +223,7 @@ export interface FileRouteTypes { | '/lakebase' | '/metric-views' | '/policy-matrix' + | '/query-dedup' | '/reconnect' | '/serving' | '/smart-dashboard' @@ -236,6 +246,7 @@ export interface FileRouteTypes { | '/lakebase' | '/metric-views' | '/policy-matrix' + | '/query-dedup' | '/reconnect' | '/serving' | '/smart-dashboard' @@ -258,6 +269,7 @@ export interface FileRouteTypes { | '/lakebase' | '/metric-views' | '/policy-matrix' + | '/query-dedup' | '/reconnect' | '/serving' | '/smart-dashboard' @@ -281,6 +293,7 @@ export interface RootRouteChildren { LakebaseRouteRoute: typeof LakebaseRouteRoute MetricViewsRouteRoute: typeof MetricViewsRouteRoute PolicyMatrixRouteRoute: typeof PolicyMatrixRouteRoute + QueryDedupRouteRoute: typeof QueryDedupRouteRoute ReconnectRouteRoute: typeof ReconnectRouteRoute ServingRouteRoute: typeof ServingRouteRoute SmartDashboardRouteRoute: typeof SmartDashboardRouteRoute @@ -341,6 +354,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' @@ -449,6 +469,7 @@ const rootRouteChildren: RootRouteChildren = { LakebaseRouteRoute: LakebaseRouteRoute, MetricViewsRouteRoute: MetricViewsRouteRoute, 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..7a8a01b75 --- /dev/null +++ b/apps/dev-playground/client/src/routes/query-dedup.route.tsx @@ -0,0 +1,185 @@ +import { + Badge, + Button, + Card, + CardContent, + CardHeader, + CardTitle, + useAnalyticsQuery, +} from "@databricks/appkit-ui/react"; +import { createFileRoute, retainSearchParams } from "@tanstack/react-router"; +import { useEffect, useState } 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/"; + +/** + * Count analytics network requests by wrapping `window.fetch` for the lifetime + * of the route (restored on unmount), tallying 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. + */ +function useAnalyticsRequestCounter(): number { + const [count, setCount] = useState(0); + + 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") { + setCount((c) => c + 1); + } + return original(input, init); + }; + return () => { + window.fetch = original; + }; + }, []); + + return count; +} + +/** + * 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/apps/dev-playground/tests/arrow-analytics.spec.ts b/apps/dev-playground/tests/arrow-analytics.spec.ts index 408c7517e..0dd408782 100644 --- a/apps/dev-playground/tests/arrow-analytics.spec.ts +++ b/apps/dev-playground/tests/arrow-analytics.spec.ts @@ -1,6 +1,5 @@ import { expect, test } from "@playwright/test"; import { - STRICT_MODE_MULTIPLIER, setupMockAPI, trackApiCalls, waitForChartsToLoad, @@ -27,10 +26,13 @@ test.describe("Arrow Analytics", () => { await page.evaluate(() => window.scrollTo(0, document.body.scrollHeight)); await waitForChartsToLoad(page); - expect(appsListCalls.length).toBe(5 * STRICT_MODE_MULTIPLIER); - expect(spendDataCalls.length).toBe(5 * STRICT_MODE_MULTIPLIER); - expect(topContributorsCalls.length).toBe(2 * STRICT_MODE_MULTIPLIER); - expect(heatmapCalls.length).toBe(2 * STRICT_MODE_MULTIPLIER); + // Deduplicated by (queryKey, parameters, format): each query is rendered + // in both a JSON and an Arrow chart, so it settles at one request per + // format = 2. + expect(appsListCalls.length).toBe(2); + expect(spendDataCalls.length).toBe(2); + expect(topContributorsCalls.length).toBe(2); + expect(heatmapCalls.length).toBe(2); }); test("charts render with mock data (no empty states)", async ({ page }) => { diff --git a/apps/dev-playground/tests/data-visualization.spec.ts b/apps/dev-playground/tests/data-visualization.spec.ts index 55d1f4b60..1bb6269d1 100644 --- a/apps/dev-playground/tests/data-visualization.spec.ts +++ b/apps/dev-playground/tests/data-visualization.spec.ts @@ -1,9 +1,5 @@ import { expect, test } from "@playwright/test"; -import { - STRICT_MODE_MULTIPLIER, - setupMockAPI, - trackApiCalls, -} from "./utils/test-utils"; +import { setupMockAPI, trackApiCalls } from "./utils/test-utils"; test.describe("Data Visualization Route Tests", () => { test.beforeEach(async ({ page }) => { @@ -64,9 +60,11 @@ test.describe("Data Visualization Route Tests", () => { await page.evaluate(() => window.scrollTo(0, document.body.scrollHeight)); await page.waitForLoadState("networkidle"); - expect(untaggedAppsCalls.length).toBe(2 * STRICT_MODE_MULTIPLIER); - expect(spendDataCalls.length).toBe(6 * STRICT_MODE_MULTIPLIER); - expect(topContributorsCalls.length).toBe(4 * STRICT_MODE_MULTIPLIER); + // Deduplicated by (queryKey, parameters, format): every chart here uses the + // same params per key, so all charts of a key collapse to one request. + expect(untaggedAppsCalls.length).toBe(1); + expect(spendDataCalls.length).toBe(1); + expect(topContributorsCalls.length).toBe(1); }); test("can toggle code visibility", async ({ page }) => { diff --git a/apps/dev-playground/tests/utils/test-utils.ts b/apps/dev-playground/tests/utils/test-utils.ts index f608a0e07..8091cc0fb 100644 --- a/apps/dev-playground/tests/utils/test-utils.ts +++ b/apps/dev-playground/tests/utils/test-utils.ts @@ -6,15 +6,6 @@ import { mockTelemetryResponse, } from "./mock-data"; -/** - * React 19 Strict Mode doubles useEffect invocations in development mode - * to help detect side effects. This multiplier accounts for that behavior - * when asserting API call counts in tests. - * - * @see https://react.dev/reference/react/StrictMode#fixing-bugs-found-by-re-running-effects-in-development - */ -export const STRICT_MODE_MULTIPLIER = 2; - function createSSEResponse(data: unknown): string { const event = JSON.stringify({ type: "result", data }); return `data: ${event}\n\n`; 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..9332125cc --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/analytics-request-store.ts @@ -0,0 +1,323 @@ +import { ArrowClient, connectSSE } from "@/js"; +import { + type AnalyticsSseHandlerContext, + GENERIC_LOAD_ERROR, + handleAnalyticsSseError, + handleAnalyticsSseMessage, + userFacingFetchError, +} from "./analytics-sse"; +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. + */ + +/** 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. */ +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(); + +// Keyed separately from `entries`: `subscribe` can run before `retain` creates +// the entry, so listeners must survive independently of entry lifetime. +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); +} + +/** + * 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; + } + + // Adapts the shared SSE handler onto the snapshot model. No warehouse + // publisher lives here — the hook mirrors status from the snapshot — so + // `unpublishWarehouseStatus` is a no-op. + const sseContext: AnalyticsSseHandlerContext = { + source: "useAnalyticsQuery", + resource: { url: entry.options.url }, + defaultExecutionError: "Unable to execute query", + unpublishOnMalformedMessage: false, + signal, + abort: () => abortController.abort(), + setLoading: (loading) => patch(key, entry, { loading }), + setError: (error) => patch(key, entry, { error }), + setErrorCode: (errorCode) => patch(key, entry, { errorCode }), + onWarehouseStatus: (status) => + patch(key, entry, { warehouseStatus: status }), + onResult: (message) => patch(key, entry, { data: message.data }), + unpublishWarehouseStatus: () => {}, + }; + + connectSSE({ + url: entry.options.url, + payload: entry.options.payload, + signal, + onMessage: (message) => handleAnalyticsSseMessage(message.data, sseContext), + onError: (error) => handleAnalyticsSseError(error, sseContext), + }); +} + +/** + * 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 93b3dba36..e263b696f 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -4,17 +4,10 @@ import { useId, useMemo, useRef, - useState, + useSyncExternalStore, } from "react"; -import { ArrowClient, connectSSE } from "@/js"; -import { - type AnalyticsSseHandlerContext, - GENERIC_LOAD_ERROR, - getDevMode, - handleAnalyticsSseError, - handleAnalyticsSseMessage, - userFacingFetchError, -} from "./analytics-sse"; +import * as store from "./analytics-request-store"; +import { getDevMode } from "./analytics-sse"; import type { AnalyticsFormat, InferParams, @@ -22,7 +15,6 @@ import type { QueryKey, UseAnalyticsQueryOptions, UseAnalyticsQueryResult, - WarehouseStatus, } from "./types"; import { useAnalyticsWarehousePublisher } from "./use-analytics-warehouse-status"; import { useQueryHMR } from "./use-query-hmr"; @@ -64,118 +56,18 @@ function useStableParams(value: T): T { return ref.current; } -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)); - } -} - /** * Subscribe to an analytics query and return its latest result. JSON_ARRAY * results stream over SSE (with warehouse-readiness progress); ARROW_STREAM * 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 @@ -218,13 +110,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 { @@ -260,87 +145,64 @@ export function useAnalyticsQuery< } }, [stableParameters, format, maxParametersSize]); - const start = useCallback(() => { - if (payload === null) { - setError("Failed to serialize query parameters"); - return; - } - - abortControllerRef.current?.abort(); - - setLoading(true); - setError(null); - setErrorCode(null); - setData(null); - setWarehouseStatus(null); - publishWarehouseStatus(null); - - const abortController = new AbortController(); - abortControllerRef.current = abortController; + // Cache key shared across hook instances. `payload` already serializes + // `{ parameters, format }`, so identical requests collapse to one key. + // On a serialization failure (`payload === null`) the key stays unused: no + // request is retained and the store reports the stable idle snapshot. + const cacheKey = `${urlSuffix}::${payload}`; + + const subscribe = useCallback( + (listener: () => void) => store.subscribe(cacheKey, listener), + [cacheKey], + ); + const getSnapshot = useCallback( + () => store.getSnapshot(cacheKey), + [cacheKey], + ); + const snapshot = useSyncExternalStore(subscribe, getSnapshot, getSnapshot); + + const start = useCallback(() => store.start(cacheKey), [cacheKey]); + + // 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 (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: AnalyticsSseHandlerContext = { - source: "useAnalyticsQuery", - resource: { queryKey }, - defaultExecutionError: "Unable to execute query", - unpublishOnMalformedMessage: false, - signal: abortController.signal, - abort: () => abortController.abort(), - setLoading, - setError, - setErrorCode, - onWarehouseStatus: (status) => { - setWarehouseStatus(status); - publishWarehouseStatus(status); - }, - onResult: (message) => setData(message.data as ResultType), - unpublishWarehouseStatus, - }; - - connectSSE({ - url: urlSuffix, - payload, - signal: abortController.signal, - onMessage: (message) => - handleAnalyticsSseMessage(message.data, sseContext), - onError: (error) => handleAnalyticsSseError(error, sseContext), - }); }, [ - 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, so surface it here. + error: + payload === null + ? "Failed to serialize query parameters" + : snapshot.error, + errorCode: snapshot.errorCode, + warehouseStatus: snapshot.warehouseStatus, + }; }