appwrite-utils-cli 0.0.37 → 0.0.38

package/src/migrations/dataLoader.ts

@@ -93,6 +93,7 @@ export class DataLoader {
    * It iterates through the target object's keys and updates the source object if:
    * - The source object has the key.
    * - The target object's value for that key is not null, undefined, or an empty string.
+   * - If the target object has an array value, it concatenates the values and removes duplicates.
    *
    * @param source - The source object to be updated.
    * @param target - The target object with values to update the source object.
@@ -162,8 +163,26 @@ export class DataLoader {
   // Method to generate a unique ID that doesn't conflict with existing IDs
   getTrueUniqueId(collectionName: string) {
     let newId = ID.unique();
-    while (this.checkMapValuesForId(newId, collectionName)) {
+    let condition =
+      this.checkMapValuesForId(newId, collectionName) ||
+      this.userExistsMap.has(newId) ||
+      this.importMap
+        .get(this.getCollectionKey("users"))
+        ?.data.some(
+          (user) =>
+            user.finalData.docId === newId || user.finalData.userId === newId
+        );
+    while (condition) {
       newId = ID.unique();
+      condition =
+        this.checkMapValuesForId(newId, collectionName) ||
+        this.userExistsMap.has(newId) ||
+        this.importMap
+          .get(this.getCollectionKey("users"))
+          ?.data.some(
+            (user) =>
+              user.finalData.docId === newId || user.finalData.userId === newId
+          );
     }
     return newId;
   }
@@ -271,6 +290,21 @@ export class DataLoader {
         this.phoneToUserIdMap.set(user.phone, user.$id);
       }
       this.userExistsMap.set(user.$id, true);
+      let importData = this.importMap.get(this.getCollectionKey("users"));
+      if (!importData) {
+        importData = {
+          data: [],
+        };
+      }
+      importData.data.push({
+        finalData: {
+          ...user,
+          userId: user.$id,
+          docId: user.$id,
+        },
+        rawData: user,
+      });
+      this.importMap.set(this.getCollectionKey("users"), importData);
     }
     return allUsers;
   }
@@ -312,15 +346,11 @@ export class DataLoader {
         for (const createDef of createDefs) {
           if (!isUsersCollection) {
             console.log(`${collection.name} is not users collection`);
-            await this.prepareCreateData(db, collection, createDef);
+            this.prepareCreateData(db, collection, createDef);
           } else {
             // Special handling for users collection if needed
             console.log(`${collection.name} is users collection`);
-            await this.prepareUserCollectionCreateData(
-              db,
-              collection,
-              createDef
-            );
+            this.prepareUserCollectionCreateData(db, collection, createDef);
           }
         }
         for (const updateDef of updateDefs) {
@@ -331,11 +361,11 @@ export class DataLoader {
             continue;
           }
           // Prepare the update data for the collection
-          await this.prepareUpdateData(db, collection, updateDef);
+          this.prepareUpdateData(db, collection, updateDef);
         }
       }
       console.log("Running update references");
-      this.dealWithMergedUsers();
+      // this.dealWithMergedUsers();
       this.updateOldReferencesForNew();
       console.log("Done running update references");
     }
@@ -355,9 +385,15 @@ export class DataLoader {
       this.config.usersCollectionName
     );
     const usersCollectionPrimaryKeyFields = new Set();
+
     if (!this.config.collections) {
+      console.log("No collections found in configuration.");
       return;
     }
+
+    let needsUpdate = false;
+    let numUpdates = 0;
+
     // Collect primary key fields from the users collection definitions
     this.config.collections.forEach((collection) => {
       if (this.getCollectionKey(collection.name) === usersCollectionKey) {
@@ -373,45 +409,95 @@ export class DataLoader {
       }
     });
 
+    console.log(
+      `Primary key fields collected for users collection: ${[
+        ...usersCollectionPrimaryKeyFields,
+      ]}`
+    );
+
     // Iterate over all collections to update references based on merged users
     this.config.collections.forEach((collection) => {
       const collectionData = this.importMap.get(
         this.getCollectionKey(collection.name)
       );
-      if (!collectionData || !collectionData.data) return;
+
+      if (!collectionData || !collectionData.data) {
+        console.log(`No data found for collection ${collection.name}`);
+        return;
+      }
+
       const collectionImportDefs = collection.importDefs;
       if (!collectionImportDefs || !collectionImportDefs.length) {
+        console.log(
+          `No import definitions found for collection ${collection.name}`
+        );
         return;
       }
+
       collectionImportDefs.forEach((importDef) => {
         importDef.idMappings?.forEach((idMapping) => {
           if (
             this.getCollectionKey(idMapping.targetCollection) ===
             usersCollectionKey
           ) {
+            const fieldToSetKey = idMapping.fieldToSet || idMapping.sourceField;
             const targetFieldKey =
               idMapping.targetFieldToMatch || idMapping.targetField;
+
             if (usersCollectionPrimaryKeyFields.has(targetFieldKey)) {
+              console.log(
+                `Processing collection ${collection.name} with target field ${targetFieldKey}`
+              );
+
               // Process each item in the collection
               collectionData.data.forEach((item) => {
-                const oldId = item.context[idMapping.sourceField];
+                const oldId =
+                  item.finalData[idMapping.sourceField] ||
+                  item.context[idMapping.sourceField];
+
+                if (oldId === undefined || oldId === null) {
+                  console.log(
+                    `Skipping item with undefined or null oldId in collection ${collection.name}`
+                  );
+                  return;
+                }
+
                 const newId = this.mergedUserMap.get(`${oldId}`);
 
                 if (newId) {
+                  needsUpdate = true;
+                  numUpdates++;
+                  console.log(
+                    `Updating old ID ${oldId} to new ID ${newId} in collection ${collection.name}`
+                  );
+
                   // Update context to use new user ID
-                  item.finalData[
-                    idMapping.fieldToSet || idMapping.sourceField
-                  ] = newId;
+                  item.finalData[fieldToSetKey] = newId;
+                  item.context[fieldToSetKey] = newId;
+                } else {
+                  console.log(
+                    `No new ID found for old ID ${oldId} in mergedUserMap.`
+                  );
                 }
               });
             }
           }
         });
       });
+
+      if (needsUpdate) {
+        console.log(
+          `Updated ${numUpdates} references for collection ${collection.name}`
+        );
+        this.importMap.set(
+          this.getCollectionKey(collection.name),
+          collectionData
+        );
+      }
     });
   }
 
-  async updateOldReferencesForNew() {
+  updateOldReferencesForNew() {
     if (!this.config.collections) {
       return;
     }
@@ -612,12 +698,19 @@ export class DataLoader {
    * @param attributeMappings - The attribute mappings for the item.
    * @returns The transformed item with user-specific keys removed.
    */
-  async prepareUserData(
+  prepareUserData(
     item: any,
     attributeMappings: AttributeMappings,
     primaryKeyField: string,
     newId: string
-  ): Promise<any> {
+  ): {
+    transformedItem: any;
+    existingId: string | undefined;
+    userData: {
+      rawData: any;
+      finalData: z.infer<typeof AuthUserCreateSchema>;
+    };
+  } {
     let transformedItem = this.transformData(item, attributeMappings);
     const userData = AuthUserCreateSchema.safeParse(transformedItem);
     if (!userData.success) {
@@ -666,15 +759,24 @@ export class DataLoader {
         });
       if (userFound) {
         userFound.finalData.userId = existingId;
+        userFound.finalData.docId = existingId;
+        this.userExistsMap.set(existingId, true);
       }
-      return [
+
+      const userKeys = ["email", "phone", "name", "labels", "prefs"];
+      userKeys.forEach((key) => {
+        if (transformedItem.hasOwnProperty(key)) {
+          delete transformedItem[key];
+        }
+      });
+      return {
         transformedItem,
         existingId,
-        {
+        userData: {
           rawData: userFound?.rawData,
           finalData: userFound?.finalData,
         },
-      ];
+      };
     } else {
       existingId = newId;
       userData.data.userId = existingId;
@@ -696,7 +798,11 @@ export class DataLoader {
       data: [...(usersMap?.data || []), userDataToAdd],
     });
 
-    return [transformedItem, existingId, userDataToAdd];
+    return {
+      transformedItem,
+      existingId,
+      userData: userDataToAdd,
+    };
   }
 
   /**
@@ -708,11 +814,11 @@ export class DataLoader {
    * @param collection - The collection configuration.
    * @param importDef - The import definition containing the attribute mappings and other relevant info.
    */
-  async prepareUserCollectionCreateData(
+  prepareUserCollectionCreateData(
     db: ConfigDatabase,
     collection: CollectionCreate,
     importDef: ImportDef
-  ): Promise<void> {
+  ): void {
     // Load the raw data based on the import definition
     const rawData = this.loadData(importDef);
     const operationId = this.collectionImportOperations.get(
@@ -736,127 +842,157 @@ export class DataLoader {
         `No import operation found for collection ${collection.name}`
       );
     }
-    await updateOperation(this.database, operationId, {
+    updateOperation(this.database, operationId, {
       status: "ready",
       total: rawData.length,
-    });
-    // Retrieve the current user data and the current collection data from the import map
-    const currentUserData = this.importMap.get(this.getCollectionKey("users"));
-    const currentData = this.importMap.get(
-      this.getCollectionKey(collection.name)
-    );
-    // Log errors if the necessary data is not found in the import map
-    if (!currentUserData) {
-      logger.error(
-        `No data found for collection ${"users"} for createDef but it says it's supposed to have one...`
+    }).then(() => {
+      // Retrieve the current user data and the current collection data from the import map
+      const currentUserData = this.importMap.get(
+        this.getCollectionKey("users")
       );
-      return;
-    } else if (!currentData) {
-      logger.error(
-        `No data found for collection ${collection.name} for createDef but it says it's supposed to have one...`
-      );
-      return;
-    }
-    // Iterate through each item in the raw data
-    for (const item of rawData) {
-      // Prepare user data, check for duplicates, and remove user-specific keys
-      let [transformedItem, existingId, userData] = await this.prepareUserData(
-        item,
-        importDef.attributeMappings,
-        importDef.primaryKeyField,
-        this.getTrueUniqueId(this.getCollectionKey("users"))
+      const currentData = this.importMap.get(
+        this.getCollectionKey(collection.name)
       );
+      // Log errors if the necessary data is not found in the import map
+      if (!currentUserData) {
+        logger.error(
+          `No data found for collection ${"users"} for createDef but it says it's supposed to have one...`
+        );
+        return;
+      } else if (!currentData) {
+        logger.error(
+          `No data found for collection ${collection.name} for createDef but it says it's supposed to have one...`
+        );
+        return;
+      }
+      // Iterate through each item in the raw data
+      for (const item of rawData) {
+        // Prepare user data, check for duplicates, and remove user-specific keys
+        let { transformedItem, existingId, userData } = this.prepareUserData(
+          item,
+          importDef.attributeMappings,
+          importDef.primaryKeyField,
+          this.getTrueUniqueId(this.getCollectionKey("users"))
+        );
 
-      logger.info(
-        `In create user -- transformedItem: ${JSON.stringify(
-          transformedItem,
-          null,
-          2
-        )}`
-      );
+        logger.info(
+          `In create user -- transformedItem: ${JSON.stringify(
+            transformedItem,
+            null,
+            2
+          )}`
+        );
 
-      // Generate a new unique ID for the item or use existing ID
-      if (!existingId) {
-        // No existing user ID, generate a new unique ID
-        existingId = this.getTrueUniqueId(this.getCollectionKey("users"));
-        transformedItem.docId = existingId; // Assign the new ID to the transformed data's docId field
-      }
+        // Generate a new unique ID for the item or use existing ID
+        if (!existingId) {
+          // No existing user ID, generate a new unique ID
+          existingId = this.getTrueUniqueId(this.getCollectionKey("users"));
+          transformedItem.docId = existingId; // Assign the new ID to the transformed data's docId field
+        }
 
-      // Create a context object for the item, including the new ID
-      let context = this.createContext(db, collection, item, existingId);
+        // Create a context object for the item, including the new ID
+        let context = this.createContext(db, collection, item, existingId);
 
-      // Merge the transformed data into the context
-      context = { ...context, ...transformedItem, ...userData.finalData };
+        // Merge the transformed data into the context
+        context = { ...context, ...transformedItem, ...userData.finalData };
 
-      // If a primary key field is defined, handle the ID mapping and check for duplicates
-      if (importDef.primaryKeyField) {
-        const oldId = item[importDef.primaryKeyField];
+        // If a primary key field is defined, handle the ID mapping and check for duplicates
+        if (importDef.primaryKeyField) {
+          const oldId = item[importDef.primaryKeyField];
 
-        // Check if the oldId already exists to handle potential duplicates
-        if (
-          this.oldIdToNewIdPerCollectionMap
-            .get(this.getCollectionKey(collection.name))
-            ?.has(`${oldId}`)
-        ) {
-          // Found a duplicate oldId, now decide how to merge or handle these duplicates
-          for (const data of currentData.data) {
-            if (
-              data.finalData.docId === oldId ||
-              data.finalData.userId === oldId
-            ) {
-              transformedItem = this.mergeObjects(
-                data.finalData,
-                transformedItem
-              );
+          // Check if the oldId already exists to handle potential duplicates
+          if (
+            this.oldIdToNewIdPerCollectionMap
+              .get(this.getCollectionKey(collection.name))
+              ?.has(`${oldId}`)
+          ) {
+            // Found a duplicate oldId, now decide how to merge or handle these duplicates
+            for (const data of currentData.data) {
+              if (
+                data.finalData.docId === oldId ||
+                data.finalData.userId === oldId
+              ) {
+                transformedItem = this.mergeObjects(
+                  data.finalData,
+                  transformedItem
+                );
+              }
             }
+          } else {
+            // No duplicate found, simply map the oldId to the new itemId
+            collectionOldIdToNewIdMap?.set(`${oldId}`, `${existingId}`);
           }
-        } else {
-          // No duplicate found, simply map the oldId to the new itemId
-          collectionOldIdToNewIdMap?.set(`${oldId}`, `${existingId}`);
         }
-      }
-      // Merge the final user data into the context
-      context = { ...context, ...userData.finalData };
 
-      // Handle merging for currentUserData
-      for (let i = 0; i < currentUserData.data.length; i++) {
-        if (
-          (currentUserData.data[i].finalData.docId === existingId ||
-            currentUserData.data[i].finalData.userId === existingId) &&
-          !_.isEqual(currentUserData.data[i], userData)
-        ) {
-          this.mergeObjects(
-            currentUserData.data[i].finalData,
+        // Handle merging for currentUserData
+        for (let i = 0; i < currentUserData.data.length; i++) {
+          const currentUserDataItem = currentUserData.data[i];
+          const samePhones =
+            currentUserDataItem.finalData.phone &&
+            transformedItem.phone &&
+            currentUserDataItem.finalData.phone === transformedItem.phone;
+          const sameEmails =
+            currentUserDataItem.finalData.email &&
+            transformedItem.email &&
+            currentUserDataItem.finalData.email === transformedItem.email;
+          if (
+            (currentUserDataItem.finalData.docId === existingId ||
+              currentUserDataItem.finalData.userId === existingId) &&
+            (samePhones || sameEmails) &&
+            currentUserDataItem.finalData &&
             userData.finalData
-          );
-          console.log("Merging user data", currentUserData.data[i].finalData);
-          this.importMap.set(this.getCollectionKey("users"), currentUserData);
+          ) {
+            const userDataMerged = this.mergeObjects(
+              currentUserData.data[i].finalData,
+              userData.finalData
+            );
+            currentUserData.data[i].finalData = userDataMerged;
+            this.importMap.set(this.getCollectionKey("users"), currentUserData);
+          }
         }
-      }
-      // Update the attribute mappings with any actions that need to be performed post-import
-      const mappingsWithActions = this.getAttributeMappingsWithActions(
-        importDef.attributeMappings,
-        context,
-        transformedItem
-      );
-      // Update the import definition with the new attribute mappings
-      const newImportDef = {
-        ...importDef,
-        attributeMappings: mappingsWithActions,
-      };
+        // Update the attribute mappings with any actions that need to be performed post-import
+        const mappingsWithActions = this.getAttributeMappingsWithActions(
+          importDef.attributeMappings,
+          context,
+          transformedItem
+        );
+        // Update the import definition with the new attribute mappings
+        const newImportDef = {
+          ...importDef,
+          attributeMappings: mappingsWithActions,
+        };
 
-      let foundData = false;
-      for (let i = 0; i < currentData.data.length; i++) {
-        if (
-          currentData.data[i].finalData.docId === existingId ||
-          currentData.data[i].finalData.userId === existingId
-        ) {
-          currentData.data[i].finalData = this.mergeObjects(
-            currentData.data[i].finalData,
-            transformedItem
-          );
-          currentData.data[i].context = context;
-          currentData.data[i].importDef = newImportDef;
+        let foundData = false;
+        for (let i = 0; i < currentData.data.length; i++) {
+          if (
+            currentData.data[i].finalData.docId === existingId ||
+            currentData.data[i].finalData.userId === existingId
+          ) {
+            currentData.data[i].finalData = this.mergeObjects(
+              currentData.data[i].finalData,
+              transformedItem
+            );
+            currentData.data[i].context = context;
+            currentData.data[i].importDef = newImportDef;
+            this.importMap.set(
+              this.getCollectionKey(collection.name),
+              currentData
+            );
+            this.oldIdToNewIdPerCollectionMap.set(
+              this.getCollectionKey(collection.name),
+              collectionOldIdToNewIdMap!
+            );
+            foundData = true;
+          }
+        }
+        if (!foundData) {
+          // Add new data to the associated collection
+          currentData.data.push({
+            rawData: item,
+            context: context,
+            importDef: newImportDef,
+            finalData: transformedItem,
+          });
           this.importMap.set(
             this.getCollectionKey(collection.name),
             currentData
@@ -865,24 +1001,9 @@ export class DataLoader {
             this.getCollectionKey(collection.name),
             collectionOldIdToNewIdMap!
           );
-          foundData = true;
         }
       }
-      if (!foundData) {
-        // Add new data to the associated collection
-        currentData.data.push({
-          rawData: item,
-          context: context,
-          importDef: newImportDef,
-          finalData: transformedItem,
-        });
-        this.importMap.set(this.getCollectionKey(collection.name), currentData);
-        this.oldIdToNewIdPerCollectionMap.set(
-          this.getCollectionKey(collection.name),
-          collectionOldIdToNewIdMap!
-        );
-      }
-    }
+    });
   }
 
   /**
@@ -893,11 +1014,11 @@ export class DataLoader {
    * @param collection - The collection configuration.
    * @param importDef - The import definition containing the attribute mappings and other relevant info.
    */
-  async prepareCreateData(
+  prepareCreateData(
     db: ConfigDatabase,
     collection: CollectionCreate,
     importDef: ImportDef
-  ): Promise<void> {
+  ): void {
     // Load the raw data based on the import definition
     const rawData = this.loadData(importDef);
     const operationId = this.collectionImportOperations.get(
@@ -908,94 +1029,99 @@ export class DataLoader {
         `No import operation found for collection ${collection.name}`
       );
     }
-    await updateOperation(this.database, operationId, {
+    updateOperation(this.database, operationId, {
       status: "ready",
       total: rawData.length,
-    });
-    // Initialize a new map for old ID to new ID mappings
-    const oldIdToNewIdMapNew = new Map<string, string>();
-    // Retrieve or initialize the collection-specific old ID to new ID map
-    const collectionOldIdToNewIdMap =
-      this.oldIdToNewIdPerCollectionMap.get(
-        this.getCollectionKey(collection.name)
-      ) ||
-      this.oldIdToNewIdPerCollectionMap
-        .set(this.getCollectionKey(collection.name), oldIdToNewIdMapNew)
-        .get(this.getCollectionKey(collection.name));
-    console.log(
-      `${collection.name} -- collectionOldIdToNewIdMap: ${collectionOldIdToNewIdMap}`
-    );
-    // Iterate through each item in the raw data
-    for (const item of rawData) {
-      // Generate a new unique ID for the item
-      const itemIdNew = this.getTrueUniqueId(
-        this.getCollectionKey(collection.name)
-      );
-      // Retrieve the current collection data from the import map
-      const currentData = this.importMap.get(
-        this.getCollectionKey(collection.name)
-      );
-      // Create a context object for the item, including the new ID
-      let context = this.createContext(db, collection, item, itemIdNew);
-      // Transform the item data based on the attribute mappings
-      const transformedData = this.transformData(
-        item,
-        importDef.attributeMappings
+    }).then(() => {
+      // Initialize a new map for old ID to new ID mappings
+      const oldIdToNewIdMapNew = new Map<string, string>();
+      // Retrieve or initialize the collection-specific old ID to new ID map
+      const collectionOldIdToNewIdMap =
+        this.oldIdToNewIdPerCollectionMap.get(
+          this.getCollectionKey(collection.name)
+        ) ||
+        this.oldIdToNewIdPerCollectionMap
+          .set(this.getCollectionKey(collection.name), oldIdToNewIdMapNew)
+          .get(this.getCollectionKey(collection.name));
+      console.log(
+        `${collection.name} -- collectionOldIdToNewIdMap: ${collectionOldIdToNewIdMap}`
       );
-      // If a primary key field is defined, handle the ID mapping and check for duplicates
-      if (importDef.primaryKeyField) {
-        const oldId = item[importDef.primaryKeyField];
-        if (collectionOldIdToNewIdMap?.has(`${oldId}`)) {
+      // Iterate through each item in the raw data
+      for (const item of rawData) {
+        // Generate a new unique ID for the item
+        const itemIdNew = this.getTrueUniqueId(
+          this.getCollectionKey(collection.name)
+        );
+        // Retrieve the current collection data from the import map
+        const currentData = this.importMap.get(
+          this.getCollectionKey(collection.name)
+        );
+        // Create a context object for the item, including the new ID
+        let context = this.createContext(db, collection, item, itemIdNew);
+        // Transform the item data based on the attribute mappings
+        const transformedData = this.transformData(
+          item,
+          importDef.attributeMappings
+        );
+        // If a primary key field is defined, handle the ID mapping and check for duplicates
+        if (importDef.primaryKeyField) {
+          const oldId = item[importDef.primaryKeyField];
+          if (collectionOldIdToNewIdMap?.has(`${oldId}`)) {
+            logger.error(
+              `Collection ${collection.name} has multiple documents with the same primary key ${oldId}`
+            );
+            continue;
+          }
+          collectionOldIdToNewIdMap?.set(`${oldId}`, `${itemIdNew}`);
+        }
+        // Merge the transformed data into the context
+        context = { ...context, ...transformedData };
+        // Validate the item before proceeding
+        const isValid = this.importDataActions.validateItem(
+          transformedData,
+          importDef.attributeMappings,
+          context
+        );
+        if (!isValid) {
+          continue;
+        }
+        // Update the attribute mappings with any actions that need to be performed post-import
+        const mappingsWithActions = this.getAttributeMappingsWithActions(
+          importDef.attributeMappings,
+          context,
+          transformedData
+        );
+        // Update the import definition with the new attribute mappings
+        const newImportDef = {
+          ...importDef,
+          attributeMappings: mappingsWithActions,
+        };
+        // If the current collection data exists, add the item with its context and final data
+        if (currentData && currentData.data) {
+          currentData.data.push({
+            rawData: item,
+            context: context,
+            importDef: newImportDef,
+            finalData: transformedData,
+          });
+          this.importMap.set(
+            this.getCollectionKey(collection.name),
+            currentData
+          );
+          this.oldIdToNewIdPerCollectionMap.set(
+            this.getCollectionKey(collection.name),
+            collectionOldIdToNewIdMap!
+          );
+        } else {
           logger.error(
-            `Collection ${collection.name} has multiple documents with the same primary key ${oldId}`
+            `No data found for collection ${collection.name} for createDef but it says it's supposed to have one...`
           );
           continue;
         }
-        collectionOldIdToNewIdMap?.set(`${oldId}`, `${itemIdNew}`);
       }
-      // Merge the transformed data into the context
-      context = { ...context, ...transformedData };
-      // Validate the item before proceeding
-      const isValid = await this.importDataActions.validateItem(
-        transformedData,
-        importDef.attributeMappings,
-        context
-      );
-      if (!isValid) {
-        continue;
-      }
-      // Update the attribute mappings with any actions that need to be performed post-import
-      const mappingsWithActions = this.getAttributeMappingsWithActions(
-        importDef.attributeMappings,
-        context,
-        transformedData
-      );
-      // Update the import definition with the new attribute mappings
-      const newImportDef = {
-        ...importDef,
-        attributeMappings: mappingsWithActions,
-      };
-      // If the current collection data exists, add the item with its context and final data
-      if (currentData && currentData.data) {
-        currentData.data.push({
-          rawData: item,
-          context: context,
-          importDef: newImportDef,
-          finalData: transformedData,
-        });
-        this.importMap.set(this.getCollectionKey(collection.name), currentData);
-        this.oldIdToNewIdPerCollectionMap.set(
-          this.getCollectionKey(collection.name),
-          collectionOldIdToNewIdMap!
-        );
-      } else {
-        logger.error(
-          `No data found for collection ${collection.name} for createDef but it says it's supposed to have one...`
-        );
-        continue;
-      }
-    }
+    });
   }
+
   /**
    * Prepares the data for updating documents within a collection.
    * This method loads the raw data based on the import definition, transforms it according to the attribute mappings,
@@ -1006,7 +1132,7 @@ export class DataLoader {
    * @param collection - The collection configuration.
    * @param importDef - The import definition containing the attribute mappings and other relevant info.
    */
-  async prepareUpdateData(
+  prepareUpdateData(
     db: ConfigDatabase,
     collection: CollectionCreate,
     importDef: ImportDef
@@ -1047,11 +1173,7 @@ export class DataLoader {
       let newId: string | undefined;
       let oldId: string | undefined;
       // Determine the new ID for the item based on the primary key field or update mapping
-      if (importDef.primaryKeyField) {
-        oldId = item[importDef.primaryKeyField];
-      } else if (importDef.updateMapping) {
-        oldId = item[importDef.updateMapping.originalIdField];
-      }
+      oldId = item[importDef.primaryKeyField];
       if (oldId) {
         newId = oldIdToNewIdMap?.get(`${oldId}`);
         if (
@@ -1074,8 +1196,13 @@ export class DataLoader {
         );
         continue;
       }
+      const itemDataToUpdate = this.importMap
+        .get(this.getCollectionKey(collection.name))
+        ?.data.find(
+          (data) => `${data.context[importDef.primaryKeyField]}` === `${oldId}`
+        );
       // Log an error and continue to the next item if no new ID is found
-      if (!newId) {
+      if (!newId && !itemDataToUpdate) {
         logger.error(
           `No new id found for collection ${
             collection.name
@@ -1086,15 +1213,24 @@ export class DataLoader {
           )} but it says it's supposed to have one...`
         );
         continue;
+      } else if (itemDataToUpdate) {
+        newId = itemDataToUpdate.finalData.docId;
+        if (!newId) {
+          logger.error(
+            `No new id found for collection ${
+              collection.name
+            } for updateDef ${JSON.stringify(
+              item,
+              null,
+              2
+            )} but it says it's supposed to have one...`
+          );
+          continue;
+        }
       }
-      const itemDataToUpdate = this.importMap
-        .get(this.getCollectionKey(collection.name))
-        ?.data.find(
-          (data) => data.rawData[importDef.primaryKeyField] === oldId
-        );
-      if (!itemDataToUpdate) {
+      if (!itemDataToUpdate || !newId) {
         logger.error(
-          `No data found for collection ${
+          `No data or ID (docId) found for collection ${
             collection.name
           } for updateDef ${JSON.stringify(
             item,
@@ -1110,9 +1246,9 @@ export class DataLoader {
       );
       // Create a context object for the item, including the new ID and transformed data
       let context = this.createContext(db, collection, item, newId);
-      context = this.mergeObjects(context, transformedData);
+      context = { ...context, ...transformedData };
       // Validate the item before proceeding
-      const isValid = await this.importDataActions.validateItem(
+      const isValid = this.importDataActions.validateItem(
         item,
         importDef.attributeMappings,
         context
@@ -1144,6 +1280,7 @@ export class DataLoader {
         );
         itemDataToUpdate.context = context;
         itemDataToUpdate.importDef = newImportDef;
+        currentData!.data.push(itemDataToUpdate);
       } else {
         // If no existing item matches, then add the new item
         currentData!.data.push({