feat: accept function in setEntryData

posva · Dec 25, 2023 · 2abb7c0 · 2abb7c0
1 parent 0ee5c8a
commit 2abb7c0
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/src/data-fetching-store.ts b/src/data-fetching-store.ts
@@ -219,15 +219,23 @@ export const useDataFetchingStore = defineStore('PiniaColada', () => {
     }
   }
 
-  function setEntryData<TResult = unknown>(key: UseQueryKey, data: TResult) {
+  function setEntryData<TResult = unknown>(
+    key: UseQueryKey,
+    data: TResult | ((data: Ref<TResult | undefined>) => void)
+  ) {
     const entry = entryStateRegistry.get(key) as
       | UseQueryStateEntry<TResult>
       | undefined
     if (!entry) {
       return
     }
 
-    entry.data.value = data
+    if (typeof data === 'function') {
+      // the remaining type is TResult & Fn, so we need a cast
+      ;(data as (data: Ref<TResult | undefined>) => void)(entry.data)
+    } else {
+      entry.data.value = data
+    }
     // TODO: complete and test
     entry.error.value = null
   }