Added tracking to fdc

invertase · Jan 31, 2025 · 76f4a48 · 76f4a48
1 parent 65e7b8f
commit 76f4a48
Show file tree

Hide file tree

Showing 4 changed files with 236 additions and 5 deletions.
diff --git a/packages/angular/src/data-connect/index.ts b/packages/angular/src/data-connect/index.ts
@@ -0,0 +1,216 @@
+import {
+  CreateMutationOptions,
+  CreateMutationResult,
+  CreateQueryOptions,
+  CreateQueryResult,
+  injectMutation,
+  injectQuery,
+  QueryClient,
+  QueryKey,
+} from "@tanstack/angular-query-experimental";
+import {
+  DataConnect,
+  executeMutation,
+  executeQuery,
+  MutationRef,
+  MutationResult,
+  QueryRef,
+  QueryResult,
+} from "firebase/data-connect";
+import { FirebaseError } from "firebase/app";
+
+import { inject, signal } from "@angular/core";
+function getQueryKey(queryRef: QueryRef<unknown, unknown>) {
+  return [queryRef.name, queryRef.variables];
+}
+type FlattenedQueryResult<Data, Variables> = Omit<
+  QueryResult<Data, Variables>,
+  "data" | "toJSON"
+> &
+  Data;
+interface CreateDataConnectQueryOptions<Data, Variables>
+  extends Omit<
+    CreateQueryOptions<
+      FlattenedQueryResult<Data, Variables>,
+      FirebaseError,
+      FlattenedQueryResult<Data, Variables>,
+      QueryKey
+    >,
+    "queryFn" | "queryKey"
+  > {
+  queryFn: () => QueryRef<Data, Variables>;
+}
+
+/**
+ * injectDataConnectQuery takes a query ref and returns a wrapper function around Tanstack's `injectQuery`
+ * @param queryRefOrOptionsFn Query Ref or callback function for calling a new query
+ * @returns {CreateQueryResult<FlattenedQueryResult<Data, Variables>>}
+ */
+export function injectDataConnectQuery<Data, Variables>(
+  queryRefOrOptionsFn:
+    | QueryRef<Data, Variables>
+    | (() => CreateDataConnectQueryOptions<Data, Variables>)
+): CreateQueryResult<FlattenedQueryResult<Data, Variables>, FirebaseError> {
+  const queryKey = signal<QueryKey>([]);
+  function fdcOptionsFn() {
+    const passedInOptions =
+      typeof queryRefOrOptionsFn === "function"
+        ? queryRefOrOptionsFn()
+        : undefined;
+
+    const modifiedFn = () => {
+      const ref: QueryRef<Data, Variables> =
+        passedInOptions?.queryFn() ||
+        (queryRefOrOptionsFn as QueryRef<Data, Variables>);
+      queryKey.set([ref.name, ref.variables]);
+      return executeQuery(ref).then((res) => {
+        const { data, ...rest } = res;
+        return {
+          ...data,
+          ...rest,
+        };
+      }) as Promise<FlattenedQueryResult<Data, Variables>>;
+    };
+    return {
+      queryKey: queryKey(),
+      ...passedInOptions,
+      queryFn: modifiedFn,
+    };
+  }
+  return injectQuery(fdcOptionsFn);
+}
+
+export type GeneratedSignature<Data, Variables> = (
+  dc: DataConnect,
+  vars: Variables
+) => MutationRef<Data, Variables>;
+export type DataConnectMutationOptionsFn<Data, Error, Variables, Arguments> =
+  () => CreateDataConnectMutationOptions<Data, Variables, Arguments>;
+export type DataConnectMutationOptionsUndefinedMutationFn<
+  Data,
+  Error,
+  Variables
+> = () => Omit<
+  ReturnType<DataConnectMutationOptionsFn<Data, Error, Variables, Variables>>,
+  "mutationFn"
+>;
+export type FlattenedMutationResult<Data, Variables> = Omit<
+  MutationResult<Data, Variables>,
+  "data" | "toJSON"
+> &
+  Data;
+
+  export interface CreateDataConnectMutationOptions<Data, Variables, Arguments = Variables> extends Omit<CreateMutationOptions<Data, FirebaseError, Arguments>, "mutationFn"> {
+    invalidate?: QueryKey | QueryRef<Data, Variables>[];
+    dataConnect?: DataConnect;
+    mutationFn: (args: Arguments) => MutationRef<Data, Variables>;
+};
+
+export function injectDataConnectMutation<Data, Variables, Arguments>(
+  factoryFn: undefined,
+  optionsFn: DataConnectMutationOptionsFn<
+    Data,
+    FirebaseError,
+    Variables,
+    Arguments
+  >
+): CreateMutationResult<
+  FlattenedMutationResult<Data, Variables>,
+  FirebaseError,
+  Arguments
+>;
+export function injectDataConnectMutation<
+  Data,
+  Variables extends undefined,
+  Arguments = void | undefined
+>(
+  factoryFn: GeneratedSignature<Data, Variables>
+): CreateMutationResult<
+  FlattenedMutationResult<Data, Variables>,
+  FirebaseError,
+  Arguments
+>;
+export function injectDataConnectMutation<Data, Variables, Arguments>(
+  factoryFn: GeneratedSignature<Data, Variables>
+): CreateMutationResult<
+  FlattenedMutationResult<Data, Variables>,
+  FirebaseError,
+  Arguments
+>;
+export function injectDataConnectMutation<Data, Variables, Arguments>(
+  factoryFn: GeneratedSignature<Data, Variables>,
+  optionsFn?: DataConnectMutationOptionsUndefinedMutationFn<
+    Data,
+    FirebaseError,
+    Arguments
+  >
+): CreateMutationResult<
+  FlattenedMutationResult<Data, Variables>,
+  FirebaseError,
+  Arguments
+>;
+/**
+ * injectDataConnectMutation takes a mutation ref factory function and returns a tanstack wrapper around `injectMutation`
+ * @param factoryFn generated SDK factory function
+ * @param optionsFn options function to create a new mutation
+ * @returns {CreateMutationResult<FlattenedMutationResult<Data, Variables>, FirebaseError, Arguments>}
+ */
+export function injectDataConnectMutation<
+  Data,
+  Variables,
+  Arguments extends Variables
+>(
+  factoryFn: GeneratedSignature<Data, Variables> | undefined,
+  optionsFn?:
+    | DataConnectMutationOptionsFn<Data, FirebaseError, Variables, Arguments>
+    | DataConnectMutationOptionsUndefinedMutationFn<
+        Data,
+        FirebaseError,
+        Variables
+      >
+): CreateMutationResult<
+  FlattenedMutationResult<Data, Variables>,
+  FirebaseError,
+  Arguments
+> {
+  const injectCb = () => {
+    const queryClient = inject(QueryClient);
+    const providedOptions = optionsFn ? optionsFn() : undefined;
+    const modifiedFn = (args: Arguments) => {
+      const dataConnect = inject(DataConnect);
+      const ref =
+        (providedOptions &&
+          "mutationFn" in providedOptions &&
+          providedOptions.mutationFn(args as Arguments)) ||
+        factoryFn!(dataConnect, args as Variables);
+
+      return executeMutation(ref)
+        .then((res) => {
+          const { data, ...rest } = res;
+          return {
+            ...data,
+            ...rest,
+          };
+        })
+        .then((ret) => {
+          providedOptions?.invalidate?.forEach((qk) => {
+            let key = qk;
+            if ("name" in (qk as Object)) {
+              const queryKey = getQueryKey(qk as QueryRef<unknown, unknown>);
+              key = queryKey;
+            }
+            queryClient.invalidateQueries({
+              queryKey: key,
+            });
+          });
+          return ret;
+        }) as Promise<FlattenedMutationResult<Data, Variables>>;
+    };
+
+    return {
+      ...providedOptions,
+      mutationFn: modifiedFn,
+    };
+  };
+  return injectMutation(injectCb);
+}
diff --git a/packages/react/src/data-connect/types.ts b/packages/react/src/data-connect/types.ts
@@ -13,3 +13,12 @@ export type FlattenedMutationResult<Data, Variables> = Omit<
   "data" | "toJSON"
 > &
   Data;
