opusdev/vector-similarity-api

1import type { Document, Long } from '../bson';2import { type Connection } from '../cmap/connection';3import { CursorResponse } from '../cmap/wire_protocol/responses';4import { MongoRuntimeError } from '../error';5import type { Server, ServerCommandOptions } from '../sdam/server';6import { type TimeoutContext } from '../timeout';7import { maxWireVersion, type MongoDBNamespace } from '../utils';8import { AbstractOperation, Aspect, defineAspects, type OperationOptions } from './operation';9 10/** @internal */11export interface GetMoreOptions extends OperationOptions {12  /** Set the batchSize for the getMoreCommand when iterating over the query results. */13  batchSize?: number;14  /**15   * Comment to apply to the operation.16   *17   * getMore only supports 'comment' in server versions 4.4 and above.18   */19  comment?: unknown;20  /** Number of milliseconds to wait before aborting the query. */21  maxTimeMS?: number;22  /** TODO(NODE-4413): Address bug with maxAwaitTimeMS not being passed in from the cursor correctly */23  maxAwaitTimeMS?: number;24}25 26/**27 * GetMore command: https://www.mongodb.com/docs/manual/reference/command/getMore/28 * @internal29 */30export interface GetMoreCommand {31  getMore: Long;32  collection: string;33  batchSize?: number;34  maxTimeMS?: number;35  /** Only supported on wire versions 10 or greater */36  comment?: unknown;37}38 39/** @internal */40export class GetMoreOperation extends AbstractOperation<CursorResponse> {41  override SERVER_COMMAND_RESPONSE_TYPE = CursorResponse;42  cursorId: Long;43  override options: GetMoreOptions;44 45  constructor(ns: MongoDBNamespace, cursorId: Long, server: Server, options: GetMoreOptions) {46    super(options);47 48    this.options = options;49    this.ns = ns;50    this.cursorId = cursorId;51    this.server = server;52  }53 54  override get commandName() {55    return 'getMore' as const;56  }57 58  override buildCommand(connection: Connection): Document {59    if (this.cursorId == null || this.cursorId.isZero()) {60      throw new MongoRuntimeError('Unable to iterate cursor with no id');61    }62 63    const collection = this.ns.collection;64    if (collection == null) {65      // Cursors should have adopted the namespace returned by MongoDB66      // which should always defined a collection name (even a pseudo one, ex. db.aggregate())67      throw new MongoRuntimeError('A collection name must be determined before getMore');68    }69 70    const getMoreCmd: GetMoreCommand = {71      getMore: this.cursorId,72      collection73    };74 75    if (typeof this.options.batchSize === 'number') {76      getMoreCmd.batchSize = Math.abs(this.options.batchSize);77    }78 79    if (typeof this.options.maxAwaitTimeMS === 'number') {80      getMoreCmd.maxTimeMS = this.options.maxAwaitTimeMS;81    }82 83    // we check for undefined specifically here to allow falsy values84    // eslint-disable-next-line no-restricted-syntax85    if (this.options.comment !== undefined && maxWireVersion(connection) >= 9) {86      getMoreCmd.comment = this.options.comment;87    }88 89    return getMoreCmd;90  }91 92  override buildOptions(timeoutContext: TimeoutContext): ServerCommandOptions {93    return {94      returnFieldSelector: null,95      documentsReturnedIn: 'nextBatch',96      timeoutContext,97      ...this.options98    };99  }100 101  override handleOk(102    response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>103  ): CursorResponse {104    return response;105  }106}107 108defineAspects(GetMoreOperation, [Aspect.READ_OPERATION, Aspect.MUST_SELECT_SAME_SERVER]);109