opusdev/vector-similarity-api

1import { type Connection } from '..';2import type { Document } from '../bson';3import {4  MIN_SUPPORTED_QE_SERVER_VERSION,5  MIN_SUPPORTED_QE_WIRE_VERSION6} from '../cmap/wire_protocol/constants';7import { MongoDBResponse } from '../cmap/wire_protocol/responses';8import { Collection } from '../collection';9import type { Db } from '../db';10import { MongoCompatibilityError } from '../error';11import type { PkFactory } from '../mongo_client';12import type { ClientSession } from '../sessions';13import { TimeoutContext } from '../timeout';14import { maxWireVersion } from '../utils';15import { CommandOperation, type CommandOperationOptions } from './command';16import { executeOperation } from './execute_operation';17import { CreateIndexesOperation } from './indexes';18import { Aspect, defineAspects } from './operation';19 20const ILLEGAL_COMMAND_FIELDS = new Set([21  'w',22  'wtimeout',23  'timeoutMS',24  'j',25  'fsync',26  'autoIndexId',27  'pkFactory',28  'raw',29  'readPreference',30  'session',31  'readConcern',32  'writeConcern',33  'raw',34  'fieldsAsRaw',35  'useBigInt64',36  'promoteLongs',37  'promoteValues',38  'promoteBuffers',39  'bsonRegExp',40  'serializeFunctions',41  'ignoreUndefined',42  'enableUtf8Validation'43]);44 45/** @public46 * Configuration options for timeseries collections47 * @see https://www.mongodb.com/docs/manual/core/timeseries-collections/48 */49export interface TimeSeriesCollectionOptions extends Document {50  timeField: string;51  metaField?: string;52  granularity?: 'seconds' | 'minutes' | 'hours' | string;53  bucketMaxSpanSeconds?: number;54  bucketRoundingSeconds?: number;55}56 57/** @public58 * Configuration options for clustered collections59 * @see https://www.mongodb.com/docs/manual/core/clustered-collections/60 */61export interface ClusteredCollectionOptions extends Document {62  name?: string;63  key: Document;64  unique: boolean;65}66 67/** @public */68export interface CreateCollectionOptions extends Omit<CommandOperationOptions, 'rawData'> {69  /** Create a capped collection */70  capped?: boolean;71  /** @deprecated Create an index on the _id field of the document. This option is deprecated in MongoDB 3.2+ and will be removed once no longer supported by the server. */72  autoIndexId?: boolean;73  /** The size of the capped collection in bytes */74  size?: number;75  /** The maximum number of documents in the capped collection */76  max?: number;77  /** Available for the MMAPv1 storage engine only to set the usePowerOf2Sizes and the noPadding flag */78  flags?: number;79  /** Allows users to specify configuration to the storage engine on a per-collection basis when creating a collection */80  storageEngine?: Document;81  /** Allows users to specify validation rules or expressions for the collection. For more information, see Document Validation */82  validator?: Document;83  /** Determines how strictly MongoDB applies the validation rules to existing documents during an update */84  validationLevel?: string;85  /** Determines whether to error on invalid documents or just warn about the violations but allow invalid documents to be inserted */86  validationAction?: string;87  /** Allows users to specify a default configuration for indexes when creating a collection */88  indexOptionDefaults?: Document;89  /** The name of the source collection or view from which to create the view. The name is not the full namespace of the collection or view (i.e., does not include the database name and implies the same database as the view to create) */90  viewOn?: string;91  /** An array that consists of the aggregation pipeline stage. Creates the view by applying the specified pipeline to the viewOn collection or view */92  pipeline?: Document[];93  /** A primary key factory function for generation of custom _id keys. */94  pkFactory?: PkFactory;95  /** A document specifying configuration options for timeseries collections. */96  timeseries?: TimeSeriesCollectionOptions;97  /** A document specifying configuration options for clustered collections. For MongoDB 5.3 and above. */98  clusteredIndex?: ClusteredCollectionOptions;99  /** The number of seconds after which a document in a timeseries or clustered collection expires. */100  expireAfterSeconds?: number;101  /** @experimental */102  encryptedFields?: Document;103  /**104   * If set, enables pre-update and post-update document events to be included for any105   * change streams that listen on this collection.106   */107  changeStreamPreAndPostImages?: { enabled: boolean };108}109 110/* @internal */111const INVALID_QE_VERSION =112  'Driver support of Queryable Encryption is incompatible with server. Upgrade server to use Queryable Encryption.';113 114/** @internal */115export class CreateCollectionOperation extends CommandOperation<Collection> {116  override SERVER_COMMAND_RESPONSE_TYPE = MongoDBResponse;117  override options: CreateCollectionOptions;118  db: Db;119  name: string;120 121  constructor(db: Db, name: string, options: CreateCollectionOptions = {}) {122    super(db, options);123 124    this.options = options;125    this.db = db;126    this.name = name;127  }128 129  override get commandName() {130    return 'create' as const;131  }132 133  override buildCommandDocument(_connection: Connection, _session?: ClientSession): Document {134    const isOptionValid = ([k, v]: [k: string, v: unknown]) =>135      v != null && typeof v !== 'function' && !ILLEGAL_COMMAND_FIELDS.has(k);136    return {137      create: this.name,138      ...Object.fromEntries(Object.entries(this.options).filter(isOptionValid))139    };140  }141 142  override handleOk(143    _response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>144  ): Collection<Document> {145    return new Collection(this.db, this.name, this.options);146  }147}148 149export async function createCollections<TSchema extends Document>(150  db: Db,151  name: string,152  options: CreateCollectionOptions153): Promise<Collection<TSchema>> {154  const timeoutContext = TimeoutContext.create({155    session: options.session,156    serverSelectionTimeoutMS: db.client.s.options.serverSelectionTimeoutMS,157    waitQueueTimeoutMS: db.client.s.options.waitQueueTimeoutMS,158    timeoutMS: options.timeoutMS159  });160 161  const encryptedFields: Document | undefined =162    options.encryptedFields ??163    db.client.s.options.autoEncryption?.encryptedFieldsMap?.[`${db.databaseName}.${name}`];164 165  if (encryptedFields) {166    class CreateSupportingFLEv2CollectionOperation extends CreateCollectionOperation {167      override buildCommandDocument(connection: Connection, session?: ClientSession): Document {168        if (169          !connection.description.loadBalanced &&170          maxWireVersion(connection) < MIN_SUPPORTED_QE_WIRE_VERSION171        ) {172          throw new MongoCompatibilityError(173            `${INVALID_QE_VERSION} The minimum server version required is ${MIN_SUPPORTED_QE_SERVER_VERSION}`174          );175        }176 177        return super.buildCommandDocument(connection, session);178      }179    }180 181    // Create auxilliary collections for queryable encryption support.182    const escCollection = encryptedFields.escCollection ?? `enxcol_.${name}.esc`;183    const ecocCollection = encryptedFields.ecocCollection ?? `enxcol_.${name}.ecoc`;184 185    for (const collectionName of [escCollection, ecocCollection]) {186      const createOp = new CreateSupportingFLEv2CollectionOperation(db, collectionName, {187        clusteredIndex: {188          key: { _id: 1 },189          unique: true190        },191        session: options.session192      });193      await executeOperation(db.client, createOp, timeoutContext);194    }195 196    if (!options.encryptedFields) {197      options = { ...options, encryptedFields };198    }199  }200 201  const coll = await executeOperation(202    db.client,203    new CreateCollectionOperation(db, name, options),204    timeoutContext205  );206 207  if (encryptedFields) {208    // Create the required index for queryable encryption support.209    const createIndexOp = CreateIndexesOperation.fromIndexSpecification(210      db,211      name,212      { __safeContent__: 1 },213      { session: options.session }214    );215    await executeOperation(db.client, createIndexOp, timeoutContext);216  }217 218  return coll as unknown as Collection<TSchema>;219}220 221defineAspects(CreateCollectionOperation, [Aspect.WRITE_OPERATION]);222