feat(schema): run schema analysis with secondaryPreferred when read pref not set COMPASS-8854 (#6665)

Anemy · web-flow · commit e8229b11ce35 · 2025-02-03T10:07:02.000-05:00
diff --git a/packages/compass-schema/src/modules/schema-analysis.ts b/packages/compass-schema/src/modules/schema-analysis.ts
@@ -64,6 +64,7 @@ export const analyzeSchema = async (
       },
       {
         abortSignal,
+        fallbackReadPreference: 'secondaryPreferred',
       }
     );
     const schemaData = await mongodbSchema(docs);
diff --git a/packages/compass-schema/src/stores/reducer.ts b/packages/compass-schema/src/stores/reducer.ts
@@ -1,5 +1,6 @@
 import type { Schema } from 'mongodb-schema';
 import type { Action, AnyAction, Reducer } from 'redux';
+import type { AggregateOptions } from 'mongodb';
 import { type AnalysisState } from '../constants/analysis-states';
 import {
   ANALYSIS_STATE_ANALYZING,
@@ -250,7 +251,7 @@ export const startAnalysis = (): SchemaThunkAction<
       fields: query.project ?? undefined,
     };
 
-    const driverOptions = {
+    const driverOptions: AggregateOptions = {
       maxTimeMS: capMaxTimeMSAtPreferenceLimit(preferences, query.maxTimeMS),
     };
 
diff --git a/packages/compass-schema/src/stores/store.spec.ts b/packages/compass-schema/src/stores/store.spec.ts
@@ -1,4 +1,5 @@
-import { activateSchemaPlugin, type SchemaStore } from './store';
+import { activateSchemaPlugin } from './store';
+import type { SchemaStore, SchemaPluginServices } from './store';
 import AppRegistry, { createActivateHelpers } from 'hadron-app-registry';
 import { expect } from 'chai';
 
@@ -7,7 +8,6 @@ import { createSandboxFromDefaultPreferences } from 'compass-preferences-model';
 import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
 import type { FieldStoreService } from '@mongodb-js/compass-field-store';
 import { createNoopTrack } from '@mongodb-js/compass-telemetry/provider';
-import type { ConnectionInfoRef } from '@mongodb-js/compass-connections/provider';
 import { startAnalysis, stopAnalysis } from './reducer';
 import Sinon from 'sinon';
 
@@ -26,46 +26,61 @@ const mockQueryBar = {
 };
 
 describe('Schema Store', function () {
-  describe('#configureStore', function () {
-    let store: SchemaStore;
-    let deactivate: () => void;
-    let sandbox: Sinon.SinonSandbox;
-    const localAppRegistry = new AppRegistry();
-    const globalAppRegistry = new AppRegistry();
-    const namespace = 'db.coll';
+  let store: SchemaStore;
+  let deactivate: () => void;
+  let sandbox: Sinon.SinonSandbox;
+  const localAppRegistry = new AppRegistry();
+  const globalAppRegistry = new AppRegistry();
+  const namespace = 'db.coll';
+  let sampleStub: Sinon.SinonStub;
+  let isCancelErrorStub: Sinon.SinonStub;
+
+  beforeEach(function () {
+    sandbox = Sinon.createSandbox();
+    sampleStub = sandbox.stub();
+    isCancelErrorStub = sandbox.stub();
+  });
+
+  async function createStore(services: Partial<SchemaPluginServices> = {}) {
+    const dataService = {
+      sample: sampleStub,
+      isCancelError: isCancelErrorStub,
+    };
     const connectionInfoRef = {
-      current: {},
-    } as ConnectionInfoRef;
-    let sampleStub: Sinon.SinonStub;
-    let isCancelErrorStub: Sinon.SinonStub;
+      current: {
+        connectionOptions: {
+          connectionString: 'mongodb://localhost:27017',
+        },
+        title: 'test',
+        id: 'test',
+      },
+    };
+
+    const plugin = activateSchemaPlugin(
+      {
+        namespace: namespace,
+      },
+      {
+        localAppRegistry: localAppRegistry,
+        globalAppRegistry: globalAppRegistry,
+        dataService,
+        logger: dummyLogger,
+        track: dummyTrack,
+        preferences: await createSandboxFromDefaultPreferences(),
+        fieldStoreService: mockFieldStoreService,
+        queryBar: mockQueryBar as any,
+        connectionInfoRef,
+        ...services,
+      },
+      createActivateHelpers()
+    );
+    store = plugin.store;
+    deactivate = () => plugin.deactivate();
+  }
 
+  describe('#configureStore', function () {
     beforeEach(async function () {
-      sandbox = Sinon.createSandbox();
-      sampleStub = sandbox.stub();
-      isCancelErrorStub = sandbox.stub();
-      const dataService = {
-        sample: sampleStub,
-        isCancelError: isCancelErrorStub,
-      };
-      const plugin = activateSchemaPlugin(
-        {
-          namespace: namespace,
-        },
-        {
-          localAppRegistry: localAppRegistry,
-          globalAppRegistry: globalAppRegistry,
-          dataService: dataService as any,
-          logger: dummyLogger,
-          track: dummyTrack,
-          preferences: await createSandboxFromDefaultPreferences(),
-          fieldStoreService: mockFieldStoreService,
-          queryBar: mockQueryBar as any,
-          connectionInfoRef,
-        },
-        createActivateHelpers()
-      );
-      store = plugin.store;
-      deactivate = () => plugin.deactivate();
+      await createStore();
     });
 
     afterEach(function () {
@@ -107,5 +122,37 @@ describe('Schema Store', function () {
       await analysisPromise;
       expect(store.getState().analysisState).to.equal('initial');
     });
+
+    it('runs the analysis with fallback read pref secondaryPreferred', async function () {
+      sampleStub.resolves([{ name: 'Hans' }, { name: 'Greta' }]);
+      await store.dispatch(startAnalysis());
+      expect(sampleStub.getCall(0).args[3])
+        .property('fallbackReadPreference')
+        .to.equal('secondaryPreferred');
+    });
+  });
+
+  describe('with a connection string with explicit read preference set', function () {
+    beforeEach(async function () {
+      await createStore({
+        connectionInfoRef: {
+          current: {
+            connectionOptions: {
+              connectionString:
+                'mongodb://localhost:27017/?readPreference=primary',
+            },
+            title: 'test',
+            id: 'test',
+          },
+        },
+      });
+    });
+
+    it('does not set read preference to secondaryPreferred', async function () {
+      await store.dispatch(startAnalysis());
+      expect(sampleStub.getCall(0).args[2]).not.to.have.property(
+        'readPreference'
+      );
+    });
   });
 });
diff --git a/packages/data-service/src/data-service.spec.ts b/packages/data-service/src/data-service.spec.ts
@@ -7,6 +7,7 @@ import { Collection, MongoServerError } from 'mongodb';
 import { MongoClient } from 'mongodb';
 import { Int32, UUID } from 'bson';
 import sinon from 'sinon';
+import ConnectionStringUrl from 'mongodb-connection-string-url';
 import type { DataService } from './data-service';
 import { DataServiceImpl } from './data-service';
 import type {
@@ -1044,6 +1045,54 @@ describe('DataService', function () {
           { allowDiskUse: false }
         );
       });
+
+      it('allows to pass fallbackReadPreference and sets the read preference when unset', function () {
+        sandbox.spy(dataService, 'aggregate');
+        void dataService.sample(
+          'db.coll',
+          {},
+          {},
+          {
+            fallbackReadPreference: 'secondaryPreferred',
+          }
+        );
+
+        // eslint-disable-next-line @typescript-eslint/unbound-method
+        expect(dataService.aggregate).to.have.been.calledWith(
+          'db.coll',
+          [{ $sample: { size: 1000 } }],
+          { allowDiskUse: true, readPreference: 'secondaryPreferred' }
+        );
+      });
+
+      it('allows to pass fallbackReadPreference and does not set the read preference when it is already set', function () {
+        sandbox.spy(dataService, 'aggregate');
+        const connectionStringReplacement = new ConnectionStringUrl(
+          cluster().connectionString
+        );
+        connectionStringReplacement.searchParams.set(
+          'readPreference',
+          'primary'
+        );
+        sandbox.replace(dataService as any, '_connectionOptions', {
+          connectionString: connectionStringReplacement.toString(),
+        });
+        void dataService.sample(
+          'db.coll',
+          {},
+          {},
+          {
+            fallbackReadPreference: 'secondaryPreferred',
+          }
+        );
+
+        // eslint-disable-next-line @typescript-eslint/unbound-method
+        expect(dataService.aggregate).to.have.been.calledWith(
+          'db.coll',
+          [{ $sample: { size: 1000 } }],
+          { allowDiskUse: true }
+        );
+      });
     });
 
     describe('#getLastSeenTopology', function () {
diff --git a/packages/data-service/src/data-service.ts b/packages/data-service/src/data-service.ts
@@ -51,6 +51,7 @@ import type {
   ClientEncryptionDataKeyProvider,
   ClientEncryptionCreateDataKeyProviderOptions,
   SearchIndexDescription,
+  ReadPreferenceMode,
 } from 'mongodb';
 import ConnectionStringUrl from 'mongodb-connection-string-url';
 import parseNamespace from 'mongodb-ns';
@@ -126,6 +127,12 @@ function isEmptyObject(obj: Record<string, unknown>) {
   return Object.keys(obj).length === 0;
 }
 
+function isReadPreferenceSet(connectionString: string): boolean {
+  return !!new ConnectionStringUrl(connectionString).searchParams.get(
+    'readPreference'
+  );
+}
+
 let id = 0;
 
 type ClientType = 'CRUD' | 'META';
@@ -665,7 +672,9 @@ export interface DataService {
     ns: string,
     args?: { query?: Filter<Document>; size?: number; fields?: Document },
     options?: AggregateOptions,
-    executionOptions?: ExecutionOptions
+    executionOptions?: ExecutionOptions & {
+      fallbackReadPreference?: ReadPreferenceMode;
+    }
   ): Promise<Document[]>;
 
   /*** Insert ***/
@@ -2219,7 +2228,9 @@ class DataServiceImpl extends WithLogContext implements DataService {
       fields,
     }: { query?: Filter<Document>; size?: number; fields?: Document } = {},
     options: AggregateOptions = {},
-    executionOptions?: ExecutionOptions
+    executionOptions?: ExecutionOptions & {
+      fallbackReadPreference?: ReadPreferenceMode;
+    }
   ): Promise<Document[]> {
     const pipeline = [];
     if (query && Object.keys(query).length > 0) {
@@ -2246,6 +2257,15 @@ class DataServiceImpl extends WithLogContext implements DataService {
       pipeline,
       {
         allowDiskUse: true,
+        // When the read preference isn't set in the connection string explicitly,
+        // then we allow consumers to default to a read preference, for instance
+        // secondaryPreferred to avoid using the primary for analyzing documents.
+        ...(executionOptions?.fallbackReadPreference &&
+        !isReadPreferenceSet(this._connectionOptions.connectionString)
+          ? {
+              readPreference: executionOptions?.fallbackReadPreference,
+            }
+          : {}),
         ...options,
       },
       executionOptions

Original file line number	Diff line number	Diff line change
`@@ -64,6 +64,7 @@ export const analyzeSchema = async (`
`64`	`64`	`},`
`65`	`65`	`{`
`66`	`66`	`abortSignal,`
	`67`	`+ fallbackReadPreference: 'secondaryPreferred',`
`67`	`68`	`}`
`68`	`69`	`);`
`69`	`70`	`const schemaData = await mongodbSchema(docs);`