From c8fa609991ac579d23a29c2eb35a1d6a3aaac917 Mon Sep 17 00:00:00 2001
From: Michael Ferris <micfer@getmaintainx.com>
Date: Fri, 15 Mar 2024 15:46:26 -0400
Subject: [PATCH] feat: Provide types and hook to set the data of a specific
 operation safely

---
 .../generators/generateReactQueryFunctions.ts | 21 ++++++++++++++++---
 plugins/typescript/src/templates/context.ts   | 19 +++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/plugins/typescript/src/generators/generateReactQueryFunctions.ts b/plugins/typescript/src/generators/generateReactQueryFunctions.ts
index 65a43898..6c0c1a82 100644
--- a/plugins/typescript/src/generators/generateReactQueryFunctions.ts
+++ b/plugins/typescript/src/generators/generateReactQueryFunctions.ts
@@ -80,6 +80,7 @@ export const generateReactQueryFunctions = async (
   const contextTypeName = `${c.pascal(filenamePrefix)}Context`;
   const nodes: ts.Node[] = [];
   const keyManagerItems: ts.TypeLiteralNode[] = [];
+  const queryOperationDataTypeItems: ts.TypeElement[] = [];
 
   const fetcherFilename = formatFilename(filenamePrefix + "-fetcher");
   const contextFilename = formatFilename(filenamePrefix + "-context");
@@ -143,8 +144,6 @@ export const generateReactQueryFunctions = async (
           ),
         });
 
-        
-
         const operationFetcherFnName = `fetch${c.pascal(operationId)}`;
         const operationQueryFnName = `${c.pascal(operationId)}Query`;
         const component: "useQuery" | "useMutate" =
@@ -157,7 +156,6 @@ export const generateReactQueryFunctions = async (
         }
 
         if (component === "useQuery") {
-          
           nodes.push(...declarationNodes);
 
           keyManagerItems.push(
@@ -185,6 +183,15 @@ export const generateReactQueryFunctions = async (
             ])
           );
 
+          queryOperationDataTypeItems.push(
+            f.createPropertySignature(
+              undefined,
+              f.createIdentifier(operationId),
+              undefined,
+              dataType
+            )
+          );
+
           nodes.push(
             ...createOperationFetcherFnNodes({
               dataType,
@@ -252,6 +259,13 @@ export const generateReactQueryFunctions = async (
         ])
   );
 
+  const queryOperationDataTypes = f.createTypeAliasDeclaration(
+    [f.createModifier(ts.SyntaxKind.ExportKeyword)],
+    "QueryDataTypes",
+    undefined,
+    f.createTypeLiteralNode(queryOperationDataTypeItems)
+  );
+
   const { nodes: usedImportsNodes, keys: usedImportsKeys } = getUsedImports(
     nodes,
     {
@@ -278,6 +292,7 @@ export const generateReactQueryFunctions = async (
       ...usedImportsNodes,
       ...nodes,
       queryKeyManager,
+      queryOperationDataTypes,
     ])
   );
 };
diff --git a/plugins/typescript/src/templates/context.ts b/plugins/typescript/src/templates/context.ts
index 30b75fd0..e45b597b 100644
--- a/plugins/typescript/src/templates/context.ts
+++ b/plugins/typescript/src/templates/context.ts
@@ -100,4 +100,23 @@ export const getContext = (prefix: string, componentsFile: string) =>
   } => {
     return Boolean((operation.variables as any).queryParams);
   };
+
+  /**
+   * Set the data of a query operation
+   * 
+   * @example
+   * const setQueryOperationData = useSetQueryOperationData();
+   * setQueryOperationData({
+   *   operationId: "getUser",
+   *   path: "/users/{id}",
+   *   variables: { pathParams: { id: "1" } }
+   * }, data /* data must have the type expected by that operation * /);
+   */
+  export function useSetQueryOperationData() {
+    const { queryKeyFn } = useAcApiContext();
+    const queryClient = useQueryClient();
+    return <Op extends QueryOperation>(operation: Op, data: QueryDataTypes[Op["operationId"]]) => {
+      queryClient.setQueryData(queryKeyFn(operation), data);
+    };
+  } 
   `;