wip: b4924ee9219475b74b0da53339239c85d6829f30

Author: JesusTheHun

packages/cbjs/src/binding.ts

@@ -529,6 +529,7 @@ export interface CppGetAllReplicasResponseEntry {
 export interface CppGetAllReplicasRequest {
   id: CppDocumentId;
   timeout?: number;
+  read_preference: CppReadPreference;
 }
 export interface CppUpsertResponse {
   // ctx
@@ -572,6 +573,7 @@ export interface CppGetAnyReplicaResponse {
 export interface CppGetAnyReplicaRequest {
   id: CppDocumentId;
   timeout?: number;
+  read_preference: CppReadPreference;
 }
 export interface CppAppendResponse {
   // ctx
@@ -773,7 +775,7 @@ export interface CppLookupInAllReplicasRequest {
   id: CppDocumentId;
   specs: CppImplSubdocCommand[];
   timeout?: number;
-  // parent_span
+  read_preference: CppReadPreference;
 }
 export interface CppAnalyticsResponse {
   // ctx
@@ -1113,7 +1115,7 @@ export interface CppLookupInAnyReplicaRequest {
   id: CppDocumentId;
   specs: CppImplSubdocCommand[];
   timeout?: number;
-  // parent_span
+  read_preference: CppReadPreference;
 }
 export interface CppMutateInResponse {
   // ctx
@@ -3557,6 +3559,13 @@ export interface CppTransaction {
     callback: (err: CppError | null, result: CppTransactionGetResult | null) => void
   ): void;
 
+  getReplicaFromPreferredServerGroup(
+    options: {
+      id: CppDocumentId;
+    },
+    callback: (err: CppError | null, result: CppTransactionGetResult | null) => void
+  ): void;
+
   insert(
     options: {
       id: CppDocumentId;

packages/cbjs/src/bindingutilities.ts

@@ -50,6 +50,7 @@ import binding, {
   CppQueryProfile,
   CppQueryScanConsistency,
   CppRangeScan,
+  CppReadPreference,
   CppReplicateTo,
   CppSamplingScan,
   CppSearchHighlightStyle,
@@ -171,7 +172,12 @@ import {
   EventingFunctionProcessingStatus,
   EventingFunctionStatus,
 } from './eventingfunctionmanager.js';
-import { DurabilityLevel, ServiceType, StoreSemantics } from './generaltypes.js';
+import {
+  DurabilityLevel,
+  ReadPreference,
+  ServiceType,
+  StoreSemantics,
+} from './generaltypes.js';
 import { MutationState } from './mutationstate.js';
 import { QueryProfileMode, QueryScanConsistency } from './querytypes.js';
 import { PrefixScan, RangeScan, SamplingScan } from './rangeScan.js';
@@ -1649,3 +1655,23 @@ export function authDomainFromCpp(domain: CppManagementRbacAuthDomain): string {
   }
   throw new InvalidArgumentError('Unrecognized CppManagementRbacAuthDomain.');
 }
+
+/**
+ * @internal
+ */
+export function readPreferenceToCpp(
+  preference: ReadPreference | undefined
+): CppReadPreference {
+  // Unspecified is allowed, and means no preference.
+  if (preference === null || preference === undefined) {
+    return binding.read_preference.no_preference;
+  }
+
+  if (preference === ReadPreference.NoPreference) {
+    return binding.read_preference.no_preference;
+  } else if (preference === ReadPreference.SelectedServerGroup) {
+    return binding.read_preference.selected_server_group;
+  }
+
+  throw new InvalidArgumentError('Unrecognized ReadPreference.');
+}

packages/cbjs/src/cluster.ts

@@ -259,6 +259,12 @@ export type ConnectOptions<T extends CouchbaseClusterTypes = any> = {
    *
    */
   configProfile?: string;
+
+  /**
+   * Specifies the preferred server group to use for replica operations that specify a non-default
+   * read preference.
+   */
+  preferredServerGroup?: string;
 };
 
 /**
@@ -290,6 +296,7 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
   private _transactions?: Transactions<T>;
   private readonly _openBuckets: Map<BucketName<T>, Promise<void>>;
   private _dnsConfig: DnsConfig | null;
+  private _preferredServerGroup: string | undefined;
 
   /**
    * @internal
@@ -438,11 +445,14 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
     this._bootstrapTimeout = options.timeouts?.bootstrapTimeout;
     this._connectTimeout = options.timeouts?.connectTimeout;
     this._resolveTimeout = options.timeouts?.resolveTimeout;
-
     this._transcoder = options.transcoder ?? new DefaultTranscoder();
     this._queryResultParser = options.queryResultParser ?? JSON.parse;
     this._hooks = options.hooks;
 
+    if (options.preferredServerGroup) {
+      this._preferredServerGroup = options.preferredServerGroup;
+    }
+
     if (options.transactions) {
       this._txnConfig = options.transactions;
     } else {
@@ -913,6 +923,10 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
         dsnObj.options.resolve_timeout = this.resolveTimeout.toString();
       }
 
+      if (this._preferredServerGroup) {
+        dsnObj.options['server_group'] = this._preferredServerGroup;
+      }
+
       const connStr = dsnObj.toString();
 
       const authOpts: CppClusterCredentials = {};

packages/cbjs/src/collection.ts

@@ -69,6 +69,7 @@ import {
   errorFromCpp,
   mutationStateToCpp,
   persistToToCpp,
+  readPreferenceToCpp,
   replicateToToCpp,
   scanTypeToCpp,
   storeSemanticToCpp,
@@ -112,7 +113,7 @@ import {
   DocumentUnretrievableError,
   InvalidArgumentError,
 } from './errors.js';
-import { DurabilityLevel, StoreSemantics } from './generaltypes.js';
+import { DurabilityLevel, ReadPreference, StoreSemantics } from './generaltypes.js';
 import { MutationState } from './mutationstate.js';
 import { CollectionQueryIndexManager } from './queryindexmanager.js';
 import { PrefixScan, RangeScan, SamplingScan } from './rangeScan.js';
@@ -289,6 +290,11 @@ export interface GetAnyReplicaOptions<ThrowIfMissing extends boolean = boolean>
    * @default true
    */
   throwIfMissing?: ThrowIfMissing;
+
+  /**
+   * Specifies how replica nodes will be filtered.
+   */
+  readPreference?: ReadPreference;
 }
 
 /**
@@ -304,6 +310,11 @@ export interface GetAllReplicasOptions {
    * The timeout for this operation, represented in milliseconds.
    */
   timeout?: number;
+
+  /**
+   * Specifies how replica nodes will be filtered.
+   */
+  readPreference?: ReadPreference;
 }
 
 /**
@@ -402,6 +413,11 @@ export interface LookupInAnyReplicaOptions<ThrowOnSpecError extends boolean = fa
    * @default false
    */
   throwOnSpecError?: ThrowOnSpecError;
+
+  /**
+   * Specifies how replica nodes will be filtered.
+   */
+  readPreference?: ReadPreference;
 }
 
 /**
@@ -420,6 +436,11 @@ export interface LookupInAllReplicasOptions<ThrowOnSpecError extends boolean = f
    * @default false
    */
   throwOnSpecError?: ThrowOnSpecError;
+
+  /**
+   * Specifies how replica nodes will be filtered.
+   */
+  readPreference?: ReadPreference;
 }
 
 /**
@@ -857,7 +878,11 @@ export class Collection<
   private _getReplica(
     key: string,
     getAllReplicas: boolean,
-    options: { transcoder?: Transcoder; timeout?: number }
+    options: {
+      transcoder?: Transcoder;
+      timeout?: number;
+      readPreference?: ReadPreference;
+    }
   ): StreamableReplicasPromise<
     [GetReplicaResult<unknown>, ...GetReplicaResult<unknown>[]],
     GetReplicaResult<unknown>
@@ -872,7 +897,11 @@ export class Collection<
   private _getReplica(
     key: string,
     getAllReplicas: boolean,
-    options?: { transcoder?: Transcoder; timeout?: number }
+    options?: {
+      transcoder?: Transcoder;
+      timeout?: number;
+      readPreference?: ReadPreference;
+    }
   ): StreamableReplicasPromise<
     [GetReplicaResult<unknown>, ...GetReplicaResult<unknown>[]],
     GetReplicaResult<unknown>
@@ -893,6 +922,7 @@ export class Collection<
     const request = {
       id: this.getDocId(key),
       timeout: timeout,
+      read_preference: readPreferenceToCpp(options.readPreference),
     };
 
     const getReplicas = getAllReplicas
@@ -2184,6 +2214,7 @@ export class Collection<
         id: this.getDocId(key),
         specs: cppSpecs,
         timeout: timeout,
+        read_preference: readPreferenceToCpp(options.readPreference),
       },
       (cppErr, res) => {
         if (cppErr) {

packages/cbjs/src/generaltypes.ts

@@ -109,3 +109,21 @@ export enum StoreSemantics {
    */
   Insert = 2,
 }
+
+/**
+ * Represents the various scan consistency options that are available when
+ * querying against the query service.
+ */
+export enum ReadPreference {
+  /**
+   * Indicates that filtering for replica set should not be enforced.
+   */
+  NoPreference = 'no_preference',
+
+  /**
+   * Indicates that any nodes that do not belong to local group selected during
+   * cluster instantiation using the `ConnectOptions.preferredServerGroup` option
+   * should be excluded.
+   */
+  SelectedServerGroup = 'selected_server_group',
+}

packages/cbjs/src/transactions.ts

@@ -515,6 +515,16 @@ export interface TransactionGetOptions {
   transcoder?: Transcoder;
 }
 
+/**
+ * @category Transactions
+ */
+export interface TransactionGetReplicaFromPreferredServerGroupOptions {
+  /**
+   * Specifies an explicit transcoder to use for this specific operation.
+   */
+  transcoder?: Transcoder;
+}
+
 /**
  * @category Transactions
  */
@@ -694,6 +704,53 @@ export class TransactionAttemptContext<
     }
   }
 
+  /**
+   * Retrieves the value of a document from the collection.
+   *
+   * @param collection The collection the document lives in.
+   * @param key The document key to retrieve.
+   * @param options Optional parameters for this operation.
+   */
+  async getReplicaFromPreferredServerGroup<
+    LInstance extends AnyCollection,
+    CKS extends CollectionKeyspace<LInstance>,
+    const LKey extends KeyspaceDocDef<
+      T,
+      CKS['bucket'],
+      CKS['scope'],
+      CKS['collection']
+    >['Key'],
+    const Key extends KeyspaceDocDef<T, B, S, C>['Key'],
+  >(
+    ...args: If<
+      IsNever<Instance>,
+      [
+        collection: LInstance,
+        key: LKey,
+        options?: TransactionGetReplicaFromPreferredServerGroupOptions,
+      ],
+      [key: Key, options?: TransactionGetOptions]
+    >
+  ): Promise<TransactionGetResult> {
+    const [collection, key, options] = args as [LInstance, LKey, TransactionGetOptions?];
+    const transcoder = options?.transcoder ?? this.transcoder;
+
+    try {
+      const get = promisify(this._impl.getReplicaFromPreferredServerGroup).bind(
+        this._impl
+      );
+      const id = collection.getDocId(key);
+      const result = await get({ id });
+
+      invariant(result);
+
+      return translateGetResult(result, transcoder) as never;
+    } catch (cppError: unknown) {
+      const err = errorFromCpp(cppError as CppError);
+      throw err;
+    }
+  }
+
   /**
    * Check if a specific document exists in the collection or not.
    *

packages/http-client/src/services/cluster/serverFeature.ts

@@ -48,6 +48,7 @@ export const ServerFeatures = {
   ScopeEventingFunctionManagement: 'scope_eventing_function_management',
   NotLockedKVStatus: 'kv_not_locked',
   BinaryTransactions: 'binary_transactions',
+  ServerGroups: 'server_groups',
 } as const;
 
 export type ServerFeature = (typeof ServerFeatures)[keyof typeof ServerFeatures];
@@ -101,6 +102,8 @@ export function versionSupports(version: string, feature: ServerFeature) {
       return gte(version, '7.6.0');
     case ServerFeatures.BinaryTransactions:
       return gte(version, '7.6.2');
+    case ServerFeatures.ServerGroups:
+      return gte(version, '7.6.2');
   }
 
   throw new Error(`Unknown feature '${feature}'`);

tests/cbjs/tests/kv.get.replica.spec.ts

@@ -21,6 +21,7 @@ import {
   DocumentNotFoundError,
   DocumentUnretrievableError,
   DurabilityLevel,
+  ReadPreference,
 } from '@cbjsdev/cbjs';
 import { getPool } from '@cbjsdev/http-client';
 import { ServerFeatures } from '@cbjsdev/http-client';
@@ -207,4 +208,30 @@ describe
 
       expect(result).toBeUndefined();
     });
+
+    test.runIf(serverSupportsFeatures(ServerFeatures.ServerGroups))(
+      'should throw DocumentUnretrievableError for getAnyReplica with read preference but no cluster preference',
+      async ({ expect, serverTestContext, testDocKey }) => {
+        // the cluster setup does not set a preferred server group, so executing getAnyReplica
+        // with a ReadPreference should fail.
+        await expect(() =>
+          serverTestContext.collection.getAnyReplica(testDocKey, {
+            readPreference: ReadPreference.SelectedServerGroup,
+          })
+        ).rejects.toThrow(DocumentUnretrievableError);
+      }
+    );
+
+    test.runIf(serverSupportsFeatures(ServerFeatures.ServerGroups))(
+      'should throw DocumentUnretrievableError for getAnyReplica with read preference but no cluster preference',
+      async ({ expect, serverTestContext, testDocKey }) => {
+        // the cluster setup does not set a preferred server group, so executing getAllReplicas
+        // with a ReadPreference should fail.
+        await expect(() =>
+          serverTestContext.collection.getAllReplicas(testDocKey, {
+            readPreference: ReadPreference.SelectedServerGroup,
+          })
+        ).rejects.toThrow(DocumentUnretrievableError);
+      }
+    );
   });