From ca7922faddae6a9dd3d4caf24e22ed302a870809 Mon Sep 17 00:00:00 2001
From: Sergey Petushkov <petushkov.sergey@gmail.com>
Date: Mon, 30 Sep 2024 17:36:37 +0200
Subject: [PATCH 1/2] chore(atlas-service, indexes): remove all the domain
 specific parts of automation agent request helper

---
 packages/atlas-service/src/atlas-service.ts   | 109 +++++++---
 .../make-automation-agent-op-request.spec.ts  | 124 ------------
 .../src/make-automation-agent-op-request.ts   | 186 ------------------
 .../src/modules/rolling-indexes-service.ts    |  57 ++++--
 4 files changed, 126 insertions(+), 350 deletions(-)
 delete mode 100644 packages/atlas-service/src/make-automation-agent-op-request.spec.ts
 delete mode 100644 packages/atlas-service/src/make-automation-agent-op-request.ts

diff --git a/packages/atlas-service/src/atlas-service.ts b/packages/atlas-service/src/atlas-service.ts
index 8abe8f2c03b..c57a506b852 100644
--- a/packages/atlas-service/src/atlas-service.ts
+++ b/packages/atlas-service/src/atlas-service.ts
@@ -9,11 +9,6 @@ import {
 import type { Logger } from '@mongodb-js/compass-logging';
 import type { PreferencesAccess } from 'compass-preferences-model';
 import type { AtlasClusterMetadata } from '@mongodb-js/connection-info';
-import type {
-  AutomationAgentRequestTypes,
-  AutomationAgentResponse,
-} from './make-automation-agent-op-request';
-import { makeAutomationAgentOpRequest } from './make-automation-agent-op-request';
 
 export type AtlasServiceOptions = {
   defaultHeaders?: Record<string, string>;
@@ -104,34 +99,90 @@ export class AtlasService {
       },
     });
   }
