fix(presentation): support references in location resolver selections (…

…#1652)

* fix(presentation): support references in location resolver selections

* fix(presentation): remove unused hook

packages/presentation/package.json

@@ -104,6 +104,7 @@
     "@types/lodash.isequal": "^4.5.8",
     "fast-deep-equal": "3.1.3",
     "framer-motion": "11.0.8",
+    "lodash.get": "^4.4.2",
     "lodash.isequal": "^4.5.0",
     "mendoza": "3.0.7",
     "mnemonist": "0.39.8",
@@ -117,6 +118,7 @@
     "@repo/visual-editing-helpers": "0.6.17",
     "@sanity/client": "^6.19.1",
     "@sanity/pkg-utils": "6.9.3",
+    "@types/lodash.get": "^4.4.9",
     "happy-dom": "^14.12.0",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",

packages/presentation/src/internals.ts

@@ -14,9 +14,11 @@ export {
   getPreviewValueWithFallback,
   getPublishedId,
   isRecord,
+  isReference,
   type Path,
   pathToString,
   Preview,
+  type Previewable,
   PreviewCard,
   type PreviewValue,
   type PublishedId,

packages/presentation/src/useDocumentLocations.ts

@@ -1,17 +1,144 @@
+import get from 'lodash.get'
 import {useEffect, useMemo, useState} from 'react'
-import {isObservable, map, type Observable, of} from 'rxjs'
-import {type SanityDocument} from 'sanity'
+import {isObservable, map, Observable, of, switchMap} from 'rxjs'
 
-import {useDocumentStore} from './internals'
+import {
+  type DocumentStore,
+  isRecord,
+  isReference,
+  type Previewable,
+  type SanityDocument,
+  useDocumentStore,
+} from './internals'
 import type {
   DocumentLocationResolver,
+  DocumentLocationResolverObject,
   DocumentLocationResolvers,
   DocumentLocationsState,
   DocumentLocationsStatus,
 } from './types'
+import {props} from './util/props'
 
 const INITIAL_STATE: DocumentLocationsState = {locations: []}
 
+function getDocumentId(value: Previewable) {
+  if (isReference(value)) {
+    return value._ref
+  }
+  return '_id' in value ? value._id : undefined
+}
+
+function cleanPreviewable(id: string | undefined, previewable: Previewable) {
+  const clean: Record<string, unknown> = id ? {...previewable, _id: id} : {...previewable}
+
+  if (clean['_type'] === 'reference') {
+    delete clean['_type']
+    delete clean['_ref']
+    delete clean['_weak']
+    delete clean['_dataset']
+    delete clean['_projectId']
+    delete clean['_strengthenOnPublish']
+  }
+
+  return clean
+}
+
+function listen(id: string, fields: string[], store: DocumentStore) {
+  const projection = fields.join(', ')
+  const query = `*[_id==$id][0]{${projection}}`
+  const params = {id}
+  return store.listenQuery(query, params, {
+    perspective: 'previewDrafts',
+  }) as Observable<SanityDocument | null>
+}
+
+function observeDocument(
+  value: Previewable | null,
+  paths: string[][],
+  store: DocumentStore,
+): Observable<Record<string, unknown> | null> {
+  if (!value || typeof value !== 'object') {
+    return of(value)
+  }
+
+  const id = getDocumentId(value)
+  const currentValue = cleanPreviewable(id, value)
+
+  const headlessPaths = paths.filter((path) => !(path[0] in currentValue))
+
+  if (id && headlessPaths.length) {
+    const fields = [...new Set(headlessPaths.map((path: string[]) => path[0]))]
+    return listen(id, fields, store).pipe(
+      switchMap((snapshot) => {
+        if (snapshot) {
+          return observeDocument(snapshot, paths, store)
+        }
+        return of(null)
+      }),
+    )
+  }
+
+  const leads: Record<string, string[][]> = {}
+  paths.forEach((path) => {
+    const [head, ...tail] = path
+    if (!leads[head]) {
+      leads[head] = []
+    }
+    leads[head].push(tail)
+  })
+  const next = Object.keys(leads).reduce((res: Record<string, unknown>, head) => {
+    const tails = leads[head].filter((tail) => tail.length > 0)
+    if (tails.length === 0) {
+      res[head] = isRecord(value) ? (value as Record<string, unknown>)[head] : undefined
+    } else {
+      res[head] = observeDocument((value as any)[head], tails, store)
+    }
+    return res
+  }, currentValue)
+
+  return of(next).pipe(props({wait: true}))
+}
+
+function observeForLocations(
+  doc: {id: string; type: string},
+  resolver:
+    | DocumentLocationsState
+    | DocumentLocationResolver
+    | DocumentLocationResolverObject<string>
+    | undefined,
+  documentStore: DocumentStore,
+) {
+  if (!resolver) return of(undefined)
+  const {id, type} = doc
+  // Original/advanced resolver which requires explicit use of Observables
+  if (typeof resolver === 'function') {
+    const params = {id, type}
+
+    const context = {documentStore}
+    const _result = resolver(params, context)
+    return isObservable(_result) ? _result : of(_result)
+  }
+
+  // Simplified resolver pattern which abstracts away Observable logic
+  if ('select' in resolver && 'resolve' in resolver) {
+    const {select} = resolver
+    const paths = Object.values(select).map((value) => String(value).split('.')) || []
+    const doc = {_type: 'reference', _ref: id}
+    return observeDocument(doc, paths, documentStore).pipe(
+      map((doc) => {
+        return Object.keys(select).reduce<Record<string, unknown>>((acc, key) => {
+          acc[key] = get(doc, select[key])
+          return acc
+        }, {})
+      }),
+      map(resolver.resolve),
+    )
+  }
+
+  // Resolver is explicitly provided state
+  return of(resolver)
+}
+
 export function useDocumentLocations(props: {
   id: string
   resolvers?: DocumentLocationResolver | DocumentLocationResolvers
@@ -43,15 +170,7 @@ export function useDocumentLocations(props: {
 
     // Simplified resolver pattern which abstracts away Observable logic
     if ('select' in resolver && 'resolve' in resolver) {
-      const projection = Object.entries(resolver.select)
-        .map(([key, value]) => `"${key}": ${value}`)
-        .join(', ')
-      const query = `*[_id==$id][0]{${projection}}`
-      const params = {id}
-      const doc$ = documentStore.listenQuery(query, params, {
-        perspective: 'previewDrafts',
-      }) as Observable<SanityDocument | null>
-      return doc$.pipe(map(resolver.resolve))
+      return observeForLocations({id, type}, resolver, documentStore)
     }
 
     // Resolver is explicitly provided state

packages/presentation/src/useMainDocument.ts

@@ -5,7 +5,6 @@ import {useClient} from 'sanity'
 import {useRouter} from 'sanity/router'
 
 import {API_VERSION} from './constants'
-import {useDocumentStore} from './internals'
 import type {
   DocumentResolver,
   DocumentResolverContext,
@@ -96,7 +95,6 @@ export function useMainDocument(props: {
   const {navigate, resolvers = [], path, previewUrl} = props
 
   const {state: routerState} = useRouter()
-  const documentStore = useDocumentStore()
   const client = useClient({apiVersion: API_VERSION})
 
   const [mainDocumentState, setMainDocumentState] = useState<MainDocumentState | undefined>(
@@ -179,7 +177,7 @@ export function useMainDocument(props: {
     }
     clearState()
     return undefined
-  }, [client, clearState, documentStore, navigate, resolvers, url])
+  }, [client, clearState, navigate, resolvers, url])
 
   return mainDocumentState
 }

packages/presentation/src/util/props.ts

@@ -0,0 +1,48 @@
+import {
+  combineLatest,
+  from,
+  isObservable,
+  map,
+  mergeAll,
+  Observable,
+  of,
+  scan,
+  switchMap,
+} from 'rxjs'
+
+type Props<K extends keyof any, T> = {
+  [P in K]: T | Observable<T>
+}
+
+function keysOf<T extends object>(value: T) {
+  return Object.keys(value) as (keyof T)[]
+}
+
+function setKey(source: Record<string, unknown>, key: any, value: unknown) {
+  return {
+    ...source,
+    [key]: value,
+  }
+}
+
+export function props<K extends keyof any, T>(options: {wait?: boolean} = {}) {
+  return (source: Observable<Props<K, T>>): Observable<Record<string, unknown>> => {
+    return new Observable<Props<K, T>>((observer) => source.subscribe(observer)).pipe(
+      switchMap((object) => {
+        const keyObservables = keysOf(object).map((key) => {
+          const value = object[key]
+          return isObservable(value) ? from(value).pipe(map((val) => [key, val])) : of([key, value])
+        })
+
+        return options.wait
+          ? combineLatest(keyObservables).pipe(
+              map((pairs) => pairs.reduce((acc, [key, value]) => setKey(acc, key, value), {})),
+            )
+          : from(keyObservables).pipe(
+              mergeAll(),
+              scan((acc, [key, value]) => setKey(acc, key, value), {}),
+            )
+      }),
+    )
+  }
+}