opusdev/vector-similarity-api

1import { type Connection, type MongoError } from '..';2import { type BSONSerializeOptions, type Document, resolveBSONOptions } from '../bson';3import { type MongoDBResponse } from '../cmap/wire_protocol/responses';4import { type Abortable } from '../mongo_types';5import { ReadPreference, type ReadPreferenceLike } from '../read_preference';6import type { Server, ServerCommandOptions } from '../sdam/server';7import type { ClientSession } from '../sessions';8import { type TimeoutContext } from '../timeout';9import type { MongoDBNamespace } from '../utils';10 11export const Aspect = {12  READ_OPERATION: Symbol('READ_OPERATION'),13  WRITE_OPERATION: Symbol('WRITE_OPERATION'),14  RETRYABLE: Symbol('RETRYABLE'),15  EXPLAINABLE: Symbol('EXPLAINABLE'),16  SKIP_COLLATION: Symbol('SKIP_COLLATION'),17  CURSOR_CREATING: Symbol('CURSOR_CREATING'),18  MUST_SELECT_SAME_SERVER: Symbol('MUST_SELECT_SAME_SERVER'),19  COMMAND_BATCHING: Symbol('COMMAND_BATCHING'),20  SUPPORTS_RAW_DATA: Symbol('SUPPORTS_RAW_DATA')21} as const;22 23/** @public */24export type Hint = string | Document;25 26/** @public */27export interface OperationOptions extends BSONSerializeOptions {28  /** Specify ClientSession for this command */29  session?: ClientSession;30  willRetryWrite?: boolean;31 32  /** The preferred read preference (ReadPreference.primary, ReadPreference.primary_preferred, ReadPreference.secondary, ReadPreference.secondary_preferred, ReadPreference.nearest). */33  readPreference?: ReadPreferenceLike;34 35  /** @internal Hints to `executeOperation` that this operation should not unpin on an ended transaction */36  bypassPinningCheck?: boolean;37 38  /** @internal Hint to `executeOperation` to omit maxTimeMS */39  omitMaxTimeMS?: boolean;40 41  /**42   * @experimental43   * Specifies the time an operation will run until it throws a timeout error44   */45  timeoutMS?: number;46}47 48/**49 * This class acts as a parent class for any operation and is responsible for setting this.options,50 * as well as setting and getting a session.51 * Additionally, this class implements `hasAspect`, which determines whether an operation has52 * a specific aspect.53 * @internal54 */55export abstract class AbstractOperation<TResult = any> {56  ns!: MongoDBNamespace;57  readPreference: ReadPreference;58  server!: Server;59  bypassPinningCheck: boolean;60 61  // BSON serialization options62  bsonOptions?: BSONSerializeOptions;63 64  options: OperationOptions & Abortable;65 66  /** Specifies the time an operation will run until it throws a timeout error. */67  timeoutMS?: number;68 69  private _session: ClientSession | undefined;70 71  static aspects?: Set<symbol>;72 73  constructor(options: OperationOptions & Abortable = {}) {74    this.readPreference = this.hasAspect(Aspect.WRITE_OPERATION)75      ? ReadPreference.primary76      : (ReadPreference.fromOptions(options) ?? ReadPreference.primary);77 78    // Pull the BSON serialize options from the already-resolved options79    this.bsonOptions = resolveBSONOptions(options);80 81    this._session = options.session != null ? options.session : undefined;82 83    this.options = options;84    this.bypassPinningCheck = !!options.bypassPinningCheck;85  }86 87  /** Must match the first key of the command object sent to the server.88  Command name should be stateless (should not use 'this' keyword) */89  abstract get commandName(): string;90 91  hasAspect(aspect: symbol): boolean {92    const ctor = this.constructor as { aspects?: Set<symbol> };93    if (ctor.aspects == null) {94      return false;95    }96 97    return ctor.aspects.has(aspect);98  }99 100  // Make sure the session is not writable from outside this class.101  get session(): ClientSession | undefined {102    return this._session;103  }104 105  set session(session: ClientSession) {106    this._session = session;107  }108 109  clearSession() {110    this._session = undefined;111  }112 113  resetBatch(): boolean {114    return true;115  }116 117  get canRetryRead(): boolean {118    return this.hasAspect(Aspect.RETRYABLE) && this.hasAspect(Aspect.READ_OPERATION);119  }120 121  get canRetryWrite(): boolean {122    return this.hasAspect(Aspect.RETRYABLE) && this.hasAspect(Aspect.WRITE_OPERATION);123  }124  abstract SERVER_COMMAND_RESPONSE_TYPE: typeof MongoDBResponse;125 126  /**127   * Build a raw command document.128   */129  abstract buildCommand(connection: Connection, session?: ClientSession): Document;130 131  /**132   * Builds an instance of `ServerCommandOptions` to be used for operation execution.133   */134  abstract buildOptions(timeoutContext: TimeoutContext): ServerCommandOptions;135 136  /**137   * Given an instance of a MongoDBResponse, map the response to the correct result type.  For138   * example, a `CountOperation` might map the response as follows:139   *140   * ```typescript141   *  override handleOk(response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>): TResult {142   *    return response.toObject(this.bsonOptions).n ?? 0;143   *  }144   *145   *  // or, with type safety:146   *  override handleOk(response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>): TResult {147   *    return response.getNumber('n') ?? 0;148   *  }149   * ```150   */151  handleOk(response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>): TResult {152    return response.toObject(this.bsonOptions) as TResult;153  }154 155  /**156   * Optional.157   *158   * If the operation performs error handling, such as wrapping, renaming the error, or squashing errors159   * this method can be overridden.160   */161  handleError(error: MongoError): TResult | never {162    throw error;163  }164}165 166export function defineAspects(167  operation: { aspects?: Set<symbol> },168  aspects: symbol | symbol[] | Set<symbol>169): Set<symbol> {170  if (!Array.isArray(aspects) && !(aspects instanceof Set)) {171    aspects = [aspects];172  }173 174  aspects = new Set(aspects);175  Object.defineProperty(operation, 'aspects', {176    value: aspects,177    writable: false178  });179 180  return aspects;181}182