refactor(store) simplify key value store, add SWR through local stora…

…ge behavior

packages/sanity/src/core/store/key-value/backends/localStorage.ts

@@ -1,8 +1,8 @@
-import {type Observable, of as observableOf} from 'rxjs'
+import {supportsLocalStorage} from '../../../util/supportsLocalStorage'
+import {type KeyValueStoreValue} from '../types'
+import {createMemoryStorage} from './memoryStorage'
 
-import {type Backend, type KeyValuePair} from './types'
-
-const tryParse = (val: string) => {
+function tryParse(val: string) {
   try {
     return JSON.parse(val)
   } catch (err) {
@@ -12,42 +12,27 @@ const tryParse = (val: string) => {
   }
 }
 
-const getKey = (key: string): Observable<unknown> => {
-  const val = localStorage.getItem(key)
-
-  return observableOf(val === null ? null : tryParse(val))
-}
-
-const setKey = (key: string, nextValue: unknown): Observable<unknown> => {
-  // Can't stringify undefined, and nulls are what
-  // `getItem` returns when key does not exist
-  if (typeof nextValue === 'undefined' || nextValue === null) {
-    localStorage.removeItem(key)
-  } else {
-    localStorage.setItem(key, JSON.stringify(nextValue))
+export function createLocalStorageStore() {
+  if (!supportsLocalStorage) {
+    return createMemoryStorage()
   }
 
-  return observableOf(nextValue)
-}
-
-const getKeys = (keys: string[]): Observable<unknown[]> => {
-  const values = keys.map((key, i) => {
+  function getKey(key: string): KeyValueStoreValue | null {
     const val = localStorage.getItem(key)
-    return val === null ? null : tryParse(val)
-  })
 
-  return observableOf(values)
-}
+    return val === null ? null : tryParse(val)
+  }
 
-const setKeys = (keyValuePairs: KeyValuePair[]): Observable<unknown[]> => {
-  keyValuePairs.forEach((pair) => {
-    if (pair.value === undefined || pair.value === null) {
-      localStorage.removeItem(pair.key)
+  const setKey = function (key: string, nextValue: KeyValueStoreValue) {
+    // Can't stringify undefined, and nulls are what
+    // `getItem` returns when key does not exist
+    if (typeof nextValue === 'undefined' || nextValue === null) {
+      localStorage.removeItem(key)
     } else {
-      localStorage.setItem(pair.key, JSON.stringify(pair.value))
+      localStorage.setItem(key, JSON.stringify(nextValue))
     }
-  })
-  return observableOf(keyValuePairs.map((pair) => pair.value))
+  }
+  return {getKey, setKey}
 }
 
-export const localStorageBackend: Backend = {getKey, setKey, getKeys, setKeys}
+export const localStore = createLocalStorageStore()

packages/sanity/src/core/store/key-value/backends/memoryStorage.ts

@@ -0,0 +1,13 @@
+import {type KeyValueStoreValue} from '../types'
+
+export function createMemoryStorage() {
+  const DB = Object.create(null)
+  return {
+    getKey(key: string): KeyValueStoreValue | null {
+      return DB[key] || null
+    },
+    setKey(key: string, value: KeyValueStoreValue) {
+      DB[key] = value
+    },
+  }
+}

packages/sanity/src/core/store/key-value/backends/server.ts → packages/sanity/src/core/store/key-value/backends/serverStorage.ts

@@ -1,22 +1,20 @@
 import {type SanityClient} from '@sanity/client'
 import DataLoader from 'dataloader'
-import {catchError, from, map, of} from 'rxjs'
 
 import {DEFAULT_STUDIO_CLIENT_OPTIONS} from '../../../studioClient'
 import {type KeyValueStoreValue} from '../types'
-import {type Backend, type KeyValuePair} from './types'
+import {type KeyValuePair, type ServerStorage} from './types'
 
 /** @internal */
-export interface ServerBackendOptions {
+export interface ServerStorageOptions {
   client: SanityClient
 }
 
 /**
- * One of serveral possible backends for KeyValueStore. This backend uses the
- * Sanity client to store and retrieve key-value pairs from the /users/me/keyvalue endpoint.
+ * Backend uses the Sanity client to store and retrieve key-value pairs from the /users/me/keyvalue endpoint.
  * @internal
  */
-export function serverBackend({client: _client}: ServerBackendOptions): Backend {
+export function createServerStorage({client: _client}: ServerStorageOptions): ServerStorage {
   const client = _client.withConfig(DEFAULT_STUDIO_CLIENT_OPTIONS)
 
   const keyValueLoader = new DataLoader<string, KeyValueStoreValue | null>(async (keys) => {
@@ -40,50 +38,37 @@ export function serverBackend({client: _client}: ServerBackendOptions): Backend
       {} as Record<string, KeyValueStoreValue | null>,
     )
 
-    const result = keys.map((key) => keyValuePairs[key] || null)
-    return result
+    return keys.map((key) => keyValuePairs[key] || null)
   })
 
-  const getKeys = (keys: string[]) => {
-    return from(keyValueLoader.loadMany(keys))
+  const getKey = (key: string) => {
+    return keyValueLoader.load(key)
   }
 
-  const setKeys = (keyValuePairs: KeyValuePair[]) => {
-    return from(
-      client.request<KeyValuePair[]>({
+  const setKey = (key: string, nextValue: unknown) => {
+    return client
+      .request<KeyValuePair[]>({
         method: 'PUT',
         uri: `/users/me/keyvalue`,
-        body: keyValuePairs,
+        body: [{key, value: nextValue}],
         withCredentials: true,
-      }),
-    ).pipe(
-      map((response) => {
-        return response.map((pair) => {
+      })
+      .then(
+        (response) => {
+          const pair = response[0]
           keyValueLoader.clear(pair.key)
           keyValueLoader.prime(pair.key, pair.value)
-
           return pair.value
-        })
-      }),
-      catchError((error) => {
-        console.error('Error setting data:', error)
-        return of(Array(keyValuePairs.length).fill(null))
-      }),
-    )
-  }
-
-  const getKey = (key: string) => {
-    return getKeys([key]).pipe(map((values) => values[0]))
-  }
-
-  const setKey = (key: string, nextValue: unknown) => {
-    return setKeys([{key, value: nextValue as KeyValueStoreValue}]).pipe(map((values) => values[0]))
+        },
+        (error) => {
+          console.error('Error setting data:', error)
+          return null
+        },
+      )
   }
 
   return {
     getKey,
     setKey,
-    getKeys,
-    setKeys,
   }
 }

packages/sanity/src/core/store/key-value/backends/types.ts

@@ -1,15 +1,11 @@
-import {type Observable} from 'rxjs'
-
 import {type KeyValueStoreValue} from '../types'
 
 export interface KeyValuePair {
   key: string
   value: KeyValueStoreValue | null
 }
 
-export interface Backend {
-  getKey: (key: string) => Observable<unknown>
-  setKey: (key: string, nextValue: unknown) => Observable<unknown>
-  getKeys: (keys: string[]) => Observable<unknown[]>
-  setKeys: (keyValuePairs: KeyValuePair[]) => Observable<unknown[]>
+export interface ServerStorage {
+  getKey: (key: string) => Promise<KeyValueStoreValue | null>
+  setKey: (key: string, nextValue: unknown) => Promise<KeyValueStoreValue>
 }

packages/sanity/src/core/store/key-value/index.ts

@@ -1,2 +1,2 @@
-export * from './KeyValueStore'
+export * from './keyValueStore'
 export * from './types'

packages/sanity/src/core/store/key-value/keyValueStore.ts

@@ -0,0 +1,9 @@
+import {type SanityClient} from '@sanity/client'
+
+import {withLocalStorageSWR} from './localStorageSWR'
+import {createServerKeyValueStore} from './serverKeyValueStore'
+
+/** @internal */
+export function createKeyValueStore(options: {client: SanityClient}) {
+  return withLocalStorageSWR(createServerKeyValueStore(options))
+}

packages/sanity/src/core/store/key-value/localStorageSWR.ts

@@ -0,0 +1,39 @@
+import {isEqual} from 'lodash'
+import {fromEvent, merge, NEVER} from 'rxjs'
+import {distinctUntilChanged, filter, map, tap} from 'rxjs/operators'
+
+import {localStore} from './backends/localStorage'
+import {type KeyValueStore, type KeyValueStoreValue} from './types'
+
+// Whether or not to enable instant user sync between tabs
+// if set to true, the setting will update instantly across all tabs
+const ENABLE_CROSS_TAB_SYNC = false
+
+/**
+ * Wraps a KeyValueStore and adds Stale-While-Revalidate (SWR) behavior to it
+ */
+export function withLocalStorageSWR(wrappedStore: KeyValueStore): KeyValueStore {
+  const storageEvent = ENABLE_CROSS_TAB_SYNC ? fromEvent<StorageEvent>(window, 'storage') : NEVER
+
+  function getKey(key: string) {
+    const lsUpdates = storageEvent.pipe(
+      filter((event) => event.key === key),
+      map(() => localStore.getKey(key)),
+    )
+
+    return merge(lsUpdates, wrappedStore.getKey(key)).pipe(
+      distinctUntilChanged(isEqual),
+      tap((value) => {
+        localStore.setKey(key, value)
+      }),
+    )
+  }
+  function setKey(key: string, nextValue: KeyValueStoreValue) {
+    localStore.setKey(key, nextValue)
+    return wrappedStore.setKey(key, nextValue)
+  }
+  return {
+    getKey,
+    setKey,
+  }
+}

packages/sanity/src/core/store/key-value/serverKeyValueStore.ts

@@ -0,0 +1,49 @@
+import {type SanityClient} from '@sanity/client'
+import {isEqual} from 'lodash'
+import {concat, type Observable, Subject} from 'rxjs'
+import {distinctUntilChanged, filter, map} from 'rxjs/operators'
+
+import {createServerStorage} from './backends/serverStorage'
+import {type KeyValueStore, type KeyValueStoreValue} from './types'
+
+export function createServerKeyValueStore({client}: {client: SanityClient}): KeyValueStore {
+  const serverStorage = createServerStorage({client})
+
+  const events$ = new Subject<{
+    type: 'optimistic' | 'commit'
+    key: string
+    value: KeyValueStoreValue
+  }>()
+
+  function getKey(key: string) {
+    return serverStorage.getKey(key)
+  }
+
+  function setKey(key: string, value: KeyValueStoreValue) {
+    events$.next({type: 'optimistic', key, value})
+
+    /*
+     * The backend returns the result of the set operation, so we can just pass that along.
+     * Most utils do not use it (they will take advantage of local state first) but it reflects the
+     * backend function and could be useful for debugging.
+     */
+    return serverStorage.setKey(key, value).then((storedValue) => {
+      events$.next({type: 'commit', key, value: storedValue})
+      return storedValue
+    })
+  }
+
+  return {
+    getKey(key: string): Observable<KeyValueStoreValue | null> {
+      return concat(
+        getKey(key),
+        events$.pipe(
+          filter((event) => event.key === key),
+          map((event) => event.value),
+          distinctUntilChanged(isEqual),
+        ),
+      )
+    },
+    setKey,
+  }
+}

packages/sanity/src/core/store/key-value/types.ts

@@ -12,5 +12,5 @@ export type KeyValueStoreValue = JsonPrimitive | JsonObject | JsonArray
 /** @internal */
 export interface KeyValueStore {
   getKey(key: string): Observable<KeyValueStoreValue | null>
-  setKey(key: string, value: KeyValueStoreValue): Observable<KeyValueStoreValue>
+  setKey(key: string, value: KeyValueStoreValue): Promise<KeyValueStoreValue>
 }