@tanstack/react-query-next-experimental 5.0.0-alpha.80

package/src/HydrationStreamProvider.tsx

@@ -0,0 +1,184 @@
+'use client'
+
+import { useServerInsertedHTML } from 'next/navigation'
+import * as React from 'react'
+
+const serializedSymbol = Symbol('serialized')
+
+interface DataTransformer {
+  serialize(object: any): any
+  deserialize(object: any): any
+}
+
+type Serialized<TData> = unknown & {
+  [serializedSymbol]: TData
+}
+
+interface TypedDataTransformer<TData> {
+  serialize: (obj: TData) => Serialized<TData>
+  deserialize: (obj: Serialized<TData>) => TData
+}
+
+interface HydrationStreamContext<TShape> {
+  id: string
+  stream: {
+    /**
+     * **Server method**
+     * Push a new entry to the stream
+     * Will be ignored on the client
+     */
+    push: (...shape: TShape[]) => void
+  }
+}
+
+export interface HydrationStreamProviderProps<TShape> {
+  children: React.ReactNode
+  /**
+   * Optional transformer to serialize/deserialize the data
+   * Example devalue, superjson et al
+   */
+  transformer?: DataTransformer
+  /**
+   * **Client method**
+   * Called in the browser when new entries are received
+   */
+  onEntries: (entries: TShape[]) => void
+  /**
+   * **Server method**
+   * onFlush is called on the server when the cache is flushed
+   */
+  onFlush?: () => TShape[]
+}
+
+export function createHydrationStreamProvider<TShape>() {
+  const context = React.createContext<HydrationStreamContext<TShape>>(
+    null as any,
+  )
+  /**
+
+   * 1. (Happens on server): `useServerInsertedHTML()` is called **on the server** whenever a `Suspense`-boundary completes
+   *    - This means that we might have some new entries in the cache that needs to be flushed
+   *    - We pass these to the client by inserting a `<script>`-tag where we do `window[id].push(serializedVersionOfCache)`
+   * 2. (Happens in browser) In `useEffect()`:
+   *   - We check if `window[id]` is set to an array and call `push()` on all the entries which will call `onEntries()` with the new entries
+   *   - We replace `window[id]` with a `push()`-method that will be called whenever new entries are received
+   **/
+  function UseClientHydrationStreamProvider(props: {
+    children: React.ReactNode
+    /**
+     * Optional transformer to serialize/deserialize the data
+     * Example devalue, superjson et al
+     */
+    transformer?: DataTransformer
+    /**
+     * **Client method**
+     * Called in the browser when new entries are received
+     */
+    onEntries: (entries: TShape[]) => void
+    /**
+     * **Server method**
+     * onFlush is called on the server when the cache is flushed
+     */
+    onFlush?: () => TShape[]
+  }) {
+    // unique id for the cache provider
+    const id = `__RQ${React.useId()}`
+    const idJSON = JSON.stringify(id)
+
+    const [transformer] = React.useState(
+      () =>
+        (props.transformer ?? {
+          // noop
+          serialize: (obj: any) => obj,
+          deserialize: (obj: any) => obj,
+        }) as unknown as TypedDataTransformer<TShape>,
+    )
+
+    // <server stuff>
+    const [stream] = React.useState<TShape[]>(() => {
+      if (typeof window !== 'undefined') {
+        return {
+          push() {
+            // no-op on the client
+          },
+        } as unknown as TShape[]
+      }
+      return []
+    })
+    const count = React.useRef(0)
+    useServerInsertedHTML(() => {
+      // This only happens on the server
+      stream.push(...(props.onFlush?.() ?? []))
+
+      if (!stream.length) {
+        return null
+      }
+      // console.log(`pushing ${stream.length} entries`)
+      const serializedCacheArgs = stream
+        .map((entry) => transformer.serialize(entry))
+        .map((entry) => JSON.stringify(entry))
+        .join(',')
+
+      // Flush stream
+      stream.length = 0
+
+      const html: string[] = [
+        `window[${idJSON}] = window[${idJSON}] || [];`,
+        `window[${idJSON}].push(${serializedCacheArgs});`,
+      ]
+      return (
+        <script
+          key={count.current++}
+          dangerouslySetInnerHTML={{
+            __html: html.join(''),
+          }}
+        />
+      )
+    })
+    // </server stuff>
+
+    // <client stuff>
+    const onEntriesRef = React.useRef(props.onEntries)
+    React.useEffect(() => {
+      onEntriesRef.current = props.onEntries
+    })
+
+    React.useEffect(() => {
+      // Client: consume cache:
+      const onEntries = (...serializedEntries: Serialized<TShape>[]) => {
+        const entries = serializedEntries.map((serialized) =>
+          transformer.deserialize(serialized),
+        )
+        onEntriesRef.current(entries)
+      }
+
+      const win = window as any
+      // Register cache consumer
+      const winStream: Array<Serialized<TShape>> = win[id] ?? []
+
+      onEntries(...winStream)
+
+      // Register our own consumer
+      win[id] = {
+        push: onEntries,
+      }
+
+      return () => {
+        // Cleanup after unmount
+        win[id] = []
+      }
+    }, [id, transformer])
+    // </client stuff>
+
+    return (
+      <context.Provider value={{ stream, id }}>
+        {props.children}
+      </context.Provider>
+    )
+  }
+
+  return {
+    Provider: UseClientHydrationStreamProvider,
+    context,
+  }
+}

