trpc-rx.ts

import type {
  MutationObserverResult,
  QueryObserverResult,
} from "@tanstack/query-core";
import { QueryClient } from "@tanstack/query-core";
import type { CreateTRPCClientOptions } from "@trpc/client";
import { createTRPCClient } from "@trpc/client";
import type {
  AnyRouter,
  AnyProcedure,
  ProcedureArgs,
  inferProcedureOutput,
  AnyMutationProcedure,
  AnyQueryProcedure,
  ProcedureRouterRecord,
  inferProcedureClientError,
  inferProcedureInput,
} from "@trpc/server";
import { createRecursiveProxy } from "@trpc/server/shared";
import { BehaviorSubject, from, Observable } from "rxjs";
import { createRxQueryHooks } from "./rx-query";

export type inferRouterRxClient<TRouter extends AnyRouter> =
  DecoratedProcedureRecord<TRouter["_def"]["record"]>;

type DecorateProcedure<TProcedure extends AnyProcedure> =
  TProcedure extends AnyQueryProcedure
    ? {
        query: (
          ...args: ProcedureArgs<TProcedure["_def"]>
        ) => BehaviorSubject<
          QueryObserverResult<
            inferProcedureOutput<TProcedure>,
            inferProcedureClientError<TProcedure>
          >
        >;
        invalidate: () => Observable<void>;
        refetch: () => Observable<void>;
      }
    : TProcedure extends AnyMutationProcedure
    ? {
        mutate: () => BehaviorSubject<
          MutationObserverResult<
            inferProcedureOutput<TProcedure>,
            inferProcedureClientError<TProcedure>,
            inferProcedureInput<TProcedure>
          >
        >;
      }
    : never;

type DecoratedProcedureRecord<TProcedures extends ProcedureRouterRecord> = {
  [TKey in keyof TProcedures]: TProcedures[TKey] extends AnyRouter
    ? DecoratedProcedureRecord<TProcedures[TKey]["_def"]["record"]>
    : TProcedures[TKey] extends AnyProcedure
    ? DecorateProcedure<TProcedures[TKey]>
    : never;
};

const clientCallTypeMap = {
  query: "query",
  mutate: "mutation",
  invalidate: "invalidate",
  refetch: "refetch",
} as const;

const hookCallTypeMap = (qc: QueryClient) => {
  const hooks = createRxQueryHooks(qc);
  return { query: hooks.fromQuery, mutation: hooks.fromMutation };
};

export function createTRPCRxClient<TRouter extends AnyRouter>(
  opts: CreateTRPCClientOptions<TRouter>
) {
  const client = createTRPCClient<TRouter>(opts);
  const qc = new QueryClient();

  const proxy = createRecursiveProxy(({ path, args }) => {
    console.log(path, args);
    const pathCopy: string[] = [...path];
    const clientCallType = pathCopy.pop() as keyof typeof clientCallTypeMap;

    const procedureType = clientCallTypeMap[clientCallType];

    const fullPath = pathCopy.join(".");

    if (procedureType === "invalidate") {
      return from(qc.invalidateQueries([fullPath, ...args]));
    }
    if (procedureType === "refetch") {
      return from(qc.refetchQueries([fullPath, ...args]));
    }

    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    return (hookCallTypeMap(qc)[procedureType] as any)({
      queryKey: [fullPath, ...args],
      queryFn: () => client.query(fullPath, ...args),

      mutationFn: (input: unknown) => client.mutation(fullPath, input),

      retry: false,
    });
  });

  return proxy as inferRouterRxClient<TRouter>;
}