feat: Support query key invalidate arg

Author: Ehesp

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

@@ -1,3 +1,4 @@
+import type { QueryKey } from "@tanstack/react-query";
 import type { MutationResult, QueryResult } from "firebase/data-connect";
 
 // Flattens a QueryResult data down into a single object.
@@ -13,3 +14,8 @@ export type FlattenedMutationResult<Data, Variables> = Omit<
   "data" | "toJSON"
 > &
   Data;
+
+// 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.ts

@@ -1,5 +1,6 @@
 import {
   type UseMutationOptions,
+  type QueryKey,
   useMutation,
   useQueryClient,
 } from "@tanstack/react-query";
@@ -10,15 +11,15 @@ import {
   type QueryRef,
   executeMutation,
 } from "firebase/data-connect";
-import type { FlattenedMutationResult } from "./types";
+import { isQueryKey, type FlattenedMutationResult } from "./types";
 
 export type useDataConnectMutationOptions<
   TData = unknown,
   TError = FirebaseError,
-  Variables = unknown,
+  Variables = unknown
 > = Omit<UseMutationOptions<TData, TError, Variables>, "mutationFn"> & {
-  invalidate?: Array<
-    QueryRef<unknown, unknown> | (() => QueryRef<unknown, unknown>)
+  invalidate?: ReadonlyArray<
+    QueryRef<unknown, unknown> | (() => QueryRef<unknown, unknown>) | QueryKey
   >;
 };
 
@@ -27,19 +28,19 @@ export function useDataConnectMutation<
   Data = ReturnType<Fn> extends MutationRef<infer D, any> ? D : never,
   Variables = Fn extends (
     dc: DataConnect,
-    vars: infer V,
+    vars: infer V
   ) => MutationRef<any, any>
     ? V
     : Fn extends (vars: infer V) => MutationRef<any, any>
-      ? V
-      : never,
+    ? V
+    : never
 >(
   ref: Fn,
   options?: useDataConnectMutationOptions<
     FlattenedMutationResult<Data, Variables>,
     FirebaseError,
     Variables
-  >,
+  >
 ) {
   const queryClient = useQueryClient();
 
@@ -57,6 +58,10 @@ export function useDataConnectMutation<
               queryKey: [ref.name, ref.variables],
               exact: true,
             });
+          } else if (isQueryKey(ref)) {
+            queryClient.invalidateQueries({
+              queryKey: ref,
+            });
           } else {
             queryClient.invalidateQueries({
               queryKey: [ref.name],