Improve import onyx performance

src/components/ImportOnyxState/index.native.tsx

@@ -1,18 +1,17 @@
 import React, {useState} from 'react';
 import ReactNativeBlobUtil from 'react-native-blob-util';
-import Onyx, {useOnyx} from 'react-native-onyx';
+import {useOnyx} from 'react-native-onyx';
 import type {FileObject} from '@components/AttachmentModal';
-import {KEYS_TO_PRESERVE, setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
+import {setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
 import {setShouldForceOffline} from '@libs/actions/Network';
+import {rollbackOngoingRequest} from '@libs/actions/PersistedRequests';
 import Navigation from '@libs/Navigation/Navigation';
 import type {OnyxValues} from '@src/ONYXKEYS';
 import ONYXKEYS from '@src/ONYXKEYS';
 import ROUTES from '@src/ROUTES';
 import BaseImportOnyxState from './BaseImportOnyxState';
 import type ImportOnyxStateProps from './types';
-import {cleanAndTransformState} from './utils';
-
-const CHUNK_SIZE = 100;
+import {cleanAndTransformState, processStateImport} from './utils';
 
 function readOnyxFile(fileUri: string) {
     const filePath = decodeURIComponent(fileUri.replace('file://', ''));
@@ -25,27 +24,6 @@ function readOnyxFile(fileUri: string) {
     });
 }
 
-function chunkArray<T>(array: T[], size: number): T[][] {
-    const result = [];
-    for (let i = 0; i < array.length; i += size) {
-        result.push(array.slice(i, i + size));
-    }
-    return result;
-}
-
-function applyStateInChunks(state: OnyxValues) {
-    const entries = Object.entries(state);
-    const chunks = chunkArray(entries, CHUNK_SIZE);
-
-    let promise = Promise.resolve();
-    chunks.forEach((chunk) => {
-        const partialOnyxState = Object.fromEntries(chunk) as Partial<OnyxValues>;
-        promise = promise.then(() => Onyx.multiSet(partialOnyxState));
-    });
-
-    return promise;
-}
-
 export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
     const [isErrorModalVisible, setIsErrorModalVisible] = useState(false);
     const [session] = useOnyx(ONYXKEYS.SESSION);
@@ -55,22 +33,25 @@ export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
             return;
         }
 
+        rollbackOngoingRequest();
         setIsLoading(true);
         readOnyxFile(file.uri)
             .then((fileContent: string) => {
                 const transformedState = cleanAndTransformState<OnyxValues>(fileContent);
                 const currentUserSessionCopy = {...session};
                 setPreservedUserSession(currentUserSessionCopy);
                 setShouldForceOffline(true);
-                Onyx.clear(KEYS_TO_PRESERVE).then(() => {
-                    applyStateInChunks(transformedState).then(() => {
-                        setIsUsingImportedState(true);
-                        Navigation.navigate(ROUTES.HOME);
-                    });
-                });
+                return processStateImport(transformedState);
+            })
+            .then(() => {
+                setIsUsingImportedState(true);
+                Navigation.navigate(ROUTES.HOME);
             })
             .catch(() => {
                 setIsErrorModalVisible(true);
+            })
+            .finally(() => {
+                setIsLoading(false);
             });
     };
 

src/components/ImportOnyxState/index.tsx

@@ -1,15 +1,16 @@
 import React, {useState} from 'react';
-import Onyx, {useOnyx} from 'react-native-onyx';
+import {useOnyx} from 'react-native-onyx';
 import type {FileObject} from '@components/AttachmentModal';
