From 51291d44e0b6d3754382bafbf4e31cfee8dba4ca Mon Sep 17 00:00:00 2001
From: Nate Kelley <nate@buster.so>
Date: Wed, 12 Feb 2025 11:56:59 -0700
Subject: [PATCH] update query

---
 web/src/api/createReactQuery.ts               | 20 +++----
 .../useBusterWebSocketQuery/types.ts          | 20 ++++++-
 .../useBusterWebSocketQuery.tsx               | 55 ++++++++++---------
 3 files changed, 55 insertions(+), 40 deletions(-)

diff --git a/web/src/api/createReactQuery.ts b/web/src/api/createReactQuery.ts
index d26487bee..abaf140c6 100644
--- a/web/src/api/createReactQuery.ts
+++ b/web/src/api/createReactQuery.ts
@@ -7,22 +7,16 @@ import {
   useQuery,
   UseQueryOptions,
   keepPreviousData,
-  useInfiniteQuery
+  useInfiniteQuery,
+  QueryKey
 } from '@tanstack/react-query';
 import { useEffect } from 'react';
 import { useBusterNotifications } from '@/context/BusterNotifications';
 import { RustApiError } from './buster_rest/errors';
 import { useMemoizedFn } from 'ahooks';
 
-export interface BaseCreateQueryProps {
-  refetchOnWindowFocus?: boolean;
-  refetchOnMount?: boolean;
-  enabled?: boolean;
-  staleTime?: number;
-  accessToken?: string;
-}
-interface CreateQueryProps<T> extends UseQueryOptions<T> {
-  queryKey: (string | number | object)[];
+interface CreateQueryProps<T> extends UseQueryOptions {
+  queryKey: QueryKey;
   isUseSession?: boolean;
   useErrorNotification?: boolean;
 }
@@ -40,7 +34,7 @@ export const useCreateReactQuery = <T>({
   useErrorNotification = true,
   staleTime,
   ...rest
-}: CreateQueryProps<T> & BaseCreateQueryProps) => {
+}: CreateQueryProps<T>) => {
   const { openErrorNotification } = useBusterNotifications();
   const accessToken = useSupabaseContext((state) => state.accessToken);
   const baseEnabled = isUseSession ? !!accessToken : true;
@@ -124,7 +118,7 @@ export const useCreateReactQueryPaginated = <T>({
   page = 0,
   pageSize = 25,
   ...rest
-}: PaginatedQueryProps<T> & BaseCreateQueryProps) => {
+}: PaginatedQueryProps<T>) => {
   const accessToken = useSupabaseContext((state) => state.accessToken);
   const baseEnabled = isUseSession ? !!accessToken : true;
 
@@ -152,7 +146,7 @@ export const useCreateReactInfiniteQuery = <T>({
   initialPageParam = 0,
   getNextPageParam,
   ...rest
-}: Parameters<typeof useInfiniteQuery>[0] & BaseCreateQueryProps) => {
+}: Parameters<typeof useInfiniteQuery>[0]) => {
   const accessToken = useSupabaseContext((state) => state.accessToken);
   const baseEnabled = !!accessToken;
 
diff --git a/web/src/context/BusterWebSocket/useBusterWebSocketQuery/types.ts b/web/src/context/BusterWebSocket/useBusterWebSocketQuery/types.ts
index 83aef5a2e..d7cb81dd7 100644
--- a/web/src/context/BusterWebSocket/useBusterWebSocketQuery/types.ts
+++ b/web/src/context/BusterWebSocket/useBusterWebSocketQuery/types.ts
@@ -1,7 +1,23 @@
-import { BusterSocketRequest, BusterSocketResponse } from '@/api/buster_socket';
+import type { BusterSocketResponse, BusterSocketResponseRoute } from '@/api/buster_socket';
 import { UseQueryOptions, UseQueryResult } from '@tanstack/react-query';
 
-export type QueryKey = readonly unknown[];
+/**
+ * Infers the response data type from a BusterSocket route
+ */
+export type InferBusterSocketResponseData<TRoute extends BusterSocketResponseRoute> = Extract<
+  BusterSocketResponse,
+  { route: TRoute }
+>['callback'] extends (d: infer D) => void
+  ? D
+  : never;
+
+/**
+ * Socket response configuration with optional error handler
+ */
+export type BusterSocketResponseConfig<TRoute extends BusterSocketResponseRoute> = {
+  route: TRoute;
+  onError?: (d: unknown) => void;
+};
 
 export interface UseBusterSocketQueryOptions<TData, TError = unknown>
   extends Omit<UseQueryOptions<TData, TError>, 'queryKey' | 'queryFn'> {}
diff --git a/web/src/context/BusterWebSocket/useBusterWebSocketQuery/useBusterWebSocketQuery.tsx b/web/src/context/BusterWebSocket/useBusterWebSocketQuery/useBusterWebSocketQuery.tsx
index f5dfd25d9..cb70f015a 100644
--- a/web/src/context/BusterWebSocket/useBusterWebSocketQuery/useBusterWebSocketQuery.tsx
+++ b/web/src/context/BusterWebSocket/useBusterWebSocketQuery/useBusterWebSocketQuery.tsx
@@ -1,32 +1,28 @@
-import { useQuery, QueryKey } from '@tanstack/react-query';
-import {
+import { useQuery, QueryKey, UseQueryOptions } from '@tanstack/react-query';
+import type {
   BusterSocketRequest,
   BusterSocketResponse,
   BusterSocketResponseRoute
 } from '@/api/buster_socket';
 import { useBusterWebSocket } from '../useBusterWebSocket';
 import { transformError } from './helpers';
-import { UseBusterSocketQueryOptions, UseBusterSocketQueryResult } from './types';
+import type {
+  UseBusterSocketQueryOptions,
+  UseBusterSocketQueryResult,
+  InferBusterSocketResponseData,
+  BusterSocketResponseConfig
+} from './types';
+import { useCreateReactQuery } from '@/api/createReactQuery';
 
 export function useBusterWebSocketQuery<TRoute extends BusterSocketResponseRoute, TError = unknown>(
   queryKey: QueryKey,
   socketRequest: BusterSocketRequest,
-  socketResponse: { route: TRoute; onError?: (d: unknown) => void },
-  options: UseBusterSocketQueryOptions<
-    Extract<BusterSocketResponse, { route: TRoute }>['callback'] extends (d: infer D) => void
-      ? D
-      : never,
-    TError
-  > = {}
-): UseBusterSocketQueryResult<
-  Extract<BusterSocketResponse, { route: TRoute }>['callback'] extends (d: infer D) => void
-    ? D
-    : never,
-  TError
-> {
+  socketResponse: BusterSocketResponseConfig<TRoute>,
+  options?: UseQueryOptions<InferBusterSocketResponseData<TRoute>, TError>
+): UseBusterSocketQueryResult<InferBusterSocketResponseData<TRoute>, TError> {
   const busterSocket = useBusterWebSocket();
 
-  const queryFn = async () => {
+  const queryFn = async (): Promise<InferBusterSocketResponseData<TRoute>> => {
     try {
       const result = await busterSocket.emitAndOnce({
         emitEvent: socketRequest,
@@ -37,19 +33,26 @@ export function useBusterWebSocketQuery<TRoute extends BusterSocketResponseRoute
         } as BusterSocketResponse
       });
 
-      return result as Extract<BusterSocketResponse, { route: TRoute }>['callback'] extends (
-        d: infer D
-      ) => void
-        ? D
-        : never;
+      return result as InferBusterSocketResponseData<TRoute>;
     } catch (error) {
-      throw transformError(error);
+      throw error;
     }
   };
 
-  return useQuery({
+  // return useCreateReactQuery<InferBusterSocketResponseData<TRoute>>({
+  //   queryKey,
+  //   queryFn,
+  //   isUseSession: false
+  // });
+
+  return useQuery<
+    InferBusterSocketResponseData<TRoute>,
+    TError,
+    InferBusterSocketResponseData<TRoute>
+  >({
     queryKey,
-    queryFn
+    queryFn,
+    ...options
   });
 }
 
@@ -60,4 +63,6 @@ export const ExampleUsage = () => {
     { route: '/chats/get', payload: { id: '123' } },
     { route: '/chats/get:getChat' }
   );
+
+  useCreateReactQuery;
 };