refactor(ssr): better ssr handling

Author: posva

src/firestore/index.ts

@@ -8,6 +8,7 @@ import {
   getFirestore,
 } from 'firebase/firestore'
 import {
+  getCurrentInstance,
   getCurrentScope,
   isRef,
   onScopeDispose,
@@ -31,6 +32,7 @@ import {
   _Nullable,
   _RefWithState,
 } from '../shared'
+import { getInitialValue } from '../ssr/initialState'
 import { addPendingPromise } from '../ssr/plugin'
 import { firestoreUnbinds } from './optionsApi'
 import {
@@ -51,6 +53,12 @@ export interface _UseFirestoreRefOptions extends FirestoreRefOptions {
    * Use the `target` ref instead of creating one.
    */
   target?: Ref<unknown>
+
+  /**
+   * Optional key to handle SSR hydration. **Necessary for Queries** or when the same source is used in multiple places
+   * with different converters.
+   */
+  ssrKey?: string
 }
 
 /**
@@ -68,14 +76,16 @@ export function _useFirestoreRef(
 ) {
   let _unbind: UnbindWithReset = noop
   const options = Object.assign({}, firestoreOptions, localOptions)
+  const initialSourceValue = unref(docOrCollectionRef)
 
+  const data = options.target || ref<unknown | null>()
+  // set the initial value from SSR even if the ref comes from outside
+  data.value = getInitialValue('f', options.ssrKey, initialSourceValue)
   // TODO: allow passing pending and error refs as option for when this is called using the options api
-  const data = options.target || ref<unknown | null>(options.initialValue)
   const pending = ref(true)
   const error = ref<FirestoreError>()
   // force the type since its value is set right after and undefined isn't possible
   const promise = shallowRef() as ShallowRef<Promise<unknown | null>>
-  let isPromiseAdded = false
   const hasCurrentScope = getCurrentScope()
   let removePendingPromise = noop
 
@@ -112,12 +122,6 @@ export function _useFirestoreRef(
       )
     })
 
-    // only add the first promise to the pending ones
-    if (!isPromiseAdded && docRefValue) {
-      // TODO: is there a way to make this only for the first render?
-      removePendingPromise = addPendingPromise(p, docRefValue)
-      isPromiseAdded = true
-    }
     promise.value = p
 
     p.catch((reason: FirestoreError) => {
@@ -136,6 +140,12 @@ export function _useFirestoreRef(
     bindFirestoreRef()
   }
 
+  // only add the first promise to the pending ones
+  // TODO: can we make this tree shakeable?
+  if (initialSourceValue) {
+    removePendingPromise = addPendingPromise(promise.value, initialSourceValue)
+  }
+
   // TODO: SSR serialize the values for Nuxt to expose them later and use them
   // as initial values while specifying a wait: true to only swap objects once
   // Firebase has done its initial sync. Also, on server, you don't need to
@@ -145,9 +155,11 @@ export function _useFirestoreRef(
   // TODO: warn else
   if (hasCurrentScope) {
     onScopeDispose(unbind)
-    // wait for the promise during SSR
-    // TODO: configurable
-    onServerPrefetch(() => promise.value)
+    if (getCurrentInstance()) {
+      // wait for the promise during SSR
+      // TODO: configurable
+      onServerPrefetch(() => promise.value)
+    }
   }
 
   // TODO: rename to stop

src/index.ts

@@ -51,6 +51,7 @@ export { useCurrentUser, VueFireAuth, useFirebaseAuth } from './auth'
 
 // SSR
 export { usePendingPromises } from './ssr/plugin'
+export { useSSRInitialState } from './ssr/initialState'
 
 // App Check
 export { VueFireAppCheck, useAppCheckToken } from './app-check'

src/shared.ts

@@ -11,6 +11,8 @@ import type { Ref, ShallowRef } from 'vue-demi'
 
 export const noop = () => {}
 
+export const isClient = typeof window !== 'undefined'
+
 // FIXME: replace any with unknown or T generics
 
 export interface OperationsType {
@@ -124,21 +126,9 @@ export function isFirestoreDataReference<T = unknown>(
   return isDocumentRef(source) || isCollectionRef(source)
 }
 
-// The Firestore SDK has an undocumented _query
-// object that has a method to generate a hash for a query,
-// which we need for useObservable
-// https://github.com/firebase/firebase-js-sdk/blob/5beb23cd47312ffc415d3ce2ae309cc3a3fde39f/packages/firestore/src/core/query.ts#L221
-// @internal
-export interface _FirestoreQueryWithId<T = DocumentData>
-  extends FirestoreQuery<T> {
-  _query: {
-    canonicalId(): string
-  }
-}
-
 export function isFirestoreQuery(
   source: unknown
-): source is _FirestoreQueryWithId<unknown> {
+): source is FirestoreQuery<unknown> {
   return isObject(source) && source.type === 'query'
 }
 

src/ssr/initialState.ts

@@ -0,0 +1,87 @@
+import { FirebaseApp } from 'firebase/app'
+import {
+  CollectionReference,
+  DocumentReference,
+  Query as FirestoreQuery,
+} from 'firebase/firestore'
+import { InjectionKey } from 'vue'
+import { useFirebaseApp } from '../app'
+import { isFirestoreQuery, _Nullable } from '../shared'
+
+export interface SSRStore {
+  // firestore data
+  f: Record<string, unknown>
+  // rtdb data
+  r: Record<string, unknown>
+
+  // storage
+  s: Record<string, unknown>
+
+  // auth data
+  u: Record<string, unknown>
+}
+
+// @internal
+const initialStatesMap = new WeakMap<FirebaseApp, SSRStore>()
+
+/**
+ * Allows getting the initial state set during SSR on the client.
+ *
+ * @param initialState - the initial state to set for the firebase app during SSR. Pass undefined to not set it
+ * @param firebaseApp - the firebase app to get the initial state for
+ * @returns the initial states for the current firebaseApp
+ */
+export function useSSRInitialState(
+  initialState?: SSRStore,
+  firebaseApp: FirebaseApp = useFirebaseApp()
+): SSRStore {
+  // get initial state based on the current firebase app
+  if (!initialStatesMap.has(firebaseApp)) {
+    initialStatesMap.set(
+      firebaseApp,
+      initialState || { f: {}, r: {}, s: {}, u: {} }
+    )
+  }
+
+  return initialStatesMap.get(firebaseApp)!
+}
+
+type FirestoreDataSource =
+  | DocumentReference<unknown>
+  | CollectionReference<unknown>
+  | FirestoreQuery<unknown>
+
+export function getInitialValue(
+  type: 'f' | 'r',
+  ssrKey?: string | undefined,
+  dataSource?: _Nullable<FirestoreDataSource>
+) {
+  const initialState: Record<string, unknown> = useSSRInitialState()[type] || {}
+  const key = ssrKey || getFirestoreSourceKey(dataSource)
+
+  // TODO: warn for queries on the client if there are other keys and this is during hydration
+
+  // returns undefined if no key, otherwise initial state or undefined
+  // undefined should be treated as no initial state
+  return key && initialState[key]
+}
+
+export function setInitialValue(
+  type: 'f' | 'r',
+  value: unknown,
+  ssrKey?: string | undefined,
+  dataSource?: _Nullable<FirestoreDataSource>
+) {
+  const initialState: Record<string, unknown> = useSSRInitialState()[type]
+  const key = ssrKey || getFirestoreSourceKey(dataSource)
+
+  if (key) {
+    initialState[key] = value
+  }
+}
+
+function getFirestoreSourceKey(
+  source: _Nullable<FirestoreDataSource>
+): string | undefined {
+  return !source || isFirestoreQuery(source) ? undefined : source.path
+}

src/ssr/plugin.ts

@@ -7,6 +7,7 @@ import {
 } from 'firebase/firestore'
 import { useFirebaseApp, _FirebaseAppInjectionKey } from '../app'
 import { getDataSourcePath, noop } from '../shared'
+import { setInitialValue } from './initialState'
 
 export const appPendingPromises = new WeakMap<
   FirebaseApp,
@@ -33,17 +34,23 @@ export function addPendingPromise(
   }
   const pendingPromises = appPendingPromises.get(app)!
 
-  ssrKey = getDataSourcePath(dataSource)
-  if (ssrKey) {
-    pendingPromises.set(ssrKey, promise)
+  const key = ssrKey || getDataSourcePath(dataSource)
+  if (key) {
+    pendingPromises.set(key, promise)
+
+    // TODO: skip this outside of SSR
+    promise.then((value) => {
+      // TODO: figure out 'f', probably based on the type of dataSource
+      setInitialValue('f', value, key! /* dataSource */)
+    })
   } else {
     // TODO: warn if in SSR context in other contexts than vite
     if (process.env.NODE_ENV !== 'production' /* && import.meta.env?.SSR */) {
       console.warn('[VueFire]: Could not get the path of the data source')
     }
   }
 
-  return ssrKey ? () => pendingPromises.delete(ssrKey!) : noop
+  return key ? () => pendingPromises.delete(key!) : noop
 }
 
 /**