-import {KEYS_TO_PRESERVE, setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
+import {setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
 import {setShouldForceOffline} from '@libs/actions/Network';
+import {rollbackOngoingRequest} from '@libs/actions/PersistedRequests';
 import Navigation from '@libs/Navigation/Navigation';
 import type {OnyxValues} from '@src/ONYXKEYS';
 import ONYXKEYS from '@src/ONYXKEYS';
 import ROUTES from '@src/ROUTES';
 import BaseImportOnyxState from './BaseImportOnyxState';
 import type ImportOnyxStateProps from './types';
-import {cleanAndTransformState} from './utils';
+import {cleanAndTransformState, processStateImport} from './utils';
 
 export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
     const [isErrorModalVisible, setIsErrorModalVisible] = useState(false);
@@ -20,26 +21,31 @@ export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
             return;
         }
 
+        rollbackOngoingRequest();
         setIsLoading(true);
+
         const blob = new Blob([file as BlobPart]);
         const response = new Response(blob);
 
         response
             .text()
             .then((text) => {
                 const fileContent = text;
+
                 const transformedState = cleanAndTransformState<OnyxValues>(fileContent);
+
                 const currentUserSessionCopy = {...session};
                 setPreservedUserSession(currentUserSessionCopy);
                 setShouldForceOffline(true);
-                Onyx.clear(KEYS_TO_PRESERVE).then(() => {
-                    Onyx.multiSet(transformedState).then(() => {
-                        setIsUsingImportedState(true);
-                        Navigation.navigate(ROUTES.HOME);
-                    });
-                });
+
+                return processStateImport(transformedState);
+            })
+            .then(() => {
+                setIsUsingImportedState(true);
+                Navigation.navigate(ROUTES.HOME);
             })
-            .catch(() => {
+            .catch((error) => {
+                console.error('Error importing state:', error);
                 setIsErrorModalVisible(true);
                 setIsLoading(false);
             });

src/components/ImportOnyxState/utils.ts

@@ -1,5 +1,9 @@
 import cloneDeep from 'lodash/cloneDeep';
+import type {OnyxCollection, OnyxEntry, OnyxKey} from 'react-native-onyx';
+import Onyx from 'react-native-onyx';
 import type {UnknownRecord} from 'type-fest';
+import {KEYS_TO_PRESERVE} from '@libs/actions/App';
+import type {OnyxCollectionValuesMapping, OnyxValues} from '@src/ONYXKEYS';
 import ONYXKEYS from '@src/ONYXKEYS';
 
 // List of Onyx keys from the .txt file we want to keep for the local override
@@ -12,7 +16,7 @@ function isRecord(value: unknown): value is Record<string, unknown> {
 function transformNumericKeysToArray(data: UnknownRecord): UnknownRecord | unknown[] {
     const dataCopy = cloneDeep(data);
     if (!isRecord(dataCopy)) {
-        return Array.isArray(dataCopy) ? (dataCopy as UnknownRecord[]).map(transformNumericKeysToArray) : (dataCopy as UnknownRecord);
+        return Array.isArray(dataCopy) ? (dataCopy as UnknownRecord[]).map(transformNumericKeysToArray) : dataCopy;
     }
 
     const keys = Object.keys(dataCopy);
@@ -50,4 +54,42 @@ function cleanAndTransformState<T>(state: string): T {
     return transformedState;
 }
 
-export {transformNumericKeysToArray, cleanAndTransformState};
+const processStateImport = (transformedState: OnyxValues) => {
+    const collectionKeys = [...new Set(Object.values(ONYXKEYS.COLLECTION))];
+    const collectionsMap = new Map<keyof OnyxCollectionValuesMapping, OnyxCollection<OnyxKey>>();
+    const regularState: Partial<Record<OnyxKey, OnyxEntry<OnyxKey>>> = {};
+
+    Object.entries(transformedState).forEach(([key, value]) => {
+        const collectionKey = collectionKeys.find((cKey) => key.startsWith(cKey));
+        if (collectionKey) {
+            if (!collectionsMap.has(collectionKey)) {
+                collectionsMap.set(collectionKey, {});
+            }
+
+            const collection = collectionsMap.get(collectionKey);
+            if (!collection) {
+                return;
+            }
+
+            collection[key] = value as OnyxEntry<OnyxKey>;
+        } else {
+            regularState[key as OnyxKey] = value as OnyxEntry<OnyxKey>;
+        }
+    });
+
+    return Onyx.clear(KEYS_TO_PRESERVE)
+        .then(() => {
+            const collectionPromises = Array.from(collectionsMap.entries()).map(([baseKey, items]) => {
+                return items ? Onyx.setCollection(baseKey, items) : Promise.resolve();
+            });
+            return Promise.all(collectionPromises);
+        })
+        .then(() => {
+            if (Object.keys(regularState).length > 0) {
+                return Onyx.multiSet(regularState as Partial<OnyxValues>);
+            }
+            return Promise.resolve();
+        });
+};
+
+export {cleanAndTransformState, processStateImport, transformNumericKeysToArray};

src/libs/actions/App.ts

@@ -570,38 +570,41 @@ function clearOnyxAndResetApp(shouldNavigateToHomepage?: boolean) {
     const isStateImported = isUsingImportedState;
     const shouldUseStagingServer = preservedShouldUseStagingServer;
     const sequentialQueue = getAll();
-    Onyx.clear(KEYS_TO_PRESERVE).then(() => {
-        // Network key is preserved, so when using imported state, we should stop forcing offline mode so that the app can re-fetch the network
-        if (isStateImported) {
-            setShouldForceOffline(false);
-        }
-
-        if (shouldNavigateToHomepage) {
-            Navigation.navigate(ROUTES.HOME);
-        }
 
-        if (preservedUserSession) {
-            Onyx.set(ONYXKEYS.SESSION, preservedUserSession);
-            Onyx.set(ONYXKEYS.PRESERVED_USER_SESSION, null);
-        }
+    Onyx.clear(KEYS_TO_PRESERVE)
+        .then(() => {
+            // Network key is preserved, so when using imported state, we should stop forcing offline mode so that the app can re-fetch the network
+            if (isStateImported) {
+                setShouldForceOffline(false);
+            }
 
-        if (shouldUseStagingServer) {
-            Onyx.set(ONYXKEYS.USER, {shouldUseStagingServer});
-        }
+            if (shouldNavigateToHomepage) {
+                Navigation.navigate(ROUTES.HOME);
+            }
 
-        // Requests in a sequential queue should be called even if the Onyx state is reset, so we do not lose any pending data.
-        // However, the OpenApp request must be called before any other request in a queue to ensure data consistency.
-        // To do that, sequential queue is cleared together with other keys, and then it's restored once the OpenApp request is resolved.
-        openApp().then(() => {
-            if (!sequentialQueue || isStateImported) {
-                return;
+            if (preservedUserSession) {
+                Onyx.set(ONYXKEYS.SESSION, preservedUserSession);
+                Onyx.set(ONYXKEYS.PRESERVED_USER_SESSION, null);
             }
 
-            sequentialQueue.forEach((request) => {
-                save(request);
+            if (shouldUseStagingServer) {
+                Onyx.set(ONYXKEYS.USER, {shouldUseStagingServer});
+            }
+        })
+        .then(() => {
+            // Requests in a sequential queue should be called even if the Onyx state is reset, so we do not lose any pending data.
+            // However, the OpenApp request must be called before any other request in a queue to ensure data consistency.
+            // To do that, sequential queue is cleared together with other keys, and then it's restored once the OpenApp request is resolved.
+            openApp().then(() => {
+                if (!sequentialQueue || isStateImported) {
+                    return;
+                }
+
+                sequentialQueue.forEach((request) => {
+                    save(request);
+                });
             });
         });
-    });
     clearSoundAssetsCache();
 }