opusdev/vector-similarity-api

1import {2  isRetryableReadError,3  isRetryableWriteError,4  MongoCompatibilityError,5  MONGODB_ERROR_CODES,6  MongoError,7  MongoErrorLabel,8  MongoExpiredSessionError,9  MongoInvalidArgumentError,10  MongoNetworkError,11  MongoNotConnectedError,12  MongoRuntimeError,13  MongoServerError,14  MongoTransactionError,15  MongoUnexpectedServerResponseError16} from '../error';17import type { MongoClient } from '../mongo_client';18import { ReadPreference } from '../read_preference';19import type { ServerDescription } from '../sdam/server_description';20import {21  sameServerSelector,22  secondaryWritableServerSelector,23  type ServerSelector24} from '../sdam/server_selection';25import type { Topology } from '../sdam/topology';26import type { ClientSession } from '../sessions';27import { TimeoutContext } from '../timeout';28import { abortable, supportsRetryableWrites } from '../utils';29import { AggregateOperation } from './aggregate';30import { AbstractOperation, Aspect } from './operation';31 32const MMAPv1_RETRY_WRITES_ERROR_CODE = MONGODB_ERROR_CODES.IllegalOperation;33const MMAPv1_RETRY_WRITES_ERROR_MESSAGE =34  'This MongoDB deployment does not support retryable writes. Please add retryWrites=false to your connection string.';35 36type ResultTypeFromOperation<TOperation extends AbstractOperation> = ReturnType<37  TOperation['handleOk']38>;39 40/**41 * Executes the given operation with provided arguments.42 * @internal43 *44 * @remarks45 * Allows for a single point of entry to provide features such as implicit sessions, which46 * are required by the Driver Sessions specification in the event that a ClientSession is47 * not provided.48 *49 * The expectation is that this function:50 * - Connects the MongoClient if it has not already been connected, see {@link autoConnect}51 * - Creates a session if none is provided and cleans up the session it creates52 * - Tries an operation and retries under certain conditions, see {@link tryOperation}53 *54 * @typeParam T - The operation's type55 * @typeParam TResult - The type of the operation's result, calculated from T56 *57 * @param client - The MongoClient to execute this operation with58 * @param operation - The operation to execute59 */60export async function executeOperation<61  T extends AbstractOperation,62  TResult = ResultTypeFromOperation<T>63>(client: MongoClient, operation: T, timeoutContext?: TimeoutContext | null): Promise<TResult> {64  if (!(operation instanceof AbstractOperation)) {65    // TODO(NODE-3483): Extend MongoRuntimeError66    throw new MongoRuntimeError('This method requires a valid operation instance');67  }68 69  const topology =70    client.topology == null71      ? await abortable(autoConnect(client), operation.options)72      : client.topology;73 74  // The driver sessions spec mandates that we implicitly create sessions for operations75  // that are not explicitly provided with a session.76  let session = operation.session;77  let owner: symbol | undefined;78 79  if (session == null) {80    owner = Symbol();81    session = client.startSession({ owner, explicit: false });82  } else if (session.hasEnded) {83    throw new MongoExpiredSessionError('Use of expired sessions is not permitted');84  } else if (session.snapshotEnabled && !topology.capabilities.supportsSnapshotReads) {85    throw new MongoCompatibilityError('Snapshot reads require MongoDB 5.0 or later');86  } else if (session.client !== client) {87    throw new MongoInvalidArgumentError('ClientSession must be from the same MongoClient');88  }89 90  operation.session ??= session;91 92  const readPreference = operation.readPreference ?? ReadPreference.primary;93  const inTransaction = !!session?.inTransaction();94 95  const hasReadAspect = operation.hasAspect(Aspect.READ_OPERATION);96 97  if (98    inTransaction &&99    !readPreference.equals(ReadPreference.primary) &&100    (hasReadAspect || operation.commandName === 'runCommand')101  ) {102    throw new MongoTransactionError(103      `Read preference in a transaction must be primary, not: ${readPreference.mode}`104    );105  }106 107  if (session?.isPinned && session.transaction.isCommitted && !operation.bypassPinningCheck) {108    session.unpin();109  }110 111  timeoutContext ??= TimeoutContext.create({112    session,113    serverSelectionTimeoutMS: client.s.options.serverSelectionTimeoutMS,114    waitQueueTimeoutMS: client.s.options.waitQueueTimeoutMS,115    timeoutMS: operation.options.timeoutMS116  });117 118  try {119    return await tryOperation(operation, {120      topology,121      timeoutContext,122      session,123      readPreference124    });125  } finally {126    if (session?.owner != null && session.owner === owner) {127      await session.endSession();128    }129  }130}131 132/**133 * Connects a client if it has not yet been connected134 * @internal135 */136export async function autoConnect(client: MongoClient): Promise<Topology> {137  if (client.topology == null) {138    if (client.s.hasBeenClosed) {139      throw new MongoNotConnectedError('Client must be connected before running operations');140    }141    client.s.options.__skipPingOnConnect = true;142    try {143      await client.connect();144      if (client.topology == null) {145        throw new MongoRuntimeError(146          'client.connect did not create a topology but also did not throw'147        );148      }149      return client.topology;150    } finally {151      delete client.s.options.__skipPingOnConnect;152    }153  }154  return client.topology;155}156 157/** @internal */158type RetryOptions = {159  session: ClientSession | undefined;160  readPreference: ReadPreference;161  topology: Topology;162  timeoutContext: TimeoutContext;163};164 165/**166 * Executes an operation and retries as appropriate167 * @internal168 *169 * @remarks170 * Implements behaviour described in [Retryable Reads](https://github.com/mongodb/specifications/blob/master/source/retryable-reads/retryable-reads.md) and [Retryable171 * Writes](https://github.com/mongodb/specifications/blob/master/source/retryable-writes/retryable-writes.md) specification172 *173 * This function:174 * - performs initial server selection175 * - attempts to execute an operation176 * - retries the operation if it meets the criteria for a retryable read or a retryable write177 *178 * @typeParam T - The operation's type179 * @typeParam TResult - The type of the operation's result, calculated from T180 *181 * @param operation - The operation to execute182 * */183async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFromOperation<T>>(184  operation: T,185  { topology, timeoutContext, session, readPreference }: RetryOptions186): Promise<TResult> {187  let selector: ReadPreference | ServerSelector;188 189  if (operation.hasAspect(Aspect.MUST_SELECT_SAME_SERVER)) {190    // GetMore and KillCursor operations must always select the same server, but run through191    // server selection to potentially force monitor checks if the server is192    // in an unknown state.193    selector = sameServerSelector(operation.server?.description);194  } else if (operation instanceof AggregateOperation && operation.hasWriteStage) {195    // If operation should try to write to secondary use the custom server selector196    // otherwise provide the read preference.197    selector = secondaryWritableServerSelector(topology.commonWireVersion, readPreference);198  } else {199    selector = readPreference;200  }201 202  let server = await topology.selectServer(selector, {203    session,204    operationName: operation.commandName,205    timeoutContext,206    signal: operation.options.signal207  });208 209  const hasReadAspect = operation.hasAspect(Aspect.READ_OPERATION);210  const hasWriteAspect = operation.hasAspect(Aspect.WRITE_OPERATION);211  const inTransaction = session?.inTransaction() ?? false;212 213  const willRetryRead = topology.s.options.retryReads && !inTransaction && operation.canRetryRead;214 215  const willRetryWrite =216    topology.s.options.retryWrites &&217    !inTransaction &&218    supportsRetryableWrites(server) &&219    operation.canRetryWrite;220 221  const willRetry =222    operation.hasAspect(Aspect.RETRYABLE) &&223    session != null &&224    ((hasReadAspect && willRetryRead) || (hasWriteAspect && willRetryWrite));225 226  if (hasWriteAspect && willRetryWrite && session != null) {227    operation.options.willRetryWrite = true;228    session.incrementTransactionNumber();229  }230 231  const maxTries = willRetry ? (timeoutContext.csotEnabled() ? Infinity : 2) : 1;232  let previousOperationError: MongoError | undefined;233  let previousServer: ServerDescription | undefined;234 235  for (let tries = 0; tries < maxTries; tries++) {236    if (previousOperationError) {237      if (hasWriteAspect && previousOperationError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {238        throw new MongoServerError({239          message: MMAPv1_RETRY_WRITES_ERROR_MESSAGE,240          errmsg: MMAPv1_RETRY_WRITES_ERROR_MESSAGE,241          originalError: previousOperationError242        });243      }244 245      if (operation.hasAspect(Aspect.COMMAND_BATCHING) && !operation.canRetryWrite) {246        throw previousOperationError;247      }248 249      if (hasWriteAspect && !isRetryableWriteError(previousOperationError))250        throw previousOperationError;251 252      if (hasReadAspect && !isRetryableReadError(previousOperationError)) {253        throw previousOperationError;254      }255 256      if (257        previousOperationError instanceof MongoNetworkError &&258        operation.hasAspect(Aspect.CURSOR_CREATING) &&259        session != null &&260        session.isPinned &&261        !session.inTransaction()262      ) {263        session.unpin({ force: true, forceClear: true });264      }265 266      server = await topology.selectServer(selector, {267        session,268        operationName: operation.commandName,269        previousServer,270        signal: operation.options.signal271      });272 273      if (hasWriteAspect && !supportsRetryableWrites(server)) {274        throw new MongoUnexpectedServerResponseError(275          'Selected server does not support retryable writes'276        );277      }278    }279 280    operation.server = server;281 282    try {283      // If tries > 0 and we are command batching we need to reset the batch.284      if (tries > 0 && operation.hasAspect(Aspect.COMMAND_BATCHING)) {285        operation.resetBatch();286      }287 288      try {289        const result = await server.command(operation, timeoutContext);290        return operation.handleOk(result);291      } catch (error) {292        return operation.handleError(error);293      }294    } catch (operationError) {295      if (!(operationError instanceof MongoError)) throw operationError;296      if (297        previousOperationError != null &&298        operationError.hasErrorLabel(MongoErrorLabel.NoWritesPerformed)299      ) {300        throw previousOperationError;301      }302      previousServer = server.description;303      previousOperationError = operationError;304 305      // Reset timeouts306      timeoutContext.clear();307    }308  }309 310  throw (311    previousOperationError ??312    new MongoRuntimeError('Tried to propagate retryability error, but no error was found.')313  );314}315