package/src/ReactQueryStreamedHydration.tsx

@@ -0,0 +1,93 @@
+'use client'
+
+import type {
+  DehydratedState,
+  DehydrateOptions,
+  HydrateOptions,
+  QueryClient,
+} from '@tanstack/react-query'
+import {
+  defaultShouldDehydrateQuery,
+  dehydrate,
+  hydrate,
+  useQueryClient,
+} from '@tanstack/react-query'
+import * as React from 'react'
+import type { HydrationStreamProviderProps } from './HydrationStreamProvider'
+import { createHydrationStreamProvider } from './HydrationStreamProvider'
+
+const stream = createHydrationStreamProvider<DehydratedState>()
+
+/**
+ * This component is responsible for:
+ * - hydrating the query client on the server
+ * - dehydrating the query client on the server
+ */
+export function ReactQueryStreamedHydration(props: {
+  children: React.ReactNode
+  queryClient?: QueryClient
+  options?: {
+    hydrate?: HydrateOptions
+    dehydrate?: DehydrateOptions
+  }
+  transformer?: HydrationStreamProviderProps<DehydratedState>['transformer']
+}) {
+  const queryClient = useQueryClient(props.queryClient)
+
+  /**
+   * We need to track which queries were added/updated during the render
+   */
+  const [trackedKeys] = React.useState(() => new Set<string>())
+
+  // <server only>
+  if (typeof window === 'undefined') {
+    // Do we need to care about unsubscribing? I don't think so to be honest
+    queryClient.getQueryCache().subscribe((event) => {
+      switch (event.type) {
+        case 'added':
+        case 'updated':
+          // console.log('tracking', event.query.queryHash, 'b/c of a', event.type)
+          trackedKeys.add(event.query.queryHash)
+      }
+    })
+  }
+  // </server only>
+
+  return (
+    <stream.Provider
+      // Happens on server:
+      onFlush={() => {
+        /**
+         * Dehydrated state of the client where we only include the queries that were added/updated since the last flush
+         */
+        const shouldDehydrate =
+          props.options?.dehydrate?.shouldDehydrateQuery ??
+          defaultShouldDehydrateQuery
+
+        const dehydratedState = dehydrate(queryClient, {
+          ...props.options?.dehydrate,
+          shouldDehydrateQuery(query) {
+            return trackedKeys.has(query.queryHash) && shouldDehydrate(query)
+          },
+        })
+        trackedKeys.clear()
+
+        if (!dehydratedState.queries.length) {
+          return []
+        }
+
+        return [dehydratedState]
+      }}
+      // Happens in browser:
+      onEntries={(entries) => {
+        for (const hydratedState of entries) {
+          hydrate(queryClient, hydratedState, props.options?.hydrate)
+        }
+      }}
+      // Handle BigInts etc using superjson
+      transformer={props.transformer}
+    >
+      {props.children}
+    </stream.Provider>
+  )
+}

package/src/index.ts

@@ -0,0 +1 @@
+export { ReactQueryStreamedHydration } from './ReactQueryStreamedHydration'