+
+export enum CallerSdkType {
+  Base, // Core JS SDK
+  Generated, // Generated JS SDK
+  TanstackReactCore, // Tanstack non-generated React SDK
+  GeneratedReact, // Generated React SDK
+  TanstackAngularCore, // Tanstack non-generated Angular SDK
+  GeneratedAngular // Generated Angular SDK
+}
diff --git a/packages/react/src/data-connect/useDataConnectMutation.ts b/packages/react/src/data-connect/useDataConnectMutation.ts
@@ -10,7 +10,7 @@ import {
   type QueryRef,
   executeMutation,
 } from "firebase/data-connect";
-import type { FlattenedMutationResult } from "./types";
+import { CallerSdkType, type FlattenedMutationResult } from "./types";
 
 export type useDataConnectMutationOptions<
   TData = unknown,
@@ -48,10 +48,10 @@ export function useDataConnectMutation<
     FlattenedMutationResult<Data, Variables>,
     FirebaseError,
     Variables
-  >
+  >,
+  _callerSdkType: CallerSdkType = CallerSdkType.TanstackReactCore
 ) {
   const queryClient = useQueryClient();
-
   return useMutation<
     FlattenedMutationResult<Data, Variables>,
     FirebaseError,
@@ -78,6 +78,9 @@ export function useDataConnectMutation<
     },
     mutationFn: async (variables) => {
       const mutationRef = typeof ref === "function" ? ref(variables) : ref;
+
+      // @ts-expect-error function is hidden under `DataConnect`.
+      mutationRef.dataConnect._setCallerSdkType(_callerSdkType);
       const response = await executeMutation<Data, Variables>(mutationRef);
 
       return {

diff --git a/packages/react/src/data-connect/useDataConnectQuery.ts b/packages/react/src/data-connect/useDataConnectQuery.ts
@@ -6,7 +6,8 @@ import {
   executeQuery,
 } from "firebase/data-connect";
 import type { PartialBy } from "../../utils";
-import type { FlattenedQueryResult } from "./types";
+import { CallerSdkType, type FlattenedQueryResult } from "./types";
+
 
 export type useDataConnectQueryOptions<
   TData = unknown,
@@ -19,6 +20,7 @@ export function useDataConnectQuery<Data = unknown, Variables = unknown>(
     FlattenedQueryResult<Data, Variables>,
     FirebaseError
   >,
+  _callerSdkType: CallerSdkType = CallerSdkType.TanstackReactCore
 ) {
   let queryRef: QueryRef<Data, Variables>;
   let initialData: FlattenedQueryResult<Data, Variables> | undefined;
@@ -34,7 +36,8 @@ export function useDataConnectQuery<Data = unknown, Variables = unknown>(
   } else {
     queryRef = refOrResult;
   }
-
+  // @ts-expect-error function is hidden under `DataConnect`.
+  queryRef.dataConnect._setCallerSdkType(_callerSdkType);
   return useQuery<FlattenedQueryResult<Data, Variables>, FirebaseError>({
     ...options,
     initialData,