react query hydration helpers

apps/nextjs/src/app/_components/posts.tsx

@@ -1,7 +1,5 @@
 "use client";
 
-import { use } from "react";
-
 import type { RouterOutputs } from "@acme/api";
 import { cn } from "@acme/ui";
 import { Button } from "@acme/ui/button";
@@ -81,16 +79,11 @@ export function CreatePostForm() {
   );
 }
 
-export function PostList(props: {
-  posts: Promise<RouterOutputs["post"]["all"]>;
-}) {
-  // TODO: Make `useSuspenseQuery` work without having to pass a promise from RSC
-  const initialData = use(props.posts);
-  const { data: posts } = api.post.all.useQuery(undefined, {
-    initialData,
-  });
+export function PostList() {
+  const { data: posts } = api.post.all.useQuery();
+  console.log("posts", posts);
 
-  if (posts.length === 0) {
+  if (posts?.length === 0) {
     return (
       <div className="relative flex w-full flex-col gap-4">
         <PostCardSkeleton pulse={false} />
@@ -106,7 +99,7 @@ export function PostList(props: {
 
   return (
     <div className="flex w-full flex-col gap-4">
-      {posts.map((p) => {
+      {posts?.map((p) => {
         return <PostCard key={p.id} post={p} />;
       })}
     </div>

apps/nextjs/src/app/page.tsx

@@ -1,6 +1,6 @@
 import { Suspense } from "react";
 
-import { api } from "~/trpc/server";
+import { api, HydrateClient, setQueryData } from "~/trpc/server";
 import { AuthShowcase } from "./_components/auth-showcase";
 import {
   CreatePostForm,
@@ -12,7 +12,6 @@ export const runtime = "edge";
 
 export default async function HomePage() {
   // You can await this here if you don't want to show Suspense fallback below
-  const posts = api.post.all();
 
   return (
     <main className="container h-screen py-16">
@@ -33,10 +32,21 @@ export default async function HomePage() {
               </div>
             }
           >
-            <PostList posts={posts} />
+            <PostListWrapped />
           </Suspense>
         </div>
       </div>
     </main>
   );
 }
+
+async function PostListWrapped() {
+  const posts = await api.post.all();
+  setQueryData((t) => [t.post.all, undefined], posts); // would be cool if this was automatically put to the cache
+
+  return (
+    <HydrateClient>
+      <PostList />
+    </HydrateClient>
+  );
+}

apps/nextjs/src/trpc/hydration-helpers.tsx

@@ -0,0 +1,69 @@
+import type { QueryClient } from "@tanstack/react-query";
+import type {
+  AnyTRPCProcedure,
+  AnyTRPCQueryProcedure,
+  AnyTRPCRouter,
+  TRPCRouterRecord,
+} from "@trpc/server";
+import { dehydrate, HydrationBoundary } from "@tanstack/react-query";
+import {
+  createFlatProxy,
+  createRecursiveProxy,
+} from "@trpc/server/unstable-core-do-not-import";
+import { getQueryKeyInternal } from "node_modules/@trpc/react-query/dist/internals/getQueryKey";
+
+/**
+ * Some internal shit used to hydrate query client on the client
+ * tRPC should probably have something like this built-in. The current
+ * `getQueryKey` requires React createContext to be defined...
+ *
+ * This also provides some DX helpers to more easily perform hydration
+ * for queries fetched in RSC.
+ */
+
+type DecorateProcedure<TProcedure extends AnyTRPCProcedure> = (
+  input: TProcedure["_def"]["_input_in"],
+) => TProcedure["_def"]["_output_out"];
+
+type DecorateRouterRecord<TRecord extends TRPCRouterRecord> = {
+  [TKey in keyof TRecord]: TRecord[TKey] extends AnyTRPCQueryProcedure
+    ? DecorateProcedure<TRecord[TKey]>
+    : TRecord[TKey] extends TRPCRouterRecord
+      ? DecorateRouterRecord<TRecord[TKey]>
+      : never;
+};
+
+export function createHydrationHelpers<TRouter extends AnyTRPCRouter>(
+  getQueryClient: () => QueryClient,
+) {
+  type Proxy = DecorateRouterRecord<TRouter["_def"]["record"]>;
+
+  const proxy = createFlatProxy<Proxy>((key) =>
+    createRecursiveProxy(({ path }) => {
+      const fullPath = [key, ...path];
+      console.log("path", fullPath);
+      return fullPath;
+    }),
+  );
+
+  function setQueryData<T extends AnyTRPCQueryProcedure>(
+    cb: (p: Proxy) => [DecorateProcedure<T>, T["_def"]["_input_in"]],
+    data: T["_def"]["_output_out"],
+  ) {
+    const [proc, input] = cb(proxy);
+    const path = proc(input);
+    const key = getQueryKeyInternal(path, input, "query");
+    getQueryClient().setQueryData(key, data);
+  }
+
+  function HydrateClient(props: { children: React.ReactNode }) {
+    const state = dehydrate(getQueryClient());
+    console.log("state", state);
+
+    return (
+      <HydrationBoundary state={state}>{props.children}</HydrationBoundary>
+    );
+  }
+
+  return { setQueryData, HydrateClient };
+}

apps/nextjs/src/trpc/react.tsx

@@ -11,7 +11,18 @@ import type { AppRouter } from "@acme/api";
 export const api = createTRPCReact<AppRouter>();
 
 export function TRPCReactProvider(props: { children: React.ReactNode }) {
-  const [queryClient] = useState(() => new QueryClient());
+  const [queryClient] = useState(
+    () =>
+      new QueryClient({
+        defaultOptions: {
+          queries: {
+            // Since queries are prefetched on the server, we set a stale time so that
+            // queries aren't immediately refetched on the client
+            staleTime: 60 * 1000,
+          },
+        },
+      }),
+  );
 
   const [trpcClient] = useState(() =>
     api.createClient({

apps/nextjs/src/trpc/server.ts

@@ -1,9 +1,13 @@
 import { cache } from "react";
 import { headers } from "next/headers";
+import { QueryClient } from "@tanstack/react-query";
 
+import type { AppRouter } from "@acme/api";
 import { createCaller, createTRPCContext } from "@acme/api";
 import { auth } from "@acme/auth";
 
+import { createHydrationHelpers } from "./hydration-helpers";
+
 /**
  * This wraps the `createTRPCContext` helper and provides the required context for the tRPC API when
  * handling a tRPC call from a React Server Component.
@@ -19,3 +23,7 @@ const createContext = cache(async () => {
 });
 
 export const api = createCaller(createContext);
+
+const getQueryClient = cache(() => new QueryClient());
+export const { HydrateClient, setQueryData } =
+  createHydrationHelpers<AppRouter>(getQueryClient);