refactor: fallback to modifyOneMethod with query COMPASS-6058 (#4373)

* refactor: fallback to modifyOneMethod with query COMPASS-6058

* refactor: remove unused variable

* refactor: bring back comments

Author: alenakhineika

packages/compass-crud/src/stores/crud-store.spec.ts

@@ -2017,7 +2017,6 @@ describe('store', function () {
     let updateOneFake;
     let replaceOneFake;
 
-    const document = { _id: 1234, name: 'document_1234' };
     const updatedDocument = { _id: 1234, name: 'document_12345' };
 
     beforeEach(function () {
@@ -2123,12 +2122,12 @@ describe('store', function () {
       });
     });
 
-    it('retries findOneAndUpdate with FLE returnDocument: "after"', async function () {
+    it('does updateOne with FLE returnDocument: "after"', async function () {
       const err = Object.assign(new Error('failed'), { code: 6371402 });
       findFake.callsFake(() => Promise.resolve([updatedDocument]));
       findOneAndReplaceFake.resolves({});
       findOneAndUpdateFake.onCall(0).rejects(err);
-      findOneAndUpdateFake.onCall(1).resolves(document);
+      updateOneFake.onCall(0).resolves({});
       const [error, d] = await findAndModifyWithFLEFallback(
         dataServiceStub,
         'db.coll',
@@ -2138,7 +2137,7 @@ describe('store', function () {
       );
       expect(error).to.equal(undefined);
       expect(d).to.deep.equal(updatedDocument);
-      expect(findOneAndUpdateFake).to.have.callCount(2);
+      expect(findOneAndUpdateFake).to.have.callCount(1);
 
       expect(findOneAndUpdateFake.firstCall.args[0]).to.equal('db.coll');
       expect(findOneAndUpdateFake.firstCall.args[1]).to.deep.equal({
@@ -2152,23 +2151,18 @@ describe('store', function () {
         promoteValues: false,
       });
 
-      expect(findOneAndUpdateFake.secondCall.args[0]).to.equal('db.coll');
-      expect(findOneAndUpdateFake.secondCall.args[1]).to.deep.equal({
+      expect(updateOneFake.firstCall.args[0]).to.equal('db.coll');
+      expect(updateOneFake.firstCall.args[1]).to.deep.equal({
         _id: 1234,
       });
-      expect(findOneAndUpdateFake.secondCall.args[2]).to.deep.equal({
+      expect(updateOneFake.firstCall.args[2]).to.deep.equal({
         name: 'document_12345',
       });
-      expect(findOneAndUpdateFake.secondCall.args[3]).to.deep.equal({
-        returnDocument: 'before',
-        promoteValues: false,
-      });
 
       expect(findFake).to.have.callCount(1);
       expect(findFake.firstCall.args[0]).to.equal('db.coll');
       expect(findFake.firstCall.args[1]).to.deep.equal({ _id: 1234 });
       expect(findFake.firstCall.args[2]).to.deep.equal({
-        returnDocument: 'before',
         promoteValues: false,
       });
     });
@@ -2178,7 +2172,7 @@ describe('store', function () {
       findFake.yields(new Error('find failed'));
       findOneAndReplaceFake.resolves({});
       findOneAndUpdateFake.onCall(0).rejects(err);
-      findOneAndUpdateFake.onCall(1).resolves(document);
+      updateOneFake.onCall(0).resolves({});
       const [error, d] = await findAndModifyWithFLEFallback(
         dataServiceStub,
         'db.coll',
@@ -2188,7 +2182,8 @@ describe('store', function () {
       );
       expect(error).to.equal(err);
       expect(d).to.equal(undefined);
-      expect(findOneAndUpdateFake).to.have.callCount(2);
+      expect(findOneAndUpdateFake).to.have.callCount(1);
+      expect(updateOneFake).to.have.callCount(1);
       expect(findFake).to.have.callCount(1);
     });
 

packages/compass-crud/src/stores/crud-store.ts

@@ -1575,67 +1575,49 @@ export async function findAndModifyWithFLEFallback(
   object: { $set?: BSONObject; $unset?: BSONObject } | BSONObject | BSONArray,
   modificationType: 'update' | 'replace'
 ): Promise<ErrorOrResult> {
-  let options: { returnDocument: 'before' | 'after'; promoteValues: false } = {
-    returnDocument: 'after',
-    promoteValues: false,
-  };
-  let error: (Error & { codeName?: string; code?: any }) | undefined;
-  let document;
-
   const findOneAndModifyMethod =
     modificationType === 'update' ? 'findOneAndUpdate' : 'findOneAndReplace';
-  const modifyOneMethod =
-    modificationType === 'update' ? 'updateOne' : 'replaceOne';
+  let error: (Error & { codeName?: string; code?: any }) | undefined;
 
   try {
-    document = await ds[findOneAndModifyMethod](ns, query, object, options);
+    return [
+      undefined,
+      await ds[findOneAndModifyMethod](ns, query, object, {
+        returnDocument: 'after',
+        promoteValues: false,
+      }),
+    ] as ErrorOrResult;
   } catch (e) {
     error = e as Error;
   }
 
-  if (!error) {
-    return [undefined, document] as ErrorOrResult;
-  }
+  if (
+    error.codeName === 'ShardKeyNotFound' ||
+    +(error?.code ?? 0) === 63714_02 // 6371402 is "'findAndModify with encryption only supports new: false'"
+  ) {
+    const modifyOneMethod =
+      modificationType === 'update' ? 'updateOne' : 'replaceOne';
 
-  if (error.codeName === 'ShardKeyNotFound') {
     try {
-      const docs = await ds.find(ns, query, options);
-      [document] = docs;
+      await ds[modifyOneMethod](ns, query, object);
     } catch (e) {
-      return [e, undefined] as ErrorOrResult;
-    }
-
-    if (document) {
-      try {
-        await ds[modifyOneMethod](ns, { _id: document._id }, object);
-      } catch (e) {
-        return [e, undefined] as ErrorOrResult;
-      }
-    }
-  } else if (+(error?.code ?? 0) === 63714_02) {
-    // 6371402 is "'findAndModify with encryption only supports new: false'"
-    options = {
-      returnDocument: 'before',
-      promoteValues: false,
-    };
-    try {
-      document = await ds[findOneAndModifyMethod](ns, query, object, options);
-    } catch (e) {
-      // Return the current findOneAndModify error here
-      // since we already know the exact error from the previous attempt
-      // and want to know what went wrong with the second attempt with the fallback options,
+      // Return the modifyOneMethod error here
+      // since we already know the original error from findOneAndModifyMethod
+      // and want to know what went wrong with the fallback method,
       // e.g. return the `Found indexed encrypted fields but could not find __safeContent__` error.
       return [e, undefined] as ErrorOrResult;
     }
-  }
 
-  try {
-    if (document) {
-      const docs = await ds.find(ns, { _id: document._id }, options);
+    try {
+      const docs = await ds.find(
+        ns,
+        { _id: query._id as any },
+        { promoteValues: false }
+      );
       return [undefined, docs[0]] as ErrorOrResult;
+    } catch (e) {
+      /* fallthrough */
     }
-  } catch (e) {
-    /* fallthrough */
   }
 
   // Race condition -- most likely, somebody else