feat(shell-api): support hedge/tagSet in cursor.readPref MONGOSH-511

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/AdminServiceProvider.kt

@@ -7,6 +7,7 @@ import org.graalvm.polyglot.Value
  */
 internal interface AdminServiceProvider {
     fun listDatabases(database: String, options: Value?): Value
+    fun readPreferenceFromOptions(options: Value?): Value
     fun getNewConnection(uri: String, options: Value?): Value
     fun getConnectionInfo(): Value
     fun authenticate(authDoc: Value): Value

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/JavaServiceProvider.kt

@@ -447,6 +447,11 @@ internal class JavaServiceProvider(private val client: MongoClient,
         Left(NotImplementedError())
     }
 
+    @HostAccess.Export
+    override fun readPreferenceFromOptions(options: Value?): Value = promise<Any?> {
+        Left(NotImplementedError())
+    }
+
     @HostAccess.Export
     override fun getConnectionInfo(): Value = promise<Any?> {
         Left(NotImplementedError())

packages/java-shell/src/test/resources/cursor/readPref.expected.txt

@@ -1,3 +1 @@
 [ { "_id": <ObjectID>, "a": "a" }, { "_id": <ObjectID>, "a": "A" }, { "_id": <ObjectID>, "a": "\u00E1" } ]
-java.lang.Exception: MongoshUnimplementedError: [COMMON-90002] the tagSet argument is not yet supported.
-java.lang.Exception: MongoshUnimplementedError: [COMMON-90002] the tagSet argument is not yet supported.

packages/java-shell/src/test/resources/cursor/readPref.js

@@ -8,12 +8,12 @@ db.coll.find({a: "a"})
     .collation({"locale": "en_US", strength: 1})
     .readPref('nearest');
 // command
-db.coll.find({a: "a"})
-    .collation({"locale": "en_US", strength: 1})
-    .readPref("secondary", [{"region": "South"}]);
+//db.coll.find({a: "a"})
+//    .collation({"locale": "en_US", strength: 1})
+//    .readPref("secondary", [{"region": "South"}]);
 // command
-db.coll.find({a: "a"})
-    .collation({"locale": "en_US", strength: 1})
-    .readPref("secondary", [{"region": "South", "datacenter": "A"}, {}]);
+//db.coll.find({a: "a"})
+//    .collation({"locale": "en_US", strength: 1})
+//    .readPref("secondary", [{"region": "South", "datacenter": "A"}, {}]);
 // clear
 db.coll.drop();

packages/service-provider-core/src/admin.ts

@@ -92,7 +92,7 @@ export default interface Admin {
    *
    * @param options
    */
-  resetConnectionOptions(options: Document): Promise<void>;
+  resetConnectionOptions(options: MongoClientOptions): Promise<void>;
 
   /**
    * Start a session.

packages/service-provider-core/src/all-transport-types.ts

@@ -35,6 +35,7 @@ export type {
   FindAndModifyOptions,
   FindOperators,
   FindOptions,
+  HedgeOptions,
   IndexDescription,
   InsertManyResult,
   InsertOneOptions,
@@ -50,12 +51,14 @@ export type {
   ReadConcernLevelId,
   ReadPreference,
   ReadPreferenceLike,
+  ReadPreferenceFromOptions,
   ReadPreferenceModeId,
   RenameOptions,
   ReplaceOptions,
   ResumeToken,
   RunCommandOptions,
   ServerSessionId,
+  TagSet,
   TransactionOptions,
   UpdateOptions,
   UpdateResult,

packages/service-provider-core/src/readable.ts

@@ -11,7 +11,9 @@ import type {
   ListIndexesOptions,
   AggregationCursor,
   FindCursor,
-  DbOptions
+  DbOptions,
+  ReadPreferenceFromOptions,
+  ReadPreferenceLike
 } from './all-transport-types';
 import { ChangeStream, ChangeStreamOptions } from './all-transport-types';
 
@@ -199,6 +201,11 @@ export default interface Readable {
     options?: ListCollectionsOptions,
     dbOptions?: DbOptions): Promise<Document[]>;
 
+  /**
+   * Create a ReadPreference object from a set of options
+   */
+  readPreferenceFromOptions(options?: Omit<ReadPreferenceFromOptions, 'session'>): ReadPreferenceLike | undefined;
+
   /**
    * Get all the collection statistics.
    *

packages/service-provider-server/src/cli-service-provider.integration.spec.ts

@@ -96,7 +96,7 @@ describe('CliServiceProvider [integration]', function() {
     it('resets the MongoClient', async() => {
       const mongoClientBefore = serviceProvider.mongoClient;
       await serviceProvider.resetConnectionOptions({
-        readPreference: { mode: 'secondaryPreferred' }
+        readPreference: 'secondaryPreferred'
       });
       expect(serviceProvider.mongoClient).to.not.equal(mongoClientBefore);
       expect(serviceProvider.getReadPreference().mode).to.equal('secondaryPreferred');

packages/service-provider-server/src/cli-service-provider.ts

@@ -16,7 +16,9 @@ import {
   Decimal128,
   BSONSymbol,
   ClientMetadata,
-  Topology
+  Topology,
+  ReadPreferenceFromOptions,
+  ReadPreferenceLike
 } from 'mongodb';
 
 import {
@@ -1104,21 +1106,16 @@ class CliServiceProvider extends ServiceProviderCore implements ServiceProvider
     return this.mongoClient.writeConcern;
   }
 
+  readPreferenceFromOptions(options?: Omit<ReadPreferenceFromOptions, 'session'>): ReadPreferenceLike | undefined {
+    return ReadPreference.fromOptions(options);
+  }
+
   /**
    * For instances where a user wants to set a option that requires a new MongoClient.
    *
    * @param options
    */
-  async resetConnectionOptions(options: Document): Promise<void> {
-    // NOTE: we keep all the original options and just overwrite the passed.
-    if (options.readPreference !== undefined) {
-      const pr = new ReadPreference(
-        options.readPreference.mode,
-        options.readPreference.tagSet,
-        options.hedgeOptions
-      );
-      options.readPreference = pr;
-    }
+  async resetConnectionOptions(options: MongoClientOptions): Promise<void> {
     const clientOptions = processDriverOptions({
       ...this.initialOptions,
       ...options

packages/shell-api/src/cursor.spec.ts

@@ -442,29 +442,32 @@ describe('Cursor', () => {
     describe('#readPref', () => {
       let spCursor: StubbedInstance<ServiceProviderCursor>;
       let shellApiCursor;
+      let fromOptionsStub;
       const value = 'primary';
+      const tagSet = [{ nodeType: 'ANALYTICS' }];
 
       beforeEach(() => {
         spCursor = stubInterface<ServiceProviderCursor>();
         shellApiCursor = new Cursor(mongo, spCursor);
+        fromOptionsStub = sinon.stub();
+        fromOptionsStub.callsFake(input => input);
+        mongo._serviceProvider = {
+          readPreferenceFromOptions: fromOptionsStub
+        };
       });
 
       it('fluidly sets the read preference', () => {
         expect(shellApiCursor.readPref(value)).to.equal(shellApiCursor);
         expect(spCursor.withReadPreference).to.have.been.calledWith(value);
       });
 
-      it('throws MongoshUnimplementedError if tagset is passed', () => {
-        try {
-          shellApiCursor.readPref(value, []);
-          expect.fail('expected error');
-        } catch (e) {
-          expect(e).to.be.instanceOf(MongoshUnimplementedError);
-          expect(e.message).to.contain('the tagSet argument is not yet supported.');
-          expect(e.code).to.equal(CommonErrors.NotImplemented);
-          expect(e.metadata?.driverCaused).to.equal(true);
-          expect(e.metadata?.api).to.equal('Cursor.readPref#tagSet');
-        }
+      it('fluidly sets the read preference with tagSet and hedge options', () => {
+        expect(shellApiCursor.readPref(value, tagSet, { enabled: true })).to.equal(shellApiCursor);
+        expect(spCursor.withReadPreference).to.have.been.calledWith({
+          readPreference: value,
+          readPreferenceTags: tagSet,
+          hedge: { enabled: true }
+        });
       });
     });
 

packages/shell-api/src/cursor.ts

@@ -22,9 +22,10 @@ import {
   CollationOptions,
   ExplainVerbosityLike,
   ReadPreferenceLike,
-  ReadConcernLevelId
+  ReadConcernLevelId,
+  TagSet,
+  HedgeOptions
 } from '@mongosh/service-provider-core';
-import { blockedByDriverMetadata } from './error-codes';
 import { iterate, validateExplainableVerbosity } from './helpers';
 import Mongo from './mongo';
 import { CursorIterationResult } from './result';
@@ -282,16 +283,20 @@ export default class Cursor extends ShellApiClass {
   }
 
   @returnType('Cursor')
-  readPref(mode: ReadPreferenceLike, tagSet?: Document[]): Cursor {
-    if (tagSet) {
-      throw new MongoshUnimplementedError(
-        'the tagSet argument is not yet supported.',
-        CommonErrors.NotImplemented,
-        blockedByDriverMetadata('Cursor.readPref#tagSet')
-      );
+  readPref(mode: ReadPreferenceLike, tagSet?: TagSet[], hedgeOptions?: HedgeOptions): Cursor {
+    let pref: ReadPreferenceLike;
+
+    // Only conditionally use readPreferenceFromOptions, for java-shell compatibility.
+    if (tagSet || hedgeOptions) {
+      pref = this._mongo._serviceProvider.readPreferenceFromOptions({
+        readPreference: mode,
+        readPreferenceTags: tagSet,
+        hedge: hedgeOptions
+      }) as ReadPreferenceLike;
+    } else {
+      pref = mode;
     }
-
-    this._cursor = this._cursor.withReadPreference(mode);
+    this._cursor = this._cursor.withReadPreference(pref);
     return this;
   }
 

packages/shell-api/src/mongo.spec.ts

@@ -369,12 +369,13 @@ describe('Mongo', () => {
     describe('setReadPref', () => {
       it('calls serviceProvider.restConnectionOptions', async() => {
         serviceProvider.resetConnectionOptions.resolves();
+        serviceProvider.readPreferenceFromOptions.callsFake(input => input as any);
         await mongo.setReadPref('primaryPreferred', []);
         expect(serviceProvider.resetConnectionOptions).to.have.been.calledWith({
           readPreference: {
-            mode: 'primaryPreferred',
-            tagSet: [],
-            hedgeOptions: undefined
+            readPreference: 'primaryPreferred',
+            readPreferenceTags: [],
+            hedge: undefined
           }
         });
       });

packages/shell-api/src/mongo.ts

@@ -23,7 +23,9 @@ import {
   ChangeStreamOptions,
   Document,
   generateUri,
+  ReadConcernLevelId,
   ReadPreference,
+  ReadPreferenceLike,
   ReadPreferenceModeId,
   ReplPlatform,
   ServiceProvider,
@@ -225,15 +227,19 @@ export default class Mongo extends ShellApiClass {
   }
 
   @returnsPromise
-  async setReadPref(mode: string, tagSet?: Record<string, string>[], hedgeOptions?: Document): Promise<void> {
+  async setReadPref(mode: ReadPreferenceLike, tagSet?: Record<string, string>[], hedgeOptions?: Document): Promise<void> {
     await this._serviceProvider.resetConnectionOptions({
-      readPreference: { mode: mode, tagSet: tagSet, hedgeOptions: hedgeOptions }
+      readPreference: this._serviceProvider.readPreferenceFromOptions({
+        readPreference: mode,
+        readPreferenceTags: tagSet,
+        hedge: hedgeOptions
+      })
     });
     this._readPreferenceWasExplicitlyRequested = true;
   }
 
   @returnsPromise
-  async setReadConcern(level: string): Promise<void> {
+  async setReadConcern(level: ReadConcernLevelId): Promise<void> {
     await this._serviceProvider.resetConnectionOptions({ readConcern: { level: level } });
   }