feat: Support dc mutation factory fn

Author: Ehesp

packages/react/src/data-connect/types.ts

@@ -18,4 +18,4 @@ export type FlattenedMutationResult<Data, Variables> = Omit<
 // Helper function to determine if a key is a QueryKey.
 export function isQueryKey(key: unknown): key is QueryKey {
   return Array.isArray(key) && key.length > 0;
-}
+}

packages/react/src/data-connect/useDataConnectMutation.test.tsx

@@ -473,6 +473,48 @@ describe("useDataConnectMutation", () => {
     );
   });
 
+  test("invalidates queries specified in the invalidate option as a QueryKey", async () => {
+    const movieData = {
+      title: "tanstack query firebase",
+      genre: "library",
+      imageUrl: "https://invertase.io/",
+    };
+
+    const createdMovie = await createMovie(movieData);
+
+    const movieId = createdMovie?.data?.movie_insert?.id;
+
+    const { result } = renderHook(
+      () =>
+        useDataConnectMutation(createMovieRef, {
+          invalidate: [["GetMovieById", { id: movieId }]],
+        }),
+      {
+        wrapper,
+      },
+    );
+    const movie = {
+      title: "TanStack Query Firebase",
+      genre: "invalidate_option_test",
+      imageUrl: "https://test-image-url.com/",
+    };
+
+    await act(async () => {
+      await result.current.mutateAsync(movie);
+    });
+
+    await waitFor(() => {
+      expect(result.current.status).toBe("success");
+    });
+
+    expect(invalidateQueriesSpy.mock.calls).toHaveLength(1);
+    expect(invalidateQueriesSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        queryKey: ["GetMovieById", { id: movieId }],
+      }),
+    );
+  });
+
   test("invalidates queries specified in the invalidate option for create mutations with both variable and non-variable refs", async () => {
     const movieData = {
       title: "tanstack query firebase",
@@ -1015,4 +1057,65 @@ describe("useDataConnectMutation", () => {
       expect(deleteMutationResult.current.data?.movie_delete?.id).toBe(movieId);
     });
   });
+
+  test("executes mutation successfully with function ref", async () => {
+    const movie = {
+      title: "TanStack Query Firebase",
+      genre: "library",
+      imageUrl: "https://test-image-url.com/",
+    };
+
+    const { result } = renderHook(
+      () => useDataConnectMutation(() => createMovieRef(movie)),
+      {
+        wrapper,
+      },
+    );
+
+    await act(async () => {
+      await result.current.mutateAsync();
+    });
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBe(true);
+      expect(result.current.data).toHaveProperty("movie_insert");
+      expect(result.current.data?.movie_insert).toMatchObject({
+        title: movie.title,
+        genre: movie.genre,
+        imageUrl: movie.imageUrl,
+      });
+    });
+  });
+
+  test("executes mutation successfully with function ref that accepts variables", async () => {
+    const { result } = renderHook(
+      () =>
+        useDataConnectMutation((title: string) =>
+          createMovieRef({
+            title,
+            genre: "library",
+            imageUrl: "https://test-image-url.com/",
+          }),
+        ),
+      {
+        wrapper,
+      },
+    );
+
+    const movieTitle = "TanStack Query Firebase";
+
+    await act(async () => {
+      await result.current.mutateAsync(movieTitle);
+    });
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBe(true);
+      expect(result.current.data).toHaveProperty("movie_insert");
+      expect(result.current.data?.movie_insert).toMatchObject({
+        title: movieTitle,
+        genre: "library",
+        imageUrl: "https://test-image-url.com/",
+      });
+    });
+  });
 });

packages/react/src/data-connect/useDataConnectMutation.ts

@@ -24,15 +24,25 @@ export type useDataConnectMutationOptions<
 };
 
 export function useDataConnectMutation<
-  Fn extends (...args: any[]) => MutationRef<any, any>,
-  Data = ReturnType<Fn> extends MutationRef<infer D, any> ? D : never,
-  Variables = Fn extends (
-    dc: DataConnect,
-    vars: infer V
-  ) => MutationRef<any, any>
-    ? V
+  Fn extends
+    | (() => MutationRef<any, any>)
+    | ((vars: any) => MutationRef<any, any>)
+    | ((...args: any[]) => MutationRef<any, any>),
+  Data = ReturnType<
+    Fn extends (() => MutationRef<infer D, any>)
+      ? () => MutationRef<D, any>
+      : Fn extends (vars: any) => MutationRef<infer D, any>
+      ? (vars: any) => MutationRef<D, any>
+      : Fn
+  > extends MutationRef<infer D, any>
+    ? D
+    : never,
+  Variables = Fn extends () => MutationRef<any, any>
+    ? void
     : Fn extends (vars: infer V) => MutationRef<any, any>
     ? V
+    : Fn extends (dc: DataConnect, vars: infer V) => MutationRef<any, any>
+    ? V
     : never
 >(
   ref: Fn,
@@ -73,7 +83,8 @@ export function useDataConnectMutation<
       options?.onSuccess?.(...args);
     },
     mutationFn: async (variables) => {
-      const response = await executeMutation<Data, Variables>(ref(variables));
+      const mutationRef = typeof ref === "function" ? ref(variables) : ref;
+      const response = await executeMutation<Data, Variables>(mutationRef);
 
       return {
         ...response.data,