feat: Support flattened result structure

Author: Ehesp

examples/react-example/src/firebase.ts

@@ -7,5 +7,5 @@ if (getApps().length === 0) {
     projectId: "example",
   });
   const dataConnect = getDataConnect(connectorConfig);
-  connectDataConnectEmulator(dataConnect!, "localhost", 5003);
+  connectDataConnectEmulator(dataConnect!, "localhost", 9399);
 }

packages/react/package.json

@@ -32,8 +32,8 @@
   "license": "Apache-2.0",
   "devDependencies": {
     "@testing-library/react": "^16.0.1",
-    "react": "^18.3.1",
-    "@types/react": "^18.3.5",
+    "react": "^19.0.0",
+    "@types/react": "^19.0.1",
     "@dataconnect/default-connector": "workspace:*"
   },
   "peerDependencies": {

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

@@ -1,2 +1,3 @@
+export { DataConnectQueryClient } from "./query-client";
 export { useConnectQuery } from "./useConnectQuery";
 export { useConnectMutation } from "./useConnectMutation";

packages/react/src/data-connect/query-client.ts

@@ -0,0 +1,70 @@
+import {
+  QueryClient,
+  QueryKey,
+  type FetchQueryOptions,
+} from "@tanstack/react-query";
+import type { FirebaseError } from "firebase/app";
+import type { FlattenedQueryResult } from "./types";
+import { executeQuery, QueryRef, QueryResult } from "firebase/data-connect";
+
+type DataConnectQueryOptions<Data, Variables> = Omit<
+  FetchQueryOptions<
+    FlattenedQueryResult<Data, Variables>,
+    FirebaseError,
+    FlattenedQueryResult<Data, Variables>,
+    QueryKey
+  >,
+  "queryFn" | "queryKey"
+> & {
+  queryRef: QueryRef<Data, Variables>;
+  queryKey?: QueryKey;
+};
+
+export class DataConnectQueryClient extends QueryClient {
+  prefetchDataConnectQuery<Data extends Record<string, any>, Variables>(
+    refOrResult: QueryRef<Data, Variables> | QueryResult<Data, Variables>,
+    options?: DataConnectQueryOptions<Data, Variables>
+  ) {
+    let queryRef: QueryRef<Data, Variables>;
+    let initialData: FlattenedQueryResult<Data, Variables> | undefined;
+
+    if ("ref" in refOrResult) {
+      queryRef = refOrResult.ref;
+      initialData = {
+        ...refOrResult.data,
+        ref: refOrResult.ref,
+        source: refOrResult.source,
+        fetchTime: refOrResult.fetchTime,
+      };
+    } else {
+      queryRef = refOrResult;
+    }
+
+    return this.prefetchQuery<
+      FlattenedQueryResult<Data, Variables>,
+      FirebaseError,
+      FlattenedQueryResult<Data, Variables>,
+      QueryKey
+    >({
+      ...options,
+      initialData,
+      queryKey: options?.queryKey ?? [
+        queryRef.name,
+        queryRef.variables || null,
+      ],
+      queryFn: async () => {
+        const response = await executeQuery(queryRef);
+
+        const data = {
+          ...response.data,
+          ref: response.ref,
+          source: response.source,
+          fetchTime: response.fetchTime,
+        };
+
+        // Ensures no serialization issues with undefined values
+        return JSON.parse(JSON.stringify(data));
+      },
+    });
+  }
+}

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

@@ -0,0 +1,15 @@
+import { MutationResult, QueryResult } from "firebase/data-connect";
+
+// Flattens a QueryResult data down into a single object.
+// This is to prevent query.data.data, and also expose additional properties.
+export type FlattenedQueryResult<Data, Variables> = Omit<
+  QueryResult<Data, Variables>,
+  "data" | "toJSON"
+> &
+  Data;
+
+export type FlattenedMutationResult<Data, Variables> = Omit<
+  MutationResult<Data, Variables>,
+  "data" | "toJSON"
+> &
+  Data;

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

@@ -2,7 +2,6 @@ import {
   useMutation,
   useQueryClient,
   type UseMutationOptions,
-  type UseMutationResult,
 } from "@tanstack/react-query";
 import { type FirebaseError } from "firebase/app";
 import {
@@ -11,13 +10,16 @@ import {
   type DataConnect,
   type QueryRef,
 } from "firebase/data-connect";
+import { FlattenedMutationResult } from "./types";
 
 type UseConnectMutationOptions<
   TData = unknown,
   TError = FirebaseError,
   Variables = unknown
 > = Omit<UseMutationOptions<TData, TError, Variables>, "mutationFn"> & {
-  invalidate?: QueryRef<unknown, unknown>[];
+  invalidate?: Array<
+    QueryRef<unknown, unknown> | (() => QueryRef<unknown, unknown>)
+  >;
 };
 
 export function useConnectMutation<
@@ -33,26 +35,47 @@ export function useConnectMutation<
     : never
 >(
   ref: Fn,
-  options?: UseConnectMutationOptions<Data, FirebaseError, Variables>
-): UseMutationResult<Data, FirebaseError, Variables> {
+  options?: UseConnectMutationOptions<
+    FlattenedMutationResult<Data, Variables>,
+    FirebaseError,
+    Variables
+  >
+) {
   const queryClient = useQueryClient();
 
-  return useMutation<Data, FirebaseError, Variables>({
+  return useMutation<
+    FlattenedMutationResult<Data, Variables>,
+    FirebaseError,
+    Variables
+  >({
     ...options,
     onSuccess(...args) {
       if (options?.invalidate && options.invalidate.length) {
         for (const ref of options.invalidate) {
-          queryClient.invalidateQueries({
-            queryKey: [ref.name, ref.variables],
-          });
+          if ("variables" in ref) {
+            queryClient.invalidateQueries({
+              queryKey: [ref.name, ref.variables || null],
+              exact: true,
+            });
+          } else {
+            queryClient.invalidateQueries({
+              queryKey: [ref().name],
+            });
+          }
         }
       }
 
       options?.onSuccess?.(...args);
     },
-    mutationFn: async (variables: Variables) => {
-      const { data } = await executeMutation<Data, Variables>(ref(variables));
-      return data;
+    mutationFn: async (variables) => {
+      const response = await executeMutation<Data, Variables>(ref(variables));
+
+      return {
+        ...response.data,
+        ref: response.ref,
+        source: response.source,
+        fetchTime: response.fetchTime,
+      };
     },
   });
 }

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

@@ -1,6 +1,7 @@
 import { useQuery, type UseQueryOptions } from "@tanstack/react-query";
-import { type FirebaseError } from "firebase/app";
+import type { FirebaseError } from "firebase/app";
 import type { PartialBy } from "../../utils";
+import type { FlattenedQueryResult } from "./types";
 import {
   type QueryRef,
   type QueryResult,
@@ -12,30 +13,41 @@ type UseConnectQueryOptions<
   TError = FirebaseError
 > = PartialBy<Omit<UseQueryOptions<TData, TError>, "queryFn">, "queryKey">;
 
-export function useConnectQuery<
-  Data extends Record<string, any>,
-  Variables = unknown
->(
+export function useConnectQuery<Data = unknown, Variables = unknown>(
   refOrResult: QueryRef<Data, Variables> | QueryResult<Data, Variables>,
-  options?: UseConnectQueryOptions<Data, FirebaseError>
+  options?: UseConnectQueryOptions<
+    FlattenedQueryResult<Data, Variables>,
+    FirebaseError
+  >
 ) {
   let queryRef: QueryRef<Data, Variables>;
-  let initialData: Data | undefined;
+  let initialData: FlattenedQueryResult<Data, Variables> | undefined;
 
   if ("ref" in refOrResult) {
     queryRef = refOrResult.ref;
-    initialData = refOrResult.data;
+    initialData = {
+      ...refOrResult.data,
+      ref: refOrResult.ref,
+      source: refOrResult.source,
+      fetchTime: refOrResult.fetchTime,
+    };
   } else {
     queryRef = refOrResult;
   }
 
-  return useQuery<Data, FirebaseError>({
-    initialData,
+  return useQuery<FlattenedQueryResult<Data, Variables>, FirebaseError>({
     ...options,
-    queryKey: options?.queryKey ?? [queryRef.name, queryRef.variables],
+    initialData,
+    queryKey: options?.queryKey ?? [queryRef.name, queryRef.variables || null],
     queryFn: async () => {
-      const { data } = await executeQuery<Data, Variables>(queryRef);
-      return data;
+      const response = await executeQuery<Data, Variables>(queryRef);
+
+      return {
+        ...response.data,
+        ref: response.ref,
+        source: response.source,
+        fetchTime: response.fetchTime,
+      };
     },
   });
 }