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..959ca9f2e --- /dev/null +++ b/apps/dev-playground/client/src/routes/query-dedup.route.tsx @@ -0,0 +1,186 @@ +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 20e28209b..c83bdae8f 100644 --- a/apps/dev-playground/tests/arrow-analytics.spec.ts +++ b/apps/dev-playground/tests/arrow-analytics.spec.ts @@ -1,7 +1,6 @@ import { expect, test } from "@playwright/test"; import { - STRICT_MODE_MULTIPLIER, setupMockAPI, trackApiCalls, waitForChartsToLoad, @@ -28,10 +27,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 85373bb0e..949777f9d 100644 --- a/apps/dev-playground/tests/data-visualization.spec.ts +++ b/apps/dev-playground/tests/data-visualization.spec.ts @@ -1,10 +1,6 @@ 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 }) => { @@ -65,9 +61,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 7026a260c..9596bd43a 100644 --- a/apps/dev-playground/tests/utils/test-utils.ts +++ b/apps/dev-playground/tests/utils/test-utils.ts @@ -7,15 +7,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__/request-store.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/request-store.test.ts new file mode 100644 index 000000000..966e07836 --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/__tests__/request-store.test.ts @@ -0,0 +1,91 @@ +import { beforeEach, describe, expect, test, vi } from "vitest"; + +import { createRequestStore, type RequestControls } from "../request-store"; + +interface Snap { + value: number | null; +} +const IDLE: Snap = { value: null }; + +// A store whose `run` just records how often it fired (no real transport), so +// these tests exercise the generic lifecycle in isolation from SSE/Arrow. +function makeStore() { + const store = createRequestStore(IDLE); + const run = vi.fn((_c: RequestControls) => {}); + return { store, run }; +} + +describe("createRequestStore", () => { + let store: ReturnType["store"]; + let run: ReturnType["run"]; + + beforeEach(() => { + ({ store, run } = makeStore()); + }); + + test("two retains on the same key start the request once", () => { + const r1 = store.retain("k", run); + const r2 = store.retain("k", run); + expect(run).toHaveBeenCalledTimes(1); + r1(); + r2(); + }); + + test("distinct keys start separate requests", () => { + store.retain("a", run); + store.retain("b", run); + expect(run).toHaveBeenCalledTimes(2); + }); + + test("re-retaining within a tick after release reuses the request", () => { + const release = store.retain("k", run); + release(); + store.retain("k", run); + expect(run).toHaveBeenCalledTimes(1); + }); + + test("re-retaining after the deferred teardown starts a fresh request", async () => { + const release = store.retain("k", run); + release(); + // Let the deferred teardown run: the entry is dropped. + await new Promise((resolve) => setTimeout(resolve, 0)); + store.retain("k", run); + expect(run).toHaveBeenCalledTimes(2); + }); + + test("patch fans the new snapshot out to every subscriber of a key", () => { + const listener = vi.fn(); + store.subscribe("k", listener); + store.retain("k", (c) => c.patch({ value: 42 })); + + expect(listener).toHaveBeenCalled(); + expect(store.getSnapshot("k").value).toBe(42); + }); + + test("getSnapshot returns the idle snapshot for a key with no entry", () => { + expect(store.getSnapshot("missing")).toBe(IDLE); + }); + + test("autoStart:false defers the run until start() is called", () => { + store.retain("k", run, false); + expect(run).not.toHaveBeenCalled(); + + store.start("k"); + expect(run).toHaveBeenCalledTimes(1); + }); + + test("reset aborts in-flight runs and clears entries", () => { + let captured: AbortSignal | undefined; + store.retain("k", (c) => { + captured = c.signal; + }); + expect(captured?.aborted).toBe(false); + + store.reset(); + + expect(captured?.aborted).toBe(true); + // Entry is gone: a fresh retain starts a new run. + store.retain("k", run); + expect(run).toHaveBeenCalledTimes(1); + }); +}); 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 cf9796d00..e6697eaa5 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(() => { @@ -451,4 +467,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..203bc0aea --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/analytics-request-store.ts @@ -0,0 +1,226 @@ +import { ArrowClient, connectSSE } from "@/js"; + +import { + type AnalyticsSseHandlerContext, + GENERIC_LOAD_ERROR, + handleAnalyticsSseError, + handleAnalyticsSseMessage, + userFacingFetchError, +} from "./analytics-sse"; +import { + createRequestStore, + type RequestControls, + type RequestRunner, +} from "./request-store"; +import type { WarehouseStatus } from "./types"; + +/** + * Shared in-flight request store for `useAnalyticsQuery`: an instance of the + * generic {@link createRequestStore} lifecycle wired to the analytics + * transports. Multiple hook instances resolving to the same request (query + * key, parameters, format, dev mode) share one network request and see the + * same result and mid-flight `warehouse_status` updates. + * + * The lifecycle (dedup, refcount, deferred teardown) lives in the factory; this + * module only supplies the snapshot shape and how a request runs — SSE via + * `analytics-sse.ts` or a direct Arrow fetch. + */ + +/** 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, +}; + +type Controls = RequestControls; + +/** + * 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( + controls: Controls, + options: AnalyticsRequestOptions, +): Promise { + const { signal } = controls; + try { + const response = await fetch(options.url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: 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. + } + controls.patch({ 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, + ); + controls.patch({ loading: false, data: table }); + } catch (error) { + if (signal.aborted) return; + controls.patch({ loading: false, error: userFacingFetchError(error) }); + } +} + +/** + * Build the runner for a request: reset to loading, then run the + * format-appropriate transport, reporting state through `controls.patch`. + */ +function runAnalyticsRequest( + options: AnalyticsRequestOptions, +): RequestRunner { + return (controls) => { + controls.patch(LOADING_SNAPSHOT); + + // 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); + return; + } + + // Adapt 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: options.url }, + defaultExecutionError: "Unable to execute query", + unpublishOnMalformedMessage: false, + signal: controls.signal, + abort: controls.abort, + setLoading: (loading) => controls.patch({ loading }), + setError: (error) => controls.patch({ error }), + setErrorCode: (errorCode) => controls.patch({ errorCode }), + onWarehouseStatus: (status) => + controls.patch({ warehouseStatus: status }), + onResult: (message) => controls.patch({ data: message.data }), + unpublishWarehouseStatus: () => {}, + }; + + connectSSE({ + url: options.url, + payload: options.payload, + signal: controls.signal, + onMessage: (message) => + handleAnalyticsSseMessage(message.data, sseContext), + onError: (error) => handleAnalyticsSseError(error, sseContext), + }); + }; +} + +const store = createRequestStore(EMPTY_SNAPSHOT); + +/** + * Register a subscriber for `key`, starting the shared request on first use. + * Returns a `release` function that must be called on unmount. + */ +export function retain( + key: string, + options: AnalyticsRequestOptions, + autoStart = true, +): () => void { + return store.retain(key, runAnalyticsRequest(options), autoStart); +} + +export const start = store.start; +export const subscribe = store.subscribe; +export const getSnapshot = store.getSnapshot; + +/** Test-only: abort every in-flight request and clear the store. */ +export const resetAnalyticsRequestStore = store.reset; diff --git a/packages/appkit-ui/src/react/hooks/request-store.ts b/packages/appkit-ui/src/react/hooks/request-store.ts new file mode 100644 index 000000000..485a0f7fe --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/request-store.ts @@ -0,0 +1,164 @@ +/** + * Generic keyed request store: coalesces identical in-flight requests so N + * subscribers of the same key share one run, and fans state updates back out + * via `useSyncExternalStore`. + * + * Owns only the lifecycle — refcount, deferred teardown, subscribe/notify — and + * is transport-agnostic: the caller's `run` performs the actual fetch (SSE, + * Arrow, plain fetch, …) and reports state through `controls.patch`. The + * snapshot type `S` and its reset policy live entirely with the caller. + * + * Dedup-only: an entry lives exactly as long as it has subscribers. When the + * last one releases, teardown is deferred a tick (so a React StrictMode + * unmount→remount reuses the in-flight request instead of aborting it); if no + * one re-subscribes by then, the request is aborted and the entry dropped. + */ + +/** What a `run` uses to drive its request and report state. */ +export interface RequestControls { + /** Aborted when the request is superseded or torn down. */ + signal: AbortSignal; + /** Abort this run's transport (e.g. to close a stream on a fatal frame). */ + abort(): void; + /** Merge fields into the entry's snapshot and notify subscribers. */ + patch(next: Partial): void; +} + +/** Starts a request and reports state through `controls`. */ +export type RequestRunner = (controls: RequestControls) => void; + +interface RequestStore { + /** + * Register a subscriber for `key`, creating and starting the shared request + * on first use. Returns a `release` function to call on unmount. + * + * @param run Runs the request; stored on the entry and re-invoked by + * `start`. Only the first caller's `run` is used (later joiners share it). + * @param autoStart Start the request on creation. Default true. + */ + retain(key: string, run: RequestRunner, autoStart?: boolean): () => void; + /** (Re)start the request for `key`: abort any in-flight run, then re-run. */ + start(key: string): void; + subscribe(key: string, listener: () => void): () => void; + getSnapshot(key: string): S; + /** Test-only: abort every in-flight request and clear the store. */ + reset(): void; +} + +interface Entry { + snapshot: S; + refCount: number; + abortController: AbortController | null; + teardownTimer: ReturnType | null; + /** True once `start` has run at least once; guards re-run on late `retain`. */ + started: boolean; + run: RequestRunner; +} + +export function createRequestStore(idle: S): RequestStore { + const entries = new Map>(); + + // Keyed separately from `entries`: `subscribe` can run before `retain` + // creates the entry, so listeners must survive independently of entry life. + const listenersByKey = new Map void>>(); + + function notify(key: string): void { + const listeners = listenersByKey.get(key); + if (!listeners) return; + for (const listener of listeners) listener(); + } + + function start(key: string): void { + const entry = entries.get(key); + if (!entry) return; + + entry.abortController?.abort(); + entry.started = true; + + const abortController = new AbortController(); + entry.abortController = abortController; + + entry.run({ + signal: abortController.signal, + abort: () => abortController.abort(), + patch(next) { + entry.snapshot = { ...entry.snapshot, ...next }; + notify(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 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); + } + + return { + retain(key, run, autoStart = true) { + let entry = entries.get(key); + if (!entry) { + entry = { + snapshot: idle, + refCount: 0, + abortController: null, + teardownTimer: null, + started: false, + run, + }; + 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); + }, + + start, + + subscribe(key, listener) { + 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); + }; + }, + + getSnapshot(key) { + return entries.get(key)?.snapshot ?? idle; + }, + + reset() { + 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 58bbfb0ce..e5f9e5e4b 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -4,19 +4,11 @@ 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, @@ -24,7 +16,6 @@ import type { QueryKey, UseAnalyticsQueryOptions, UseAnalyticsQueryResult, - WarehouseStatus, } from "./types"; import { useAnalyticsWarehousePublisher } from "./use-analytics-warehouse-status"; import { useQueryHMR } from "./use-query-hmr"; @@ -66,118 +57,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 @@ -220,13 +111,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 { @@ -262,87 +146,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, + }; }