TanStack Start - Create tRPC-like API Router

_credit.txt

Most of this code came from the TanStack Start quickstart guide, plus a thread from the TanStack Discord

https://discord.com/channels/719702312431386674/1322986702553223288
https://tanstack.com/router/latest/docs/framework/react/start/overview

In particular, users matthewsomethin, ben-pr-p contributed the code above, and I pieced some of it together for my own purpose + documenting it here

api.ts

import {
  useMutation,
  UseMutationOptions,
  useQuery,
  UseQueryOptions,
} from "@tanstack/react-query";
import { UseQueryResult, UseMutationResult } from "@tanstack/react-query";
import { updateCount, getCount, getGreeting, getServerTime, getUserById } from "./serverFns";

type AsyncFunction = (...args: any[]) => Promise<any>;

type WrappedQuery<T extends AsyncFunction> = T & {
  useQuery: (
    args: Parameters<T>[0]["data"],
    options?: Omit<
      UseQueryOptions<
        Awaited<ReturnType<T>>,
        Error,
        Awaited<ReturnType<T>>,
        [string, Parameters<T>[0]]
      >,
      "queryKey" | "queryFn"
    >
  ) => UseQueryResult<Awaited<ReturnType<T>>>;
  useMutation: (
    options?: Omit<
      UseMutationOptions<Awaited<ReturnType<T>>, Error, Parameters<T>[0]>,
      "mutationFn"
    >
  ) => UseMutationResult<Awaited<ReturnType<T>>, Error, Parameters<T>[0]>;
};

type WrapperResult<T extends Record<string, AsyncFunction>> = {
  [K in keyof T]: WrappedQuery<T[K]>;
};

export function RouterWrapper<T extends Record<string, AsyncFunction>>(
  functions: T
): WrapperResult<T> {
  const wrappedFunctions: Partial<WrapperResult<T>> = {};

  for (const key in functions) {
    const func = functions[key];
    if (typeof func !== "function") continue;

    const wrappedFunc = func as WrappedQuery<typeof func>;

    wrappedFunc.useQuery = (
      args: Parameters<typeof func>[0]["data"],
      options?: Omit<
        UseQueryOptions<
          Awaited<ReturnType<typeof func>>,
          Error,
          Awaited<ReturnType<typeof func>>,
          [string, Parameters<typeof func>[0]]
        >,
        "queryKey" | "queryFn"
      >
    ) => {
      console.log("RouterWrapper useQuery", {
        key,
        args,
        options,
        func,
      });
      return useQuery({
        queryKey: [key, args],
        queryFn: () => func({ data: args }),
        ...options,
      });
    };

    wrappedFunc.useMutation = (
      options?: Omit<
        UseMutationOptions<
          Awaited<ReturnType<typeof func>>,
          Error,
          Parameters<typeof func>[0]
        >,
        "mutationFn"
      >
    ) =>
      useMutation({
        mutationFn: (args: Parameters<typeof func>[0]) => func(args),
        ...options,
      });

    wrappedFunctions[key] = wrappedFunc;
  }
  return wrappedFunctions as WrapperResult<T>;
}

const allServerFns = {
  updateCount, getCount, getGreeting, getServerTime, getUserById
};

export const api = RouterWrapper(allServerFns);

type UnwrapPromise<T> = T extends Promise<infer U> ? U : T;
type AsyncFunctions = { [key: string]: (...args: any[]) => Promise<any> };
type AsyncFnReturnTypes<T extends AsyncFunctions> = {
  [K in keyof T]: UnwrapPromise<ReturnType<T[K]>>;
};
export type ApiReturnTypes = AsyncFnReturnTypes<typeof allServerFns>;

Index.tsx

// app/routes/index.tsx
import { createFileRoute, Link, useRouter } from "@tanstack/react-router";
import React from "react";
import { queryClient } from "./__root";
import { api } from "../server/api";

export const Route = createFileRoute("/")({
  component: Home,
});

function Home() {
  const { data: count } = api.getCount.useQuery({});
  const { data: serverTime } = api.getServerTime.useQuery({});
  const { data: greeting } = api.getGreeting.useQuery({
    data: { name: "world" },
  });
  const { data: user } = api.getUserById.useQuery({ data: { id: 1 } });

  const updateCountMutation = api.updateCount.useMutation();

  console.log("serverHooks", {
    user,
    time: serverTime,
    clientTime: new Date().toISOString(),
  });

  return (
    <>
      <div>
        <button
          type="button"
          onClick={async () => {
            await updateCountMutation.mutateAsync({});
            queryClient.invalidateQueries({ queryKey: ["getCount"] });
          }}>
          Add 1 to {count}?
        </button>
      </div>
      <div>
        <span>{serverTime}</span>
      </div>
      <div>
        <span>{greeting}</span>
      </div>
      <div>
        <span>{user?.name}</span>, <span>{user?.id}</span>
      </div>
      <div>
        <Link to="/about">About</Link>
      </div>
    </>
  );
}

serverFns.ts

import { createServerFn } from "@tanstack/start";
import { z } from "zod";
import * as fs from "node:fs";

const filePath = "count.txt";
let count: number = 0;

const sampleUserDataset = [
  { id: 1, name: "Tommy" },
  { id: 2, name: "John" },
  { id: 3, name: "Jane" },
];

export const getUserById = createServerFn({
  method: "GET",
})
  .validator(z.object({ id: z.number() }))
  .handler(async ({ data }) => {
    const { id } = data;

    return sampleUserDataset.find((user) => user.id === id);
  });

export const getServerTime = createServerFn({
  method: "GET",
}).handler(async () => {
  return new Date().toISOString();
});

export const getGreeting = createServerFn({
  method: "GET",
})
  .validator(z.object({ name: z.string() }))
  .handler(async ({ data }) => {
    const { name } = data;
    return `Hello, ${name}!`;
  });

export const getCount = createServerFn({
  method: "GET",
}).handler(async () => {
  count = parseInt(
    await fs.promises.readFile(filePath, "utf-8").catch(() => "0")
  );
  return count;
});

export const updateCount = createServerFn({
  method: "POST",
}).handler(async () => {
  count += 1;
  await fs.promises.writeFile(filePath, count.toString());
});