opusdev/vector-similarity-api

1import { type Connection } from '..';2import type { Document } from '../bson';3import { MongoDBResponse } from '../cmap/wire_protocol/responses';4import type { Collection } from '../collection';5import type { ClientSession } from '../sessions';6import type { MongoDBNamespace } from '../utils';7import { CommandOperation, type CommandOperationOptions } from './command';8import { Aspect, defineAspects } from './operation';9 10/** @public */11export interface CountOptions extends CommandOperationOptions {12  /** The number of documents to skip. */13  skip?: number;14  /** The maximum amounts to count before aborting. */15  limit?: number;16  /**17   * Number of milliseconds to wait before aborting the query.18   */19  maxTimeMS?: number;20  /** An index name hint for the query. */21  hint?: string | Document;22}23 24/** @internal */25export class CountOperation extends CommandOperation<number> {26  override SERVER_COMMAND_RESPONSE_TYPE = MongoDBResponse;27  override options: CountOptions;28  collectionName?: string;29  query: Document;30 31  constructor(namespace: MongoDBNamespace, filter: Document, options: CountOptions) {32    super({ s: { namespace: namespace } } as unknown as Collection, options);33 34    this.options = options;35    this.collectionName = namespace.collection;36    this.query = filter;37  }38 39  override get commandName() {40    return 'count' as const;41  }42 43  override buildCommandDocument(_connection: Connection, _session?: ClientSession): Document {44    const options = this.options;45    const cmd: Document = {46      count: this.collectionName,47      query: this.query48    };49 50    if (typeof options.limit === 'number') {51      cmd.limit = options.limit;52    }53 54    if (typeof options.skip === 'number') {55      cmd.skip = options.skip;56    }57 58    if (options.hint != null) {59      cmd.hint = options.hint;60    }61 62    if (typeof options.maxTimeMS === 'number') {63      cmd.maxTimeMS = options.maxTimeMS;64    }65 66    return cmd;67  }68 69  override handleOk(response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>): number {70    return response.getNumber('n') ?? 0;71  }72}73 74defineAspects(CountOperation, [Aspect.READ_OPERATION, Aspect.RETRYABLE, Aspect.SUPPORTS_RAW_DATA]);75