feat: Remove experimental streamedQuery

Author: pulpdrew

packages/app/package.json

@@ -41,8 +41,8 @@
     "@mantine/spotlight": "7.9.2",
     "@microsoft/fetch-event-source": "^2.0.1",
     "@tabler/icons-react": "^3.5.0",
-    "@tanstack/react-query": "^5.90.2",
-    "@tanstack/react-query-devtools": "^5.90.2",
+    "@tanstack/react-query": "^5.56.2",
+    "@tanstack/react-query-devtools": "^5.56.2",
     "@tanstack/react-table": "^8.7.9",
     "@tanstack/react-virtual": "^3.0.1",
     "@uiw/codemirror-theme-atomone": "^4.23.3",

packages/app/src/hooks/useChartConfig.tsx

@@ -16,11 +16,7 @@ import {
   ChartConfigWithDateRange,
   ChartConfigWithOptDateRange,
 } from '@hyperdx/common-utils/dist/types';
-import {
-  experimental_streamedQuery as streamedQuery,
-  useQuery,
-  UseQueryOptions,
-} from '@tanstack/react-query';
+import { useQuery, UseQueryOptions } from '@tanstack/react-query';
 
 import {
   convertDateRangeToGranularityString,
@@ -52,6 +48,11 @@ type TQueryFnData = Pick<ResponseJSON<any>, 'data' | 'meta' | 'rows'> & {
   isComplete: boolean;
 };
 
+type TChunk = {
+  chunk: ResponseJSON<Record<string, string | number>>;
+  isComplete: boolean;
+};
+
 const shouldUseChunking = (
   config: ChartConfigWithOptDateRange,
 ): config is ChartConfigWithDateRange & {
@@ -158,6 +159,19 @@ async function* fetchDataInChunks(
   }
 }
 
+/** Append the given chunk to the given accumulated result */
+function appendChunk(
+  accumulated: TQueryFnData,
+  { chunk, isComplete }: TChunk,
+): TQueryFnData {
+  return {
+    data: [...(chunk.data || []), ...(accumulated?.data || [])],
+    meta: chunk.meta,
+    rows: (accumulated?.rows || 0) + (chunk.rows || 0),
+    isComplete,
+  };
+}
+
 /**
  * A hook providing data queried based on the provided chart config.
  *
@@ -186,36 +200,54 @@ export function useQueriedChartConfig(
     // Include disableQueryChunking in the query key to ensure that queries with the
     // same config but different disableQueryChunking values do not share a query
     queryKey: [config, options?.disableQueryChunking ?? false],
-    queryFn: streamedQuery({
-      streamFn: context =>
-        fetchDataInChunks(
-          config,
-          clickhouseClient,
-          context.signal,
-          options?.disableQueryChunking,
-        ),
-      /**
-       * This mode ensures that data remains in the cache until the next full streamed result is available.
-       * By default, the cache would be cleared before new data starts arriving, which results in the query briefly
-       * going back into the loading/pending state when multiple observers are sharing the query result resulting
-       * in flickering or render loops.
-       */
-      refetchMode: 'replace',
-      initialValue: {
+    // TODO: Replace this with `streamedQuery` when it is no longer experimental. Use 'replace' refetch mode.
+    // https://tanstack.com/query/latest/docs/reference/streamedQuery
+    queryFn: async context => {
+      const query = context.client
+        .getQueryCache()
+        .find({ queryKey: context.queryKey, exact: true });
+      const isRefetch = !!query && query.state.data !== undefined;
+
+      const emptyValue: TQueryFnData = {
         data: [],
         meta: [],
         rows: 0,
         isComplete: false,
-      } as TQueryFnData,
-      reducer: (acc, { chunk, isComplete }) => {
-        return {
-          data: [...(chunk.data || []), ...(acc?.data || [])],
-          meta: chunk.meta,
-          rows: (acc?.rows || 0) + (chunk.rows || 0),
-          isComplete,
-        };
-      },
-    }),
+      };
+
+      const chunks = fetchDataInChunks(
+        config,
+        clickhouseClient,
+        context.signal,
+        options?.disableQueryChunking,
+      );
+
+      let accumulatedChunks: TQueryFnData = emptyValue;
+      for await (const chunk of chunks) {
+        if (context.signal.aborted) {
+          break;
+        }
+
+        accumulatedChunks = appendChunk(accumulatedChunks, chunk);
+
+        // When refetching, the cache is not updated until all chunks are fetched.
+        if (!isRefetch) {
+          context.client.setQueryData<TQueryFnData>(
+            context.queryKey,
+            accumulatedChunks,
+          );
+        }
+      }
+
+      if (isRefetch && !context.signal.aborted) {
+        context.client.setQueryData<TQueryFnData>(
+          context.queryKey,
+          accumulatedChunks,
+        );
+      }
+
+      return context.client.getQueryData(context.queryKey)!;
+    },
     retry: 1,
     refetchOnWindowFocus: false,
     ...options,

