opusdev/vector-similarity-api

1import { type Connection } from '..';2import type { BSONSerializeOptions, Document } from '../bson';3import { MIN_SUPPORTED_RAW_DATA_WIRE_VERSION } from '../cmap/wire_protocol/constants';4import { MongoInvalidArgumentError } from '../error';5import {6  decorateWithExplain,7  Explain,8  type ExplainOptions,9  validateExplainTimeoutOptions10} from '../explain';11import { ReadConcern } from '../read_concern';12import type { ReadPreference } from '../read_preference';13import type { ServerCommandOptions } from '../sdam/server';14import type { ClientSession } from '../sessions';15import { type TimeoutContext } from '../timeout';16import { commandSupportsReadConcern, maxWireVersion, MongoDBNamespace } from '../utils';17import { WriteConcern, type WriteConcernOptions } from '../write_concern';18import type { ReadConcernLike } from './../read_concern';19import { AbstractOperation, Aspect, type OperationOptions } from './operation';20 21/** @public */22export interface CollationOptions {23  locale: string;24  caseLevel?: boolean;25  caseFirst?: string;26  strength?: number;27  numericOrdering?: boolean;28  alternate?: string;29  maxVariable?: string;30  backwards?: boolean;31  normalization?: boolean;32}33 34/** @public */35export interface CommandOperationOptions36  extends OperationOptions,37    WriteConcernOptions,38    ExplainOptions {39  /** Specify a read concern and level for the collection. (only MongoDB 3.2 or higher supported) */40  readConcern?: ReadConcernLike;41  /** Collation */42  collation?: CollationOptions;43  /**44   * maxTimeMS is a server-side time limit in milliseconds for processing an operation.45   */46  maxTimeMS?: number;47  /**48   * Comment to apply to the operation.49   *50   * In server versions pre-4.4, 'comment' must be string.  A server51   * error will be thrown if any other type is provided.52   *53   * In server versions 4.4 and above, 'comment' can be any valid BSON type.54   */55  comment?: unknown;56  /**57   * @deprecated58   * This option is deprecated and will be removed in a future release as it is not used59   * in the driver. Use MongoClientOptions or connection string parameters instead.60   * */61  retryWrites?: boolean;62 63  // Admin command overrides.64  dbName?: string;65  authdb?: string;66  /**67   * @deprecated68   * This option is deprecated and will be removed in an upcoming major version.69   */70  noResponse?: boolean;71 72  /**73   * Used when the command needs to grant access to the underlying namespaces for time series collections.74   * Only available on server versions 8.2 and above and is not meant for public use.75   * @internal76   * @sinceServerVersion 8.277   **/78  rawData?: boolean;79}80 81/** @internal */82export interface OperationParent {83  s: { namespace: MongoDBNamespace };84  readConcern?: ReadConcern;85  writeConcern?: WriteConcern;86  readPreference?: ReadPreference;87  bsonOptions?: BSONSerializeOptions;88  timeoutMS?: number;89}90 91/** @internal */92export abstract class CommandOperation<T> extends AbstractOperation<T> {93  override options: CommandOperationOptions;94  readConcern?: ReadConcern;95  writeConcern?: WriteConcern;96  explain?: Explain;97 98  constructor(parent?: OperationParent, options?: CommandOperationOptions) {99    super(options);100    this.options = options ?? {};101 102    // NOTE: this was explicitly added for the add/remove user operations, it's likely103    //       something we'd want to reconsider. Perhaps those commands can use `Admin`104    //       as a parent?105    const dbNameOverride = options?.dbName || options?.authdb;106    if (dbNameOverride) {107      this.ns = new MongoDBNamespace(dbNameOverride, '$cmd');108    } else {109      this.ns = parent110        ? parent.s.namespace.withCollection('$cmd')111        : new MongoDBNamespace('admin', '$cmd');112    }113 114    this.readConcern = ReadConcern.fromOptions(options);115    this.writeConcern = WriteConcern.fromOptions(options);116 117    if (this.hasAspect(Aspect.EXPLAINABLE)) {118      this.explain = Explain.fromOptions(options);119      if (this.explain) validateExplainTimeoutOptions(this.options, this.explain);120    } else if (options?.explain != null) {121      throw new MongoInvalidArgumentError(`Option "explain" is not supported on this command`);122    }123  }124 125  override get canRetryWrite(): boolean {126    if (this.hasAspect(Aspect.EXPLAINABLE)) {127      return this.explain == null;128    }129    return super.canRetryWrite;130  }131 132  abstract buildCommandDocument(connection: Connection, session?: ClientSession): Document;133 134  override buildOptions(timeoutContext: TimeoutContext): ServerCommandOptions {135    return {136      ...this.options,137      ...this.bsonOptions,138      timeoutContext,139      readPreference: this.readPreference,140      session: this.session141    };142  }143 144  override buildCommand(connection: Connection, session?: ClientSession): Document {145    const command = this.buildCommandDocument(connection, session);146 147    const inTransaction = this.session && this.session.inTransaction();148 149    if (this.readConcern && commandSupportsReadConcern(command) && !inTransaction) {150      Object.assign(command, { readConcern: this.readConcern });151    }152 153    if (this.writeConcern && this.hasAspect(Aspect.WRITE_OPERATION) && !inTransaction) {154      WriteConcern.apply(command, this.writeConcern);155    }156 157    if (158      this.options.collation &&159      typeof this.options.collation === 'object' &&160      !this.hasAspect(Aspect.SKIP_COLLATION)161    ) {162      Object.assign(command, { collation: this.options.collation });163    }164 165    if (typeof this.options.maxTimeMS === 'number') {166      command.maxTimeMS = this.options.maxTimeMS;167    }168 169    if (170      this.options.rawData != null &&171      this.hasAspect(Aspect.SUPPORTS_RAW_DATA) &&172      maxWireVersion(connection) >= MIN_SUPPORTED_RAW_DATA_WIRE_VERSION173    ) {174      command.rawData = this.options.rawData;175    }176 177    if (this.hasAspect(Aspect.EXPLAINABLE) && this.explain) {178      return decorateWithExplain(command, this.explain);179    }180 181    return command;182  }183}184