fixup: move the read preference defaulting to the data service

Author: Anemy

packages/compass-schema/src/modules/schema-analysis.ts

@@ -64,6 +64,7 @@ export const analyzeSchema = async (
       },
       {
         abortSignal,
+        fallbackReadPreference: 'secondaryPreferred',
       }
     );
     const schemaData = await mongodbSchema(docs);

packages/compass-schema/src/stores/reducer.ts

@@ -1,6 +1,5 @@
 import type { Schema } from 'mongodb-schema';
 import type { Action, AnyAction, Reducer } from 'redux';
-import ConnectionString from 'mongodb-connection-string-url';
 import type { AggregateOptions } from 'mongodb';
 import { type AnalysisState } from '../constants/analysis-states';
 import {
@@ -214,12 +213,6 @@ export const stopAnalysis = (): SchemaThunkAction<void> => {
   };
 };
 
-function isReadPreferenceSet(connectionString: string): boolean {
-  return !!new ConnectionString(connectionString).searchParams.get(
-    'readPreference'
-  );
-}
-
 export const startAnalysis = (): SchemaThunkAction<
   Promise<void>,
   AnalysisStartedAction | AnalysisFinishedAction | AnalysisFailedAction
@@ -260,16 +253,6 @@ export const startAnalysis = (): SchemaThunkAction<
 
     const driverOptions: AggregateOptions = {
       maxTimeMS: capMaxTimeMSAtPreferenceLimit(preferences, query.maxTimeMS),
-      // When the read preference isn't set in the connection string explicitly,
-      // then we default to secondaryPreferred to avoid using the primary for
-      // schema analysis.
-      ...(isReadPreferenceSet(
-        connectionInfoRef.current.connectionOptions.connectionString
-      )
-        ? {}
-        : {
-            readPreference: 'secondaryPreferred',
-          }),
     };
 
     try {

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

@@ -123,11 +123,11 @@ describe('Schema Store', function () {
       expect(store.getState().analysisState).to.equal('initial');
     });
 
-    it('runs the analysis with read pref secondaryPreferred', async function () {
+    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[2])
-        .property('readPreference')
+      expect(sampleStub.getCall(0).args[3])
+        .property('fallbackReadPreference')
         .to.equal('secondaryPreferred');
     });
   });

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 {
@@ -866,6 +867,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 () {

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';
@@ -122,6 +123,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';
@@ -661,7 +668,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 ***/
@@ -2182,7 +2191,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) {
@@ -2209,6 +2220,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