opusdev/vector-similarity-api

1import type { ObjectId } from '../bson';2import type { Collection } from '../collection';3import type { FindCursor } from '../cursor/find_cursor';4import type { Db } from '../db';5import { MongoOperationTimeoutError, MongoRuntimeError } from '../error';6import { type Filter, TypedEventEmitter } from '../mongo_types';7import type { ReadPreference } from '../read_preference';8import type { Sort } from '../sort';9import { CSOTTimeoutContext } from '../timeout';10import { noop, resolveOptions } from '../utils';11import { WriteConcern, type WriteConcernOptions } from '../write_concern';12import type { FindOptions } from './../operations/find';13import {14  GridFSBucketReadStream,15  type GridFSBucketReadStreamOptions,16  type GridFSBucketReadStreamOptionsWithRevision,17  type GridFSFile18} from './download';19import {20  GridFSBucketWriteStream,21  type GridFSBucketWriteStreamOptions,22  type GridFSChunk23} from './upload';24 25const DEFAULT_GRIDFS_BUCKET_OPTIONS: {26  bucketName: string;27  chunkSizeBytes: number;28} = {29  bucketName: 'fs',30  chunkSizeBytes: 255 * 102431};32 33/** @public */34export interface GridFSBucketOptions extends WriteConcernOptions {35  /** The 'files' and 'chunks' collections will be prefixed with the bucket name followed by a dot. */36  bucketName?: string;37  /** Number of bytes stored in each chunk. Defaults to 255KB */38  chunkSizeBytes?: number;39  /** Read preference to be passed to read operations */40  readPreference?: ReadPreference;41  /**42   * @experimental43   * Specifies the lifetime duration of a gridFS stream. If any async operations are in progress44   * when this timeout expires, the stream will throw a timeout error.45   */46  timeoutMS?: number;47}48 49/** @internal */50export interface GridFSBucketPrivate {51  db: Db;52  options: {53    bucketName: string;54    chunkSizeBytes: number;55    readPreference?: ReadPreference;56    writeConcern: WriteConcern | undefined;57    timeoutMS?: number;58  };59  _chunksCollection: Collection<GridFSChunk>;60  _filesCollection: Collection<GridFSFile>;61  checkedIndexes: boolean;62  calledOpenUploadStream: boolean;63}64 65/** @public */66export type GridFSBucketEvents = {67  index(): void;68};69 70/**71 * Constructor for a streaming GridFS interface72 * @public73 */74export class GridFSBucket extends TypedEventEmitter<GridFSBucketEvents> {75  /** @internal */76  s: GridFSBucketPrivate;77 78  /**79   * When the first call to openUploadStream is made, the upload stream will80   * check to see if it needs to create the proper indexes on the chunks and81   * files collections. This event is fired either when 1) it determines that82   * no index creation is necessary, 2) when it successfully creates the83   * necessary indexes.84   * @event85   */86  static readonly INDEX = 'index' as const;87 88  constructor(db: Db, options?: GridFSBucketOptions) {89    super();90    this.on('error', noop);91    this.setMaxListeners(0);92    const privateOptions = resolveOptions(db, {93      ...DEFAULT_GRIDFS_BUCKET_OPTIONS,94      ...options,95      writeConcern: WriteConcern.fromOptions(options)96    });97    this.s = {98      db,99      options: privateOptions,100      _chunksCollection: db.collection<GridFSChunk>(privateOptions.bucketName + '.chunks'),101      _filesCollection: db.collection<GridFSFile>(privateOptions.bucketName + '.files'),102      checkedIndexes: false,103      calledOpenUploadStream: false104    };105  }106 107  /**108   * Returns a writable stream (GridFSBucketWriteStream) for writing109   * buffers to GridFS. The stream's 'id' property contains the resulting110   * file's id.111   *112   * @param filename - The value of the 'filename' key in the files doc113   * @param options - Optional settings.114   */115 116  openUploadStream(117    filename: string,118    options?: GridFSBucketWriteStreamOptions119  ): GridFSBucketWriteStream {120    return new GridFSBucketWriteStream(this, filename, {121      timeoutMS: this.s.options.timeoutMS,122      ...options123    });124  }125 126  /**127   * Returns a writable stream (GridFSBucketWriteStream) for writing128   * buffers to GridFS for a custom file id. The stream's 'id' property contains the resulting129   * file's id.130   */131  openUploadStreamWithId(132    id: ObjectId,133    filename: string,134    options?: GridFSBucketWriteStreamOptions135  ): GridFSBucketWriteStream {136    return new GridFSBucketWriteStream(this, filename, {137      timeoutMS: this.s.options.timeoutMS,138      ...options,139      id140    });141  }142 143  /** Returns a readable stream (GridFSBucketReadStream) for streaming file data from GridFS. */144  openDownloadStream(145    id: ObjectId,146    options?: GridFSBucketReadStreamOptions147  ): GridFSBucketReadStream {148    return new GridFSBucketReadStream(149      this.s._chunksCollection,150      this.s._filesCollection,151      this.s.options.readPreference,152      { _id: id },153      { timeoutMS: this.s.options.timeoutMS, ...options }154    );155  }156 157  /**158   * Deletes a file with the given id159   *160   * @param id - The id of the file doc161   */162  async delete(id: ObjectId, options?: { timeoutMS: number }): Promise<void> {163    const { timeoutMS } = resolveOptions(this.s.db, options);164    let timeoutContext: CSOTTimeoutContext | undefined = undefined;165 166    if (timeoutMS) {167      timeoutContext = new CSOTTimeoutContext({168        timeoutMS,169        serverSelectionTimeoutMS: this.s.db.client.s.options.serverSelectionTimeoutMS170      });171    }172 173    const { deletedCount } = await this.s._filesCollection.deleteOne(174      { _id: id },175      { timeoutMS: timeoutContext?.remainingTimeMS }176    );177 178    const remainingTimeMS = timeoutContext?.remainingTimeMS;179    if (remainingTimeMS != null && remainingTimeMS <= 0)180      throw new MongoOperationTimeoutError(`Timed out after ${timeoutMS}ms`);181    // Delete orphaned chunks before returning FileNotFound182    await this.s._chunksCollection.deleteMany({ files_id: id }, { timeoutMS: remainingTimeMS });183 184    if (deletedCount === 0) {185      // TODO(NODE-3483): Replace with more appropriate error186      // Consider creating new error MongoGridFSFileNotFoundError187      throw new MongoRuntimeError(`File not found for id ${id}`);188    }189  }190 191  /** Convenience wrapper around find on the files collection */192  find(filter: Filter<GridFSFile> = {}, options: FindOptions = {}): FindCursor<GridFSFile> {193    return this.s._filesCollection.find(filter, options);194  }195 196  /**197   * Returns a readable stream (GridFSBucketReadStream) for streaming the198   * file with the given name from GridFS. If there are multiple files with199   * the same name, this will stream the most recent file with the given name200   * (as determined by the `uploadDate` field). You can set the `revision`201   * option to change this behavior.202   */203  openDownloadStreamByName(204    filename: string,205    options?: GridFSBucketReadStreamOptionsWithRevision206  ): GridFSBucketReadStream {207    let sort: Sort = { uploadDate: -1 };208    let skip = undefined;209    if (options && options.revision != null) {210      if (options.revision >= 0) {211        sort = { uploadDate: 1 };212        skip = options.revision;213      } else {214        skip = -options.revision - 1;215      }216    }217    return new GridFSBucketReadStream(218      this.s._chunksCollection,219      this.s._filesCollection,220      this.s.options.readPreference,221      { filename },222      { timeoutMS: this.s.options.timeoutMS, ...options, sort, skip }223    );224  }225 226  /**227   * Renames the file with the given _id to the given string228   *229   * @param id - the id of the file to rename230   * @param filename - new name for the file231   */232  async rename(id: ObjectId, filename: string, options?: { timeoutMS: number }): Promise<void> {233    const filter = { _id: id };234    const update = { $set: { filename } };235    const { matchedCount } = await this.s._filesCollection.updateOne(filter, update, options);236    if (matchedCount === 0) {237      throw new MongoRuntimeError(`File with id ${id} not found`);238    }239  }240 241  /** Removes this bucket's files collection, followed by its chunks collection. */242  async drop(options?: { timeoutMS: number }): Promise<void> {243    const { timeoutMS } = resolveOptions(this.s.db, options);244    let timeoutContext: CSOTTimeoutContext | undefined = undefined;245 246    if (timeoutMS) {247      timeoutContext = new CSOTTimeoutContext({248        timeoutMS,249        serverSelectionTimeoutMS: this.s.db.client.s.options.serverSelectionTimeoutMS250      });251    }252 253    if (timeoutContext) {254      await this.s._filesCollection.drop({ timeoutMS: timeoutContext.remainingTimeMS });255      const remainingTimeMS = timeoutContext.getRemainingTimeMSOrThrow(256        `Timed out after ${timeoutMS}ms`257      );258      await this.s._chunksCollection.drop({ timeoutMS: remainingTimeMS });259    } else {260      await this.s._filesCollection.drop();261      await this.s._chunksCollection.drop();262    }263  }264}265