opusdev/vector-similarity-api

1import type { Document } from '../bson';2import { type Connection } from '../cmap/connection';3import { MongoDBResponse } from '../cmap/wire_protocol/responses';4import { MongoInvalidArgumentError, MongoServerError } from '../error';5import type { InferIdType } from '../mongo_types';6import type { ClientSession } from '../sessions';7import { formatSort, type Sort, type SortForCmd } from '../sort';8import {9  hasAtomicOperators,10  type MongoDBCollectionNamespace,11  type MongoDBNamespace12} from '../utils';13import { type CollationOptions, CommandOperation, type CommandOperationOptions } from './command';14import { Aspect, defineAspects, type Hint } from './operation';15 16/** @public */17export interface UpdateOptions extends CommandOperationOptions {18  /** A set of filters specifying to which array elements an update should apply */19  arrayFilters?: Document[];20  /** If true, allows the write to opt-out of document level validation */21  bypassDocumentValidation?: boolean;22  /** Specifies a collation */23  collation?: CollationOptions;24  /** Specify that the update query should only consider plans using the hinted index */25  hint?: Hint;26  /** When true, creates a new document if no document matches the query */27  upsert?: boolean;28  /** Map of parameter names and values that can be accessed using $$var (requires MongoDB 5.0). */29  let?: Document;30}31 32/**33 * @public34 * `TSchema` is the schema of the collection35 */36export interface UpdateResult<TSchema extends Document = Document> {37  /** Indicates whether this write result was acknowledged. If not, then all other members of this result will be undefined */38  acknowledged: boolean;39  /** The number of documents that matched the filter */40  matchedCount: number;41  /** The number of documents that were modified */42  modifiedCount: number;43  /** The number of documents that were upserted */44  upsertedCount: number;45  /** The identifier of the inserted document if an upsert took place */46  upsertedId: InferIdType<TSchema> | null;47}48 49/** @public */50export interface UpdateStatement {51  /** The query that matches documents to update. */52  q: Document;53  /** The modifications to apply. */54  u: Document | Document[];55  /**  If true, perform an insert if no documents match the query. */56  upsert?: boolean;57  /** If true, updates all documents that meet the query criteria. */58  multi?: boolean;59  /** Specifies the collation to use for the operation. */60  collation?: CollationOptions;61  /** An array of filter documents that determines which array elements to modify for an update operation on an array field. */62  arrayFilters?: Document[];63  /** A document or string that specifies the index to use to support the query predicate. */64  hint?: Hint;65  /** Specifies the sort order for the documents matched by the filter. */66  sort?: SortForCmd;67}68 69/**70 * @internal71 * UpdateOperation is used in bulk write, while UpdateOneOperation and UpdateManyOperation are only used in the collections API72 */73export class UpdateOperation extends CommandOperation<Document> {74  override SERVER_COMMAND_RESPONSE_TYPE = MongoDBResponse;75  override options: UpdateOptions & { ordered?: boolean };76  statements: UpdateStatement[];77 78  constructor(79    ns: MongoDBNamespace,80    statements: UpdateStatement[],81    options: UpdateOptions & { ordered?: boolean }82  ) {83    super(undefined, options);84    this.options = options;85    this.ns = ns;86 87    this.statements = statements;88  }89 90  override get commandName() {91    return 'update' as const;92  }93 94  override get canRetryWrite(): boolean {95    if (super.canRetryWrite === false) {96      return false;97    }98 99    return this.statements.every(op => op.multi == null || op.multi === false);100  }101 102  override buildCommandDocument(_connection: Connection, _session?: ClientSession): Document {103    const options = this.options;104    const command: Document = {105      update: this.ns.collection,106      updates: this.statements,107      ordered: options.ordered ?? true108    };109 110    if (typeof options.bypassDocumentValidation === 'boolean') {111      command.bypassDocumentValidation = options.bypassDocumentValidation;112    }113 114    if (options.let) {115      command.let = options.let;116    }117 118    // we check for undefined specifically here to allow falsy values119    // eslint-disable-next-line no-restricted-syntax120    if (options.comment !== undefined) {121      command.comment = options.comment;122    }123 124    return command;125  }126}127 128/** @internal */129export class UpdateOneOperation extends UpdateOperation {130  constructor(131    ns: MongoDBCollectionNamespace,132    filter: Document,133    update: Document,134    options: UpdateOptions135  ) {136    super(ns, [makeUpdateStatement(filter, update, { ...options, multi: false })], options);137 138    if (!hasAtomicOperators(update, options)) {139      throw new MongoInvalidArgumentError('Update document requires atomic operators');140    }141  }142 143  override handleOk(144    response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>145  ): UpdateResult {146    const res = super.handleOk(response);147 148    // @ts-expect-error Explain typing is broken149    if (this.explain != null) return res;150 151    if (res.code) throw new MongoServerError(res);152    if (res.writeErrors) throw new MongoServerError(res.writeErrors[0]);153 154    return {155      acknowledged: this.writeConcern?.w !== 0,156      modifiedCount: res.nModified ?? res.n,157      upsertedId:158        Array.isArray(res.upserted) && res.upserted.length > 0 ? res.upserted[0]._id : null,159      upsertedCount: Array.isArray(res.upserted) && res.upserted.length ? res.upserted.length : 0,160      matchedCount: Array.isArray(res.upserted) && res.upserted.length > 0 ? 0 : res.n161    };162  }163}164 165/** @internal */166export class UpdateManyOperation extends UpdateOperation {167  constructor(168    ns: MongoDBCollectionNamespace,169    filter: Document,170    update: Document,171    options: UpdateOptions172  ) {173    super(ns, [makeUpdateStatement(filter, update, { ...options, multi: true })], options);174 175    if (!hasAtomicOperators(update, options)) {176      throw new MongoInvalidArgumentError('Update document requires atomic operators');177    }178  }179 180  override handleOk(181    response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>182  ): UpdateResult {183    const res = super.handleOk(response);184 185    // @ts-expect-error Explain typing is broken186    if (this.explain != null) return res;187    if (res.code) throw new MongoServerError(res);188    if (res.writeErrors) throw new MongoServerError(res.writeErrors[0]);189 190    return {191      acknowledged: this.writeConcern?.w !== 0,192      modifiedCount: res.nModified ?? res.n,193      upsertedId:194        Array.isArray(res.upserted) && res.upserted.length > 0 ? res.upserted[0]._id : null,195      upsertedCount: Array.isArray(res.upserted) && res.upserted.length ? res.upserted.length : 0,196      matchedCount: Array.isArray(res.upserted) && res.upserted.length > 0 ? 0 : res.n197    };198  }199}200 201/** @public */202export interface ReplaceOptions extends CommandOperationOptions {203  /** If true, allows the write to opt-out of document level validation */204  bypassDocumentValidation?: boolean;205  /** Specifies a collation */206  collation?: CollationOptions;207  /** Specify that the update query should only consider plans using the hinted index */208  hint?: string | Document;209  /** When true, creates a new document if no document matches the query */210  upsert?: boolean;211  /** Map of parameter names and values that can be accessed using $$var (requires MongoDB 5.0). */212  let?: Document;213  /** Specifies the sort order for the documents matched by the filter. */214  sort?: Sort;215}216 217/** @internal */218export class ReplaceOneOperation extends UpdateOperation {219  constructor(220    ns: MongoDBCollectionNamespace,221    filter: Document,222    replacement: Document,223    options: ReplaceOptions224  ) {225    super(ns, [makeUpdateStatement(filter, replacement, { ...options, multi: false })], options);226 227    if (hasAtomicOperators(replacement)) {228      throw new MongoInvalidArgumentError('Replacement document must not contain atomic operators');229    }230  }231 232  override handleOk(233    response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>234  ): UpdateResult {235    const res = super.handleOk(response);236 237    // @ts-expect-error Explain typing is broken238    if (this.explain != null) return res;239    if (res.code) throw new MongoServerError(res);240    if (res.writeErrors) throw new MongoServerError(res.writeErrors[0]);241 242    return {243      acknowledged: this.writeConcern?.w !== 0,244      modifiedCount: res.nModified ?? res.n,245      upsertedId:246        Array.isArray(res.upserted) && res.upserted.length > 0 ? res.upserted[0]._id : null,247      upsertedCount: Array.isArray(res.upserted) && res.upserted.length ? res.upserted.length : 0,248      matchedCount: Array.isArray(res.upserted) && res.upserted.length > 0 ? 0 : res.n249    };250  }251}252 253export function makeUpdateStatement(254  filter: Document,255  update: Document | Document[],256  options: UpdateOptions & { multi?: boolean } & { sort?: Sort }257): UpdateStatement {258  if (filter == null || typeof filter !== 'object') {259    throw new MongoInvalidArgumentError('Selector must be a valid JavaScript object');260  }261 262  if (update == null || typeof update !== 'object') {263    throw new MongoInvalidArgumentError('Document must be a valid JavaScript object');264  }265 266  const op: UpdateStatement = { q: filter, u: update };267  if (typeof options.upsert === 'boolean') {268    op.upsert = options.upsert;269  }270 271  if (options.multi) {272    op.multi = options.multi;273  }274 275  if (options.hint) {276    op.hint = options.hint;277  }278 279  if (options.arrayFilters) {280    op.arrayFilters = options.arrayFilters;281  }282 283  if (options.collation) {284    op.collation = options.collation;285  }286 287  if (!options.multi && options.sort != null) {288    op.sort = formatSort(options.sort);289  }290 291  return op;292}293 294defineAspects(UpdateOperation, [295  Aspect.RETRYABLE,296  Aspect.WRITE_OPERATION,297  Aspect.SKIP_COLLATION,298  Aspect.SUPPORTS_RAW_DATA299]);300defineAspects(UpdateOneOperation, [301  Aspect.RETRYABLE,302  Aspect.WRITE_OPERATION,303  Aspect.EXPLAINABLE,304  Aspect.SKIP_COLLATION,305  Aspect.SUPPORTS_RAW_DATA306]);307defineAspects(UpdateManyOperation, [308  Aspect.WRITE_OPERATION,309  Aspect.EXPLAINABLE,310  Aspect.SKIP_COLLATION,311  Aspect.SUPPORTS_RAW_DATA312]);313defineAspects(ReplaceOneOperation, [314  Aspect.RETRYABLE,315  Aspect.WRITE_OPERATION,316  Aspect.SKIP_COLLATION,317  Aspect.SUPPORTS_RAW_DATA318]);319