add set SET_COLLECTION support to onyx.update

Expensify · Nov 22, 2024 · 5091684 · 5091684
1 parent c2fb5c5
commit 5091684
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 19 deletions.
diff --git a/lib/Onyx.ts b/lib/Onyx.ts
@@ -27,6 +27,7 @@ import type {
     OnyxValue,
     OnyxInput,
     OnyxMethodMap,
+    OnyxMethod,
 } from './types';
 import OnyxUtils from './OnyxUtils';
 import logMessages from './logMessages';
@@ -600,7 +601,7 @@ function updateSnapshots(data: OnyxUpdate[]) {
 function update(data: OnyxUpdate[]): Promise<void> {
     // First, validate the Onyx object is in the format we expect
     data.forEach(({onyxMethod, key, value}) => {
-        if (![OnyxUtils.METHOD.CLEAR, OnyxUtils.METHOD.SET, OnyxUtils.METHOD.MERGE, OnyxUtils.METHOD.MERGE_COLLECTION, OnyxUtils.METHOD.MULTI_SET].includes(onyxMethod)) {
+        if (!Object.values(OnyxUtils.METHOD).includes(onyxMethod)) {
             throw new Error(`Invalid onyxMethod ${onyxMethod} in Onyx update.`);
         }
         if (onyxMethod === OnyxUtils.METHOD.MULTI_SET) {
@@ -633,30 +634,32 @@ function update(data: OnyxUpdate[]): Promise<void> {
         }
     };
 
+    const handleCollectionOperation = (operation: typeof enqueueSetOperation | typeof enqueueMergeOperation, onyxMethod: OnyxMethod) => {
+        return (key: OnyxKey, value: OnyxInput<OnyxKey>) => {
+            const collection = value as Collection<CollectionKey, unknown, unknown>;
+            if (!OnyxUtils.isValidNonEmptyCollectionForMerge(collection)) {
+                Logger.logInfo(`${onyxMethod} enqueued within update() with invalid or empty value. Skipping this operation.`);
+                return;
+            }
+
+            const collectionKeys = Object.keys(collection);
+            if (OnyxUtils.doAllCollectionItemsBelongToSameParent(key, collectionKeys)) {
+                const mergedCollection: OnyxInputKeyValueMapping = collection;
+                collectionKeys.forEach((collectionKey) => operation(collectionKey, mergedCollection[collectionKey]));
+            }
+        };
+    };
+
     const promises: Array<() => Promise<void>> = [];
     let clearPromise: Promise<void> = Promise.resolve();
 
     data.forEach(({onyxMethod, key, value}) => {
         const handlers: Record<OnyxMethodMap[keyof OnyxMethodMap], (k: typeof key, v: typeof value) => void> = {
             [OnyxUtils.METHOD.SET]: enqueueSetOperation,
             [OnyxUtils.METHOD.MERGE]: enqueueMergeOperation,
-            [OnyxUtils.METHOD.MERGE_COLLECTION]: (k, v) => {
-                const collection = v as Collection<CollectionKey, unknown, unknown>;
-                if (!OnyxUtils.isValidNonEmptyCollectionForMerge(collection)) {
-                    Logger.logInfo('mergeCollection enqueued within update() with invalid or empty value. Skipping this operation.');
-                    return;
-                }
-
-                // Confirm all the collection keys belong to the same parent
-                const collectionKeys = Object.keys(collection);
-                if (OnyxUtils.doAllCollectionItemsBelongToSameParent(k, collectionKeys)) {
-                    const mergedCollection: OnyxInputKeyValueMapping = collection;
-                    collectionKeys.forEach((collectionKey) => enqueueMergeOperation(collectionKey, mergedCollection[collectionKey]));
-                }
-            },
-            [OnyxUtils.METHOD.MULTI_SET]: (k, v) => {
-                Object.entries(v as Partial<OnyxInputKeyValueMapping>).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue));
-            },
+            [OnyxUtils.METHOD.MERGE_COLLECTION]: handleCollectionOperation(enqueueMergeOperation, onyxMethod),
+            [OnyxUtils.METHOD.SET_COLLECTION]: handleCollectionOperation(enqueueSetOperation, onyxMethod),
+            [OnyxUtils.METHOD.MULTI_SET]: (k, v) => Object.entries(v as Partial<OnyxInputKeyValueMapping>).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue)),
             [OnyxUtils.METHOD.CLEAR]: () => {
                 clearPromise = clear();
             },
@@ -703,7 +706,12 @@ function update(data: OnyxUpdate[]): Promise<void> {
             promises.push(() => mergeCollection(collectionKey, batchedCollectionUpdates.merge as Collection<CollectionKey, unknown, unknown>));
         }
         if (!utils.isEmptyObject(batchedCollectionUpdates.set)) {
-            promises.push(() => multiSet(batchedCollectionUpdates.set));
+            // Check if this is from a SET_COLLECTION operation
+            if (data.some((onyxUpdate) => onyxUpdate.onyxMethod === OnyxUtils.METHOD.SET_COLLECTION && onyxUpdate.key === collectionKey)) {
+                promises.push(() => setCollection(collectionKey, batchedCollectionUpdates.set as Collection<CollectionKey, unknown, unknown>));
+            } else {
+                promises.push(() => multiSet(batchedCollectionUpdates.set));
+            }
         }
     });
 

diff --git a/lib/types.ts b/lib/types.ts
@@ -420,6 +420,10 @@ type OnyxMethodValueMap = {
         key: CollectionKeyBase;
         value: OnyxMergeCollectionInput<CollectionKeyBase>;
     };
+    [OnyxUtils.METHOD.SET_COLLECTION]: {
+        key: CollectionKeyBase;
+        value: OnyxMergeCollectionInput<CollectionKeyBase>;
+    };
 };
 
 /**