opusdev/vector-similarity-api

1import { type Connection, MongoServerError } from '..';2import type { Document } from '../bson';3import { MongoDBResponse } from '../cmap/wire_protocol/responses';4import { CursorTimeoutContext } from '../cursor/abstract_cursor';5import type { Db } from '../db';6import { MONGODB_ERROR_CODES } from '../error';7import type { ClientSession } from '../sessions';8import { TimeoutContext } from '../timeout';9import { CommandOperation, type CommandOperationOptions } from './command';10import { executeOperation } from './execute_operation';11import { Aspect, defineAspects } from './operation';12 13/** @public */14export interface DropCollectionOptions extends Omit<CommandOperationOptions, 'rawData'> {15  /** @experimental */16  encryptedFields?: Document;17}18 19/** @internal */20export class DropCollectionOperation extends CommandOperation<boolean> {21  override SERVER_COMMAND_RESPONSE_TYPE = MongoDBResponse;22 23  override options: DropCollectionOptions;24  name: string;25 26  constructor(db: Db, name: string, options: DropCollectionOptions = {}) {27    super(db, options);28    this.options = options;29    this.name = name;30  }31 32  override get commandName() {33    return 'drop' as const;34  }35 36  override buildCommandDocument(_connection: Connection, _session?: ClientSession): Document {37    return { drop: this.name };38  }39 40  override handleOk(_response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>): boolean {41    return true;42  }43}44 45export async function dropCollections(46  db: Db,47  name: string,48  options: DropCollectionOptions49): Promise<boolean> {50  const timeoutContext = TimeoutContext.create({51    session: options.session,52    serverSelectionTimeoutMS: db.client.s.options.serverSelectionTimeoutMS,53    waitQueueTimeoutMS: db.client.s.options.waitQueueTimeoutMS,54    timeoutMS: options.timeoutMS55  });56 57  const encryptedFieldsMap = db.client.s.options.autoEncryption?.encryptedFieldsMap;58  let encryptedFields: Document | undefined =59    options.encryptedFields ?? encryptedFieldsMap?.[`${db.databaseName}.${name}`];60 61  if (!encryptedFields && encryptedFieldsMap) {62    // If the MongoClient was configured with an encryptedFieldsMap,63    // and no encryptedFields config was available in it or explicitly64    // passed as an argument, the spec tells us to look one up using65    // listCollections().66    const listCollectionsResult = await db67      .listCollections(68        { name },69        {70          nameOnly: false,71          session: options.session,72          timeoutContext: new CursorTimeoutContext(timeoutContext, Symbol())73        }74      )75      .toArray();76    encryptedFields = listCollectionsResult?.[0]?.options?.encryptedFields;77  }78 79  if (encryptedFields) {80    const escCollection = encryptedFields.escCollection || `enxcol_.${name}.esc`;81    const ecocCollection = encryptedFields.ecocCollection || `enxcol_.${name}.ecoc`;82 83    for (const collectionName of [escCollection, ecocCollection]) {84      // Drop auxilliary collections, ignoring potential NamespaceNotFound errors.85      const dropOp = new DropCollectionOperation(db, collectionName, options);86      try {87        await executeOperation(db.client, dropOp, timeoutContext);88      } catch (err) {89        if (90          !(err instanceof MongoServerError) ||91          err.code !== MONGODB_ERROR_CODES.NamespaceNotFound92        ) {93          throw err;94        }95      }96    }97  }98 99  return await executeOperation(100    db.client,101    new DropCollectionOperation(db, name, options),102    timeoutContext103  );104}105 106/** @public */107export type DropDatabaseOptions = CommandOperationOptions;108 109/** @internal */110export class DropDatabaseOperation extends CommandOperation<boolean> {111  override SERVER_COMMAND_RESPONSE_TYPE = MongoDBResponse;112  override options: DropDatabaseOptions;113 114  constructor(db: Db, options: DropDatabaseOptions) {115    super(db, options);116    this.options = options;117  }118  override get commandName() {119    return 'dropDatabase' as const;120  }121 122  override buildCommandDocument(_connection: Connection, _session?: ClientSession): Document {123    return { dropDatabase: 1 };124  }125 126  override handleOk(_response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>): boolean {127    return true;128  }129}130 131defineAspects(DropCollectionOperation, [Aspect.WRITE_OPERATION]);132defineAspects(DropDatabaseOperation, [Aspect.WRITE_OPERATION]);133