yarn.lock

@@ -4595,8 +4595,8 @@ __metadata:
     "@storybook/react": "npm:^8.1.5"
     "@storybook/test": "npm:^8.1.5"
     "@tabler/icons-react": "npm:^3.5.0"
-    "@tanstack/react-query": "npm:^5.90.2"
-    "@tanstack/react-query-devtools": "npm:^5.90.2"
+    "@tanstack/react-query": "npm:^5.56.2"
+    "@tanstack/react-query-devtools": "npm:^5.56.2"
     "@tanstack/react-table": "npm:^8.7.9"
     "@tanstack/react-virtual": "npm:^3.0.1"
     "@testing-library/jest-dom": "npm:^6.4.2"
@@ -9399,40 +9399,40 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@tanstack/query-core@npm:5.90.2":
-  version: 5.90.2
-  resolution: "@tanstack/query-core@npm:5.90.2"
-  checksum: 10c0/695a7450b0bb9f6dd21bebeacfc962dfc886631a3b3a13c33a842ef719b4c3dd30c15febe8c1ade6902a85e0f387c51a97570f430cc8f5c7032ff737d6410597
+"@tanstack/query-core@npm:5.56.2":
+  version: 5.56.2
+  resolution: "@tanstack/query-core@npm:5.56.2"
+  checksum: 10c0/54ff55f02b01f6ba089f4965bfd46f430c18ce7e11d874de04c4d58cc8f698598b41e1c017ba029d08ae75e321e546b26f1ea7f788474db265eeba46e780f2f6
   languageName: node
   linkType: hard
 
-"@tanstack/query-devtools@npm:5.90.1":
-  version: 5.90.1
-  resolution: "@tanstack/query-devtools@npm:5.90.1"
-  checksum: 10c0/3b69e5441438acf0e753adbf187abf54b5b2e19d7c6d1e465d97278cb8c248bb86d3be193092d50414e4093cbf014093103517cb523daae003e53c867f3c11c2
+"@tanstack/query-devtools@npm:5.56.1":
+  version: 5.56.1
+  resolution: "@tanstack/query-devtools@npm:5.56.1"
+  checksum: 10c0/4f50fccf9e731e0fee4b7d2344cd8bf73a9c45816d7361f54ca4ec2df0a0455bf400fa3b578279edc8a4d1013c52c0d90cd50dbffb133edcf52e55868468d6c3
   languageName: node
   linkType: hard
 
-"@tanstack/react-query-devtools@npm:^5.90.2":
-  version: 5.90.2
-  resolution: "@tanstack/react-query-devtools@npm:5.90.2"
+"@tanstack/react-query-devtools@npm:^5.56.2":
+  version: 5.56.2
+  resolution: "@tanstack/react-query-devtools@npm:5.56.2"
   dependencies:
-    "@tanstack/query-devtools": "npm:5.90.1"
+    "@tanstack/query-devtools": "npm:5.56.1"
   peerDependencies:
-    "@tanstack/react-query": ^5.90.2
+    "@tanstack/react-query": ^5.56.2
     react: ^18 || ^19
-  checksum: 10c0/526d529bf995426ace7511f51a425ce92dfc1b6dd74c9956a3cd7d68950119e97291bced2ff17173bcdb329eae36c68abc211a4dec32d6e92ab537b41c0533c2
+  checksum: 10c0/2edebb04800585da1ca479f9289a9265cfcc26c5d8fb1db08e6fbe9fdf370a55a2007416b52f819c2faeab7b919cd51a4afca7731b50b6f3dd164057e4966741
   languageName: node
   linkType: hard
 
-"@tanstack/react-query@npm:^5.90.2":
-  version: 5.90.2
-  resolution: "@tanstack/react-query@npm:5.90.2"
+"@tanstack/react-query@npm:^5.56.2":
+  version: 5.56.2
+  resolution: "@tanstack/react-query@npm:5.56.2"
   dependencies:
-    "@tanstack/query-core": "npm:5.90.2"
+    "@tanstack/query-core": "npm:5.56.2"
   peerDependencies:
     react: ^18 || ^19
-  checksum: 10c0/22e76626a59890409858521b0e42b49219126a4ea5ed79eaa48a267959175dfdd28b30b9b03a415dccf703d95c18100a9d8917679818f6d2adc26d6c5f96a4d6
+  checksum: 10c0/6e883b4ca1948f990215b7bce194251faf13a79c6ecf3f3c660af6c6788ed113ab629cefdafb496dfb04866f12dd48d7314e936b75c881b6749127b6496ac8fd
   languageName: node
   linkType: hard