-  automationAgentFetch<OpType extends keyof AutomationAgentRequestTypes>(
-    atlasMetadata: Pick<
-      AtlasClusterMetadata,
-      | 'projectId'
-      | 'clusterUniqueId'
-      | 'regionalBaseUrl'
-      | 'metricsType'
-      | 'metricsId'
-    >,
-    opType: OpType,
-    opBody: Omit<
-      AutomationAgentRequestTypes[OpType],
-      'clusterId' | 'serverlessId'
-    >
-  ): Promise<AutomationAgentResponse<OpType>> {
+  async automationAgentRequest(
+    atlasMetadata: AtlasClusterMetadata,
+    opType: string,
+    opBody: Record<string, unknown>
+  ): Promise<{ _id: string; requestType: string } | undefined> {
     const opBodyClusterId =
       atlasMetadata.metricsType === 'serverless'
         ? { serverlessId: atlasMetadata.clusterUniqueId }
         : { clusterId: atlasMetadata.metricsId };
-    return makeAutomationAgentOpRequest(
-      this.authenticatedFetch.bind(this),
-      this.regionalizedCloudEndpoint(atlasMetadata),
-      atlasMetadata.projectId,
-      opType,
-      Object.assign(
-        opBodyClusterId,
-        opBody
-      ) as AutomationAgentRequestTypes[OpType]
+    const requestUrl = this.regionalizedCloudEndpoint(
+      atlasMetadata,
+      `/explorer/v1/groups/${atlasMetadata.projectId}/requests/${opType}`
     );
+    const json = await this.authenticatedFetch(requestUrl, {
+      method: 'POST',
+      headers: {
+        'content-type': 'application/json',
+      },
+      body: JSON.stringify({ ...opBodyClusterId, ...opBody }),
+    }).then((res) => {
+      if (
+        res.headers
+          .get('content-type')
+          ?.toLowerCase()
+          .includes('application/json')
+      ) {
+        return res.json();
+      }
+    });
+    assertAutomationAgentRequestResponse(json, opType);
+    return json;
+  }
+  async automationAgentAwait<T>(
+    atlasMetadata: AtlasClusterMetadata,
+    opType: string,
+    requestId: string
+  ): Promise<{
+    _id: string;
+    requestType: string;
+    response: T[];
+  }> {
+    const requestUrl = this.regionalizedCloudEndpoint(
+      atlasMetadata,
+      `/explorer/v1/groups/${atlasMetadata.projectId}/requests/${requestId}/types/${opType}/await`
+    );
+    const json = await this.authenticatedFetch(requestUrl, {
+      method: 'GET',
+    }).then((res) => {
+      return res.json();
+    });
+    assertAutomationAgentAwaitResponse<T>(json, opType);
+    return json;
+  }
+}
+
+function assertAutomationAgentRequestResponse(
+  json: any,
+  opType: string
+): asserts json is { _id: string; requestType: string } {
+  if (
+    Object.prototype.hasOwnProperty.call(json, '_id') &&
+    Object.prototype.hasOwnProperty.call(json, 'requestType') &&
+    json.requestType === opType
+  ) {
+    return;
+  }
+  throw new Error(
+    'Got unexpected backend response for automation agent request'
+  );
+}
+
+function assertAutomationAgentAwaitResponse<T>(
+  json: any,
+  opType: string
+): asserts json is { _id: string; requestType: string; response: T[] } {
+  if (
+    Object.prototype.hasOwnProperty.call(json, '_id') &&
+    Object.prototype.hasOwnProperty.call(json, 'requestType') &&
+    Object.prototype.hasOwnProperty.call(json, 'response') &&
+    json.requestType === opType
+  ) {
+    return;
   }
+  throw new Error(
+    'Got unexpected backend response for automation agent request await'
+  );
 }
diff --git a/packages/atlas-service/src/make-automation-agent-op-request.spec.ts b/packages/atlas-service/src/make-automation-agent-op-request.spec.ts
deleted file mode 100644
index c989d62fca7..00000000000
--- a/packages/atlas-service/src/make-automation-agent-op-request.spec.ts
+++ /dev/null
@@ -1,124 +0,0 @@
-import Sinon from 'sinon';
-import { makeAutomationAgentOpRequest } from './make-automation-agent-op-request';
-import { expect } from 'chai';
-
-describe('makeAutomationAgentOpRequest', function () {
-  const successSpecs = [
-    [
-      'succeeds if backend returned requestId and response',
-      { _id: 'abc', requestType: 'listIndexStats' },
-      {
-        _id: 'abc',
-        requestType: 'listIndexStats',
-        response: [{ indexName: 'test' }],
-      },
-    ],
-  ] as const;
-
-  const failSpecs = [
-    [
-      'fails if initial request fails',
-      new Error('NetworkError'),
-      {},
-      /NetworkError/,
-    ],
-    [
-      'fails if await response fails',
-      { _id: 'abc', requestType: 'listIndexStats' },
-      new Error('NetworkError'),
-      /NetworkError/,
-    ],
-    [
-      'fails if backend did not return requestId',
-      {},
-      {},
-      /Got unexpected backend response/,
-    ],
-    [
-      'fails if backend returned requestId but no response',
-      { _id: 'abc', requestType: 'listIndexStats' },
-      {},
-      /Got unexpected backend response/,
-    ],
-  ] as const;
-
-  function getMockFetch(
-    requestResponse: Record<string, unknown> | Error,
-    awaitResponse: Record<string, unknown> | Error
-  ) {
-    return Sinon.stub()
-      .onFirstCall()
-      .callsFake(() => {
-        return requestResponse instanceof Error
-          ? Promise.reject(requestResponse)
-          : Promise.resolve({
-              ok: true,
-              staus: 200,
-              json() {
-                return Promise.resolve(requestResponse);
-              },
-            });
-      })
-      .onSecondCall()
-      .callsFake(() => {
-        return awaitResponse instanceof Error
-          ? Promise.reject(awaitResponse)
-          : Promise.resolve({
-              ok: true,
-              staus: 200,
-              json() {
-                return Promise.resolve(awaitResponse);
-              },
-            });
-      });
-  }
-
-  function getRequestBodyFromFnCall(call: Sinon.SinonSpyCall<any, any>) {
-    return JSON.parse(call.args[1].body);
-  }
-
-  for (const [
-    successSpecName,
-    requestResponse,
-    awaitResponse,
-  ] of successSpecs) {
-    it(successSpecName, async function () {
-      const fetchFn = getMockFetch(requestResponse, awaitResponse);
-      const res = await makeAutomationAgentOpRequest(
-        fetchFn,
-        'http://example.com',
-        'abc',
-        'listIndexStats',
-        { clusterId: 'abc', db: 'db', collection: 'coll' }
-      );
-      expect(getRequestBodyFromFnCall(fetchFn.firstCall)).to.deep.eq({
-        clusterId: 'abc',
-        collection: 'coll',
-        db: 'db',
-      });
-      expect(res).to.deep.eq(awaitResponse.response);
-    });
-  }
-
-  for (const [
-    failSpecName,
-    requestResponse,
-    awaitResponse,
-    errorMessage,
-  ] of failSpecs) {
-    it(failSpecName, async function () {
-      try {
-        await makeAutomationAgentOpRequest(
-          getMockFetch(requestResponse, awaitResponse),
-          'http://example.com',
-          'abc',
-          'listIndexStats',
-          { clusterId: 'abc', db: 'db', collection: 'coll' }
-        );
-        expect.fail('Expected makeAutomationAgentOpRequest call to fail');
-      } catch (err) {
-        expect((err as any).message).to.match(errorMessage);
-      }
-    });
-  }
-});
diff --git a/packages/atlas-service/src/make-automation-agent-op-request.ts b/packages/atlas-service/src/make-automation-agent-op-request.ts
deleted file mode 100644
index 1d0c5c0def1..00000000000
--- a/packages/atlas-service/src/make-automation-agent-op-request.ts
+++ /dev/null
@@ -1,186 +0,0 @@
-type ClusterOrServerlessId =
-  | { serverlessId?: never; clusterId: string }
-  | { serverlessId: string; clusterId?: never };
-
-export type AutomationAgentRequestTypes = {
-  listIndexStats: ClusterOrServerlessId & {
-    db: string;
-    collection: string;
-  };
-  index: ClusterOrServerlessId & {
-    db: string;
-    collection: string;
-    keys: string;
-    options: string;
-    collationOptions: string;
-  };
-  dropIndex: ClusterOrServerlessId & {
-    db: string;
-    collection: string;
-    name: string;
-  };
-};
-
-type AutomationAgentRequestOpTypes = keyof AutomationAgentRequestTypes;
-
-type AutomationAgentRequestResponse<
-  OpType extends AutomationAgentRequestOpTypes
-> = {
-  _id: string;
-  requestType: OpType;
-};
-
-function assertAutomationAgentRequestResponse<
-  OpType extends AutomationAgentRequestOpTypes
->(
-  json: any,
-  opType: OpType
-): asserts json is AutomationAgentRequestResponse<OpType> {
-  if (
-    Object.prototype.hasOwnProperty.call(json, '_id') &&
-    Object.prototype.hasOwnProperty.call(json, 'requestType') &&
-    json.requestType === opType
-  ) {
-    return;
-  }
-  throw new Error(
-    'Got unexpected backend response for automation agent request'
-  );
-}
-
-export type AutomationAgentAwaitResponseTypes = {
-  listIndexStats: {
-    collName: string;
-    dbName: string;
-    indexName: string;
-    indexProperties: { label: string; properties: Record<string, unknown> }[];
-    indexType: { label: string };
-    keys: { name: string; value: string | number };
-    sizeBytes: number;
-    status: 'rolling build' | 'building' | 'exists';
-  }[];
-  dropIndex: never[];
-};
-
-type AutomationAgentAwaitOpTypes = keyof AutomationAgentAwaitResponseTypes;
-
-type AutomationAgentAwaitResponse<OpType extends AutomationAgentAwaitOpTypes> =
-  {
-    _id: string;
-    requestID: string;
-    requestType: OpType;
-    response: AutomationAgentAwaitResponseTypes[OpType];
-    type: OpType;
-  };
-
-function assertAutomationAgentAwaitResponse<
-  OpType extends AutomationAgentAwaitOpTypes
->(
-  json: any,
-  opType: OpType
-): asserts json is AutomationAgentAwaitResponse<OpType> {
-  if (
-    Object.prototype.hasOwnProperty.call(json, '_id') &&
-    Object.prototype.hasOwnProperty.call(json, 'requestType') &&
-    Object.prototype.hasOwnProperty.call(json, 'response') &&
-    json.requestType === opType
-  ) {
-    return;
-  }
-  throw new Error(
-    'Got unexpected backend response for automation agent request await'
-  );
-}
-
-type PickAwaitResponse<OpType extends AutomationAgentAwaitOpTypes> =
-  AutomationAgentAwaitResponse<OpType>['response'];
-
-/**
- * Helper type that maps whatever is returned by automation agent in response
- * prop as follows:
- *
- * empty array -> undefined
- * array with one item -> unwrapped item
- * array of items -> array of items
- */
-export type UnwrappedAutomationAgentAwaitResponse<
-  OpType extends AutomationAgentAwaitOpTypes
-> = PickAwaitResponse<OpType> extends never[]
-  ? undefined
-  : PickAwaitResponse<OpType> extends [infer UnwrappedResponse]
-  ? UnwrappedResponse
-  : PickAwaitResponse<OpType> extends Array<unknown>
-  ? PickAwaitResponse<OpType>
-  : never;
-
-function unwrapAutomationAgentAwaitResponse(
-  json: any,
-  opType: 'listIndexStats'
-): UnwrappedAutomationAgentAwaitResponse<'listIndexStats'>;
-function unwrapAutomationAgentAwaitResponse(
-  json: any,
-  opType: 'dropIndex'
-): UnwrappedAutomationAgentAwaitResponse<'dropIndex'>;
-function unwrapAutomationAgentAwaitResponse(json: any, opType: string): never;
-function unwrapAutomationAgentAwaitResponse(
-  json: any,
-  opType: string
-): unknown {
-  if (opType === 'dropIndex') {
-    assertAutomationAgentAwaitResponse(json, opType);
-    // `dropIndex` returns an empty array, so returning undefined here is just a
-    // bit more explicit than returning `json.response[0]` instead
-    return undefined;
-  }
-  if (opType === 'listIndexStats') {
-    assertAutomationAgentAwaitResponse(json, opType);
-    return json.response;
-  }
-  throw new Error(`Unsupported await response type: ${opType}`);
-}
-
-export type AutomationAgentResponse<
-  OpType extends AutomationAgentRequestOpTypes
-> = OpType extends AutomationAgentAwaitOpTypes
-  ? UnwrappedAutomationAgentAwaitResponse<OpType>
-  : undefined;
-
-async function makeAutomationAgentOpRequest<
-  OpType extends AutomationAgentRequestOpTypes
->(
-  fetchFn: typeof fetch,
-  baseUrl: string,
-  projectId: string,
-  opType: OpType,
-  opBody: AutomationAgentRequestTypes[OpType]
-): Promise<AutomationAgentResponse<OpType>> {
-  const requestUrl =
-    baseUrl + encodeURI(`/explorer/v1/groups/${projectId}/requests/${opType}`);
-  // Tell automation agent to run the op first, this will return the id that we
-  // can use to track the job result
-  const requestRes = await fetchFn(requestUrl, {
-    method: 'POST',
-    headers: {
-      'content-type': 'application/json',
-    },
-    body: JSON.stringify(opBody),
-  });
-  // Rolling index creation request doesn't return anything that we can "await"
-  // on (a successful response is already an acknowledgement that request to
-  // create an index was registered), so we just end here
-  if (opType === 'index') {
-    return undefined as AutomationAgentResponse<OpType>;
-  }
-  const requestJson = await requestRes.json();
-  assertAutomationAgentRequestResponse(requestJson, opType);
-  const awaitUrl =
-    baseUrl +
-    encodeURI(
-      `/explorer/v1/groups/${projectId}/requests/${requestJson._id}/types/${opType}/await`
-    );
-  const awaitRes = await fetchFn(awaitUrl, { method: 'GET' });
-  const awaitJson = await awaitRes.json();
-  return unwrapAutomationAgentAwaitResponse(awaitJson, opType);
-}
-
-export { makeAutomationAgentOpRequest };
diff --git a/packages/compass-indexes/src/modules/rolling-indexes-service.ts b/packages/compass-indexes/src/modules/rolling-indexes-service.ts
index 78e7863fcd1..0311f3859da 100644
--- a/packages/compass-indexes/src/modules/rolling-indexes-service.ts
+++ b/packages/compass-indexes/src/modules/rolling-indexes-service.ts
@@ -3,12 +3,23 @@ import type { ConnectionInfoRef } from '@mongodb-js/compass-connections/provider
 import type { CreateIndexesOptions } from 'mongodb';
 import toNS from 'mongodb-ns';
 
+type AtlasIndexStats = {
+  collName: string;
+  dbName: string;
+  indexName: string;
+  indexProperties: { label: string; properties: Record<string, unknown> }[];
+  indexType: { label: string };
+  keys: { name: string; value: string | number };
+  sizeBytes: number;
+  status: 'rolling build' | 'building' | 'exists';
+};
+
 export class RollingIndexesService {
   constructor(
     private atlasService: AtlasService,
     private connectionInfo: ConnectionInfoRef
   ) {}
-  async listRollingIndexes(namespace: string) {
+  async listRollingIndexes(namespace: string): Promise<AtlasIndexStats[]> {
     const { atlasMetadata } = this.connectionInfo.current;
     if (!atlasMetadata) {
       throw new Error(
@@ -16,16 +27,26 @@ export class RollingIndexesService {
       );
     }
     const { database: db, collection } = toNS(namespace);
-    const indexes = await this.atlasService.automationAgentFetch(
+    const req = await this.atlasService.automationAgentRequest(
       atlasMetadata,
       'listIndexStats',
       { db, collection }
     );
-    return indexes.filter((index) => {
+    if (!req) {
+      throw new Error(
+        'Unexpected response from the automation agent backend: expected to get the request metadata, got undefined'
+      );
+    }
+    const res = await this.atlasService.automationAgentAwait<AtlasIndexStats>(
+      atlasMetadata,
+      req.requestType,
+      req._id
+    );
+    return res.response.filter((index) => {
       return index.status === 'rolling build';
     });
   }
-  createRollingIndex(
+  async createRollingIndex(
     namespace: string,
     indexSpec: Record<string, string | number>,
     { collation, ...options }: CreateIndexesOptions
@@ -37,12 +58,26 @@ export class RollingIndexesService {
       );
     }
     const { database: db, collection } = toNS(namespace);
-    return this.atlasService.automationAgentFetch(atlasMetadata, 'index', {
-      db,
-      collection,
-      keys: JSON.stringify(indexSpec),
-      options: Object.keys(options).length > 0 ? JSON.stringify(options) : '',
-      collationOptions: collation ? JSON.stringify(collation) : '',
-    });
+    const res = await this.atlasService.automationAgentRequest(
+      atlasMetadata,
+      'index',
+      {
+        db,
+        collection,
+        keys: JSON.stringify(indexSpec),
+        options: Object.keys(options).length > 0 ? JSON.stringify(options) : '',
+        collationOptions: collation ? JSON.stringify(collation) : '',
+      }
+    );
+    // Creating a rolling index doesn't return an "awaitable" job from
+    // automation agent backend
+    if (res) {
+      throw new Error(
+        `Unexpected response from the server, expected undefined, but got ${JSON.stringify(
+          res
+        )}`
+      );
+    }
+    return undefined;
   }
 }

From 79c4527ae37749ffdd801a518edff62737747968 Mon Sep 17 00:00:00 2001
From: Sergey Petushkov <petushkov.sergey@gmail.com>
Date: Mon, 30 Sep 2024 18:05:49 +0200
Subject: [PATCH 2/2] chore(indexes): add some tests

---
 .../modules/rolling-indexes-service.spec.ts   | 69 +++++++++++++++++++
 .../src/modules/rolling-indexes-service.ts    |  5 +-
 2 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 packages/compass-indexes/src/modules/rolling-indexes-service.spec.ts

diff --git a/packages/compass-indexes/src/modules/rolling-indexes-service.spec.ts b/packages/compass-indexes/src/modules/rolling-indexes-service.spec.ts
new file mode 100644
index 00000000000..6fd94829d5f
--- /dev/null
+++ b/packages/compass-indexes/src/modules/rolling-indexes-service.spec.ts
@@ -0,0 +1,69 @@
+import Sinon from 'sinon';
+import { RollingIndexesService } from './rolling-indexes-service';
+import { expect } from 'chai';
+
+describe('RollingIndexesService', function () {
+  const atlasServiceStub = {
+    automationAgentRequest: Sinon.stub(),
+    automationAgentAwait: Sinon.stub(),
+  };
+  let service: RollingIndexesService;
+
+  beforeEach(() => {
+    service = new RollingIndexesService(atlasServiceStub, {
+      current: {
+        atlasMetadata: {
+          projectId: 'abc',
+          metricsType: 'cluster',
+          metricsId: '123',
+        },
+      } as any,
+    });
+  });
+
+  afterEach(() => {
+    atlasServiceStub.automationAgentRequest.reset();
+    atlasServiceStub.automationAgentAwait.reset();
+  });
+
+  describe('listRollingIndexes', function () {
+    it('should succeed if automation agent reutrned all responses as expected and filter only the rolling indexes', async function () {
+      atlasServiceStub.automationAgentRequest.resolves({
+        _id: '_id',
+        requestType: 'requestType',
+      });
+      atlasServiceStub.automationAgentAwait.resolves({
+        response: [
+          { indexName: 'abc', status: 'rolling build' },
+          { indexName: 'cba', status: 'exists' },
+        ],
+      });
+      const res = await service.listRollingIndexes('db.coll');
+      expect(res).to.deep.eq([{ indexName: 'abc', status: 'rolling build' }]);
+    });
+
+    it('should fail if automation agent backend returned unexpected result', async function () {
+      atlasServiceStub.automationAgentRequest.resolves(undefined);
+
+      try {
+        await service.listRollingIndexes('db.coll');
+        expect.fail('expected listRollingIndexes to throw');
+      } catch (err) {
+        expect(err).not.to.be.null;
+      }
+    });
+  });
+
+  describe('createRollingIndex', function () {
+    it('should fail if automation agent returned unexpected result', async function () {
+      atlasServiceStub.automationAgentRequest.resolves({ _id: '_id' });
+
+      try {
+        await service.createRollingIndex('db.coll', {}, {});
+        expect.fail('expected createRollingIndex to throw');
+      } catch (err) {
+        expect(err).not.to.be.null;
+      }
+    });
+  });
+});
diff --git a/packages/compass-indexes/src/modules/rolling-indexes-service.ts b/packages/compass-indexes/src/modules/rolling-indexes-service.ts
index 0311f3859da..df89c246bed 100644
--- a/packages/compass-indexes/src/modules/rolling-indexes-service.ts
+++ b/packages/compass-indexes/src/modules/rolling-indexes-service.ts
@@ -16,7 +16,10 @@ type AtlasIndexStats = {
 
 export class RollingIndexesService {
   constructor(
-    private atlasService: AtlasService,
+    private atlasService: Pick<
+      AtlasService,
+      'automationAgentRequest' | 'automationAgentAwait'
+    >,
     private connectionInfo: ConnectionInfoRef
   ) {}
   async listRollingIndexes(namespace: string): Promise<AtlasIndexStats[]> {