opusdev/vector-similarity-api

1import type { Document } from '../bson';2import { CursorResponse, ExplainedCursorResponse } from '../cmap/wire_protocol/responses';3import { type CursorTimeoutMode } from '../cursor/abstract_cursor';4import { MongoInvalidArgumentError } from '../error';5import { type ExplainOptions } from '../explain';6import { type MongoDBNamespace } from '../utils';7import { WriteConcern } from '../write_concern';8import { type CollationOptions, CommandOperation, type CommandOperationOptions } from './command';9import { Aspect, defineAspects, type Hint } from './operation';10 11/** @internal */12export const DB_AGGREGATE_COLLECTION = 1 as const;13 14/** @public */15export interface AggregateOptions extends Omit<CommandOperationOptions, 'explain'> {16  /** allowDiskUse lets the server know if it can use disk to store temporary results for the aggregation (requires mongodb 2.6 \>). */17  allowDiskUse?: boolean;18  /** The number of documents to return per batch. See [aggregation documentation](https://www.mongodb.com/docs/manual/reference/command/aggregate). */19  batchSize?: number;20  /** Allow driver to bypass schema validation. */21  bypassDocumentValidation?: boolean;22  /** Return the query as cursor, on 2.6 \> it returns as a real cursor on pre 2.6 it returns as an emulated cursor. */23  cursor?: Document;24  /**25   * Specifies a cumulative time limit in milliseconds for processing operations on the cursor. MongoDB interrupts the operation at the earliest following interrupt point.26   */27  maxTimeMS?: number;28  /** The maximum amount of time for the server to wait on new documents to satisfy a tailable cursor query. */29  maxAwaitTimeMS?: number;30  /** Specify collation. */31  collation?: CollationOptions;32  /** Add an index selection hint to an aggregation command */33  hint?: Hint;34  /** Map of parameter names and values that can be accessed using $$var (requires MongoDB 5.0). */35  let?: Document;36 37  out?: string;38 39  /**40   * Specifies the verbosity mode for the explain output.41   * @deprecated This API is deprecated in favor of `collection.aggregate().explain()`42   * or `db.aggregate().explain()`.43   */44  explain?: ExplainOptions['explain'];45  /** @internal */46  timeoutMode?: CursorTimeoutMode;47}48 49/** @internal */50export class AggregateOperation extends CommandOperation<CursorResponse> {51  override SERVER_COMMAND_RESPONSE_TYPE = CursorResponse;52  override options: AggregateOptions;53  target: string | typeof DB_AGGREGATE_COLLECTION;54  pipeline: Document[];55  hasWriteStage: boolean;56 57  constructor(ns: MongoDBNamespace, pipeline: Document[], options?: AggregateOptions) {58    super(undefined, { ...options, dbName: ns.db });59 60    this.options = { ...options };61 62    // Covers when ns.collection is null, undefined or the empty string, use DB_AGGREGATE_COLLECTION63    this.target = ns.collection || DB_AGGREGATE_COLLECTION;64 65    this.pipeline = pipeline;66 67    // determine if we have a write stage, override read preference if so68    this.hasWriteStage = false;69    if (typeof options?.out === 'string') {70      this.pipeline = this.pipeline.concat({ $out: options.out });71      this.hasWriteStage = true;72    } else if (pipeline.length > 0) {73      const finalStage = pipeline[pipeline.length - 1];74      if (finalStage.$out || finalStage.$merge) {75        this.hasWriteStage = true;76      }77    }78 79    if (!this.hasWriteStage) {80      delete this.options.writeConcern;81    }82 83    if (this.explain && this.writeConcern) {84      throw new MongoInvalidArgumentError(85        'Option "explain" cannot be used on an aggregate call with writeConcern'86      );87    }88 89    if (options?.cursor != null && typeof options.cursor !== 'object') {90      throw new MongoInvalidArgumentError('Cursor options must be an object');91    }92 93    this.SERVER_COMMAND_RESPONSE_TYPE = this.explain ? ExplainedCursorResponse : CursorResponse;94  }95 96  override get commandName() {97    return 'aggregate' as const;98  }99 100  override get canRetryRead(): boolean {101    return !this.hasWriteStage;102  }103 104  addToPipeline(stage: Document): void {105    this.pipeline.push(stage);106  }107 108  override buildCommandDocument(): Document {109    const options = this.options;110    const command: Document = { aggregate: this.target, pipeline: this.pipeline };111 112    if (this.hasWriteStage && this.writeConcern) {113      WriteConcern.apply(command, this.writeConcern);114    }115 116    if (options.bypassDocumentValidation === true) {117      command.bypassDocumentValidation = options.bypassDocumentValidation;118    }119 120    if (typeof options.allowDiskUse === 'boolean') {121      command.allowDiskUse = options.allowDiskUse;122    }123 124    if (options.hint) {125      command.hint = options.hint;126    }127 128    if (options.let) {129      command.let = options.let;130    }131 132    // we check for undefined specifically here to allow falsy values133    // eslint-disable-next-line no-restricted-syntax134    if (options.comment !== undefined) {135      command.comment = options.comment;136    }137 138    command.cursor = options.cursor || {};139    if (options.batchSize && !this.hasWriteStage) {140      command.cursor.batchSize = options.batchSize;141    }142 143    return command;144  }145 146  override handleOk(147    response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>148  ): CursorResponse {149    return response;150  }151}152 153defineAspects(AggregateOperation, [154  Aspect.READ_OPERATION,155  Aspect.RETRYABLE,156  Aspect.EXPLAINABLE,157  Aspect.CURSOR_CREATING,158  Aspect.SUPPORTS_RAW_DATA159]);160