DOP-3633: Use bulkWrites for asset upsertion (#814)

Author: rayangler

modules/persistence/src/services/assets/index.ts

@@ -1,5 +1,5 @@
 import AdmZip from 'adm-zip';
-import { upsert } from '../connector';
+import { bulkUpsert } from '../connector';
 
 const COLLECTION_NAME = 'assets';
 
@@ -14,8 +14,8 @@ const assetsFromZip = (zip: AdmZip) => {
 
 export const upsertAssets = async (zip: AdmZip) => {
   try {
-    const assets = await assetsFromZip(zip);
-    return Promise.all(assets.map((asset) => upsert(asset, COLLECTION_NAME, asset._id)));
+    const assets = assetsFromZip(zip);
+    return bulkUpsert(assets, COLLECTION_NAME);
   } catch (error) {
     console.error(`Error at upsertion time for ${COLLECTION_NAME}: ${error}`);
     throw error;

modules/persistence/src/services/connector/index.ts

@@ -4,7 +4,7 @@
 // If no, the helper should be implemented in that service, not here
 
 import * as mongodb from 'mongodb';
-import { ObjectId, Db } from 'mongodb';
+import { ObjectId, Db, Document } from 'mongodb';
 import { db as poolDb } from './pool';
 
 // We should only ever have one client active at a time.
@@ -51,15 +51,23 @@ export const insert = async (docs: any[], collection: string, buildId: ObjectId)
 };
 
 // Upsert wrapper, requires an _id field.
-export const upsert = async (payload: any, collection: string, _id: string | ObjectId) => {
+export const bulkUpsert = async (items: Document[], collection: string) => {
   const upsertSession = await db();
   try {
-    const query = { _id };
-    const update = { $set: payload };
-    const options = { upsert: true };
-    return await upsertSession.collection(collection).updateOne(query, update, options);
+    const operations: mongodb.AnyBulkWriteOperation[] = [];
+    items.forEach((item: Document) => {
+      const op = {
+        updateOne: {
+          filter: { _id: item._id },
+          update: { $set: item },
+          upsert: true,
+        },
+      };
+      operations.push(op);
+    });
+    return upsertSession.collection(collection).bulkWrite(operations);
   } catch (error) {
-    console.error(`Error at upsertion time for ${collection}: ${error}`);
+    console.error(`Error at bulk upsertion time for ${collection}: ${error}`);
     throw error;
   }
 };

modules/persistence/src/services/metadata/repos_branches/index.ts

@@ -1,4 +1,4 @@
-import { pool, db } from '../../connector';
+import { pool } from '../../connector';
 import { Metadata } from '..';
 import { project } from '../ToC';
 import { WithId } from 'mongodb';

modules/persistence/tests/services/connector.test.ts

@@ -1,11 +1,11 @@
 import { ObjectID } from 'bson';
-import { db, insert, upsert } from '../../src/services/connector';
+import { bulkUpsert, db, insert } from '../../src/services/connector';
 
 const mockConnect = jest.fn();
 const mockDb = jest.fn();
 const mockCollection = jest.fn();
 const mockInsertMany = jest.fn();
-const mockUpdateOne = jest.fn();
+const mockBulkWrite = jest.fn();
 const mockClose = jest.fn();
 
 // below is a "jest mock" of a mongodb client
@@ -31,8 +31,8 @@ jest.mock('mongodb', () => ({
     async insertMany() {
       return mockInsertMany();
     }
-    async updateOne(...args) {
-      return mockUpdateOne(...args);
+    async bulkWrite(...args) {
+      return mockBulkWrite(...args);
     }
     close() {
       mockClose();
@@ -118,20 +118,28 @@ describe('Connector module', () => {
     });
   });
 
-  describe('upsert', () => {
-    const payload = { name: 'upsert-doc' };
+  describe('bulkUpsert', () => {
+    const payload = { _id: 'test-id', name: 'upsert-doc' };
     const collection = 'metadata';
-    const id = 'test-id';
+
     test('it calls on collection to update one with upsert option true', async () => {
-      await upsert(payload, collection, id);
+      await bulkUpsert([payload], collection);
       expect(mockCollection).toBeCalledWith(collection);
-      expect(mockUpdateOne).toBeCalledWith({ _id: 'test-id' }, { $set: { name: 'upsert-doc' } }, { upsert: true });
+      expect(mockBulkWrite).toBeCalledWith([
+        {
+          updateOne: {
+            filter: { _id: payload._id },
+            update: { $set: payload },
+            upsert: true,
+          },
+        },
+      ]);
     });
 
     test('it throws error on updateone error', async () => {
-      mockUpdateOne.mockRejectedValueOnce(new Error('test error') as never);
+      mockBulkWrite.mockRejectedValueOnce(new Error('test error') as never);
       try {
-        await upsert(payload, collection, id);
+        await bulkUpsert([payload], collection);
       } catch (e) {
         expect(e.message).toEqual('test error');
       }