opusdev/vector-similarity-api

1import { Writable } from 'stream';2 3import { type Document, ObjectId } from '../bson';4import type { Collection } from '../collection';5import { CursorTimeoutMode } from '../cursor/abstract_cursor';6import {7  MongoAPIError,8  MONGODB_ERROR_CODES,9  MongoError,10  MongoOperationTimeoutError11} from '../error';12import { CSOTTimeoutContext } from '../timeout';13import { type Callback, resolveTimeoutOptions, squashError } from '../utils';14import type { WriteConcernOptions } from '../write_concern';15import { WriteConcern } from './../write_concern';16import type { GridFSFile } from './download';17import type { GridFSBucket } from './index';18 19/** @public */20export interface GridFSChunk {21  _id: ObjectId;22  files_id: ObjectId;23  n: number;24  data: Buffer | Uint8Array;25}26 27/** @public */28export interface GridFSBucketWriteStreamOptions extends WriteConcernOptions {29  /** Overwrite this bucket's chunkSizeBytes for this file */30  chunkSizeBytes?: number;31  /** Custom file id for the GridFS file. */32  id?: ObjectId;33  /** Object to store in the file document's `metadata` field */34  metadata?: Document;35  /**36   * String to store in the file document's `contentType` field.37   * @deprecated Will be removed in the next major version. Add a contentType field to the metadata document instead.38   */39  contentType?: string;40  /**41   * Array of strings to store in the file document's `aliases` field.42   * @deprecated Will be removed in the next major version. Add an aliases field to the metadata document instead.43   */44  aliases?: string[];45  /**46   * @experimental47   * Specifies the time an operation will run until it throws a timeout error48   */49  timeoutMS?: number;50}51 52/**53 * A writable stream that enables you to write buffers to GridFS.54 *55 * Do not instantiate this class directly. Use `openUploadStream()` instead.56 * @public57 */58export class GridFSBucketWriteStream extends Writable {59  bucket: GridFSBucket;60  /** A Collection instance where the file's chunks are stored */61  chunks: Collection<GridFSChunk>;62  /** A Collection instance where the file's GridFSFile document is stored */63  files: Collection<GridFSFile>;64  /** The name of the file */65  filename: string;66  /** Options controlling the metadata inserted along with the file */67  options: GridFSBucketWriteStreamOptions;68  /** Indicates the stream is finished uploading */69  done: boolean;70  /** The ObjectId used for the `_id` field on the GridFSFile document */71  id: ObjectId;72  /** The number of bytes that each chunk will be limited to */73  chunkSizeBytes: number;74  /** Space used to store a chunk currently being inserted */75  bufToStore: Buffer;76  /** Accumulates the number of bytes inserted as the stream uploads chunks */77  length: number;78  /** Accumulates the number of chunks inserted as the stream uploads file contents */79  n: number;80  /** Tracks the current offset into the buffered bytes being uploaded */81  pos: number;82  /** Contains a number of properties indicating the current state of the stream */83  state: {84    /** If set the stream has ended */85    streamEnd: boolean;86    /** Indicates the number of chunks that still need to be inserted to exhaust the current buffered data */87    outstandingRequests: number;88    /** If set an error occurred during insertion */89    errored: boolean;90    /** If set the stream was intentionally aborted */91    aborted: boolean;92  };93  /** The write concern setting to be used with every insert operation */94  writeConcern?: WriteConcern;95  /**96   * The document containing information about the inserted file.97   * This property is defined _after_ the finish event has been emitted.98   * It will remain `null` if an error occurs.99   *100   * @example101   * ```ts102   * fs.createReadStream('file.txt')103   *   .pipe(bucket.openUploadStream('file.txt'))104   *   .on('finish', function () {105   *     console.log(this.gridFSFile)106   *   })107   * ```108   */109  gridFSFile: GridFSFile | null = null;110  /** @internal */111  timeoutContext?: CSOTTimeoutContext;112 113  /**114   * @param bucket - Handle for this stream's corresponding bucket115   * @param filename - The value of the 'filename' key in the files doc116   * @param options - Optional settings.117   * @internal118   */119  constructor(bucket: GridFSBucket, filename: string, options?: GridFSBucketWriteStreamOptions) {120    super();121 122    options = options ?? {};123    this.bucket = bucket;124    this.chunks = bucket.s._chunksCollection;125    this.filename = filename;126    this.files = bucket.s._filesCollection;127    this.options = options;128    this.writeConcern = WriteConcern.fromOptions(options) || bucket.s.options.writeConcern;129    // Signals the write is all done130    this.done = false;131 132    this.id = options.id ? options.id : new ObjectId();133    // properly inherit the default chunksize from parent134    this.chunkSizeBytes = options.chunkSizeBytes || this.bucket.s.options.chunkSizeBytes;135    this.bufToStore = Buffer.alloc(this.chunkSizeBytes);136    this.length = 0;137    this.n = 0;138    this.pos = 0;139    this.state = {140      streamEnd: false,141      outstandingRequests: 0,142      errored: false,143      aborted: false144    };145 146    if (options.timeoutMS != null)147      this.timeoutContext = new CSOTTimeoutContext({148        timeoutMS: options.timeoutMS,149        serverSelectionTimeoutMS: resolveTimeoutOptions(this.bucket.s.db.client, {})150          .serverSelectionTimeoutMS151      });152  }153 154  /**155   * @internal156   *157   * The stream is considered constructed when the indexes are done being created158   */159  override _construct(callback: (error?: Error | null) => void): void {160    if (!this.bucket.s.calledOpenUploadStream) {161      this.bucket.s.calledOpenUploadStream = true;162 163      checkIndexes(this).then(164        () => {165          this.bucket.s.checkedIndexes = true;166          this.bucket.emit('index');167          callback();168        },169        error => {170          if (error instanceof MongoOperationTimeoutError) {171            return handleError(this, error, callback);172          }173          squashError(error);174          callback();175        }176      );177    } else {178      return process.nextTick(callback);179    }180  }181 182  /**183   * @internal184   * Write a buffer to the stream.185   *186   * @param chunk - Buffer to write187   * @param encoding - Optional encoding for the buffer188   * @param callback - Function to call when the chunk was added to the buffer, or if the entire chunk was persisted to MongoDB if this chunk caused a flush.189   */190  override _write(191    chunk: Buffer | string,192    encoding: BufferEncoding,193    callback: Callback<void>194  ): void {195    doWrite(this, chunk, encoding, callback);196  }197 198  /** @internal */199  override _final(callback: (error?: Error | null) => void): void {200    if (this.state.streamEnd) {201      return process.nextTick(callback);202    }203    this.state.streamEnd = true;204    writeRemnant(this, callback);205  }206 207  /**208   * Places this write stream into an aborted state (all future writes fail)209   * and deletes all chunks that have already been written.210   */211  async abort(): Promise<void> {212    if (this.state.streamEnd) {213      // TODO(NODE-3485): Replace with MongoGridFSStreamClosed214      throw new MongoAPIError('Cannot abort a stream that has already completed');215    }216 217    if (this.state.aborted) {218      // TODO(NODE-3485): Replace with MongoGridFSStreamClosed219      throw new MongoAPIError('Cannot call abort() on a stream twice');220    }221 222    this.state.aborted = true;223    const remainingTimeMS = this.timeoutContext?.getRemainingTimeMSOrThrow(224      `Upload timed out after ${this.timeoutContext?.timeoutMS}ms`225    );226 227    await this.chunks.deleteMany({ files_id: this.id }, { timeoutMS: remainingTimeMS });228  }229}230 231function handleError(stream: GridFSBucketWriteStream, error: Error, callback: Callback): void {232  if (stream.state.errored) {233    process.nextTick(callback);234    return;235  }236  stream.state.errored = true;237  process.nextTick(callback, error);238}239 240function createChunkDoc(filesId: ObjectId, n: number, data: Buffer): GridFSChunk {241  return {242    _id: new ObjectId(),243    files_id: filesId,244    n,245    data246  };247}248 249async function checkChunksIndex(stream: GridFSBucketWriteStream): Promise<void> {250  const index = { files_id: 1, n: 1 };251 252  let remainingTimeMS;253  remainingTimeMS = stream.timeoutContext?.getRemainingTimeMSOrThrow(254    `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`255  );256 257  let indexes;258  try {259    indexes = await stream.chunks260      .listIndexes({261        timeoutMode: remainingTimeMS != null ? CursorTimeoutMode.LIFETIME : undefined,262        timeoutMS: remainingTimeMS263      })264      .toArray();265  } catch (error) {266    if (error instanceof MongoError && error.code === MONGODB_ERROR_CODES.NamespaceNotFound) {267      indexes = [];268    } else {269      throw error;270    }271  }272 273  const hasChunksIndex = !!indexes.find(index => {274    const keys = Object.keys(index.key);275    if (keys.length === 2 && index.key.files_id === 1 && index.key.n === 1) {276      return true;277    }278    return false;279  });280 281  if (!hasChunksIndex) {282    remainingTimeMS = stream.timeoutContext?.getRemainingTimeMSOrThrow(283      `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`284    );285    await stream.chunks.createIndex(index, {286      ...stream.writeConcern,287      background: true,288      unique: true,289      timeoutMS: remainingTimeMS290    });291  }292}293 294function checkDone(stream: GridFSBucketWriteStream, callback: Callback): void {295  if (stream.done) {296    return process.nextTick(callback);297  }298 299  if (stream.state.streamEnd && stream.state.outstandingRequests === 0 && !stream.state.errored) {300    // Set done so we do not trigger duplicate createFilesDoc301    stream.done = true;302    // Create a new files doc303    const gridFSFile = createFilesDoc(304      stream.id,305      stream.length,306      stream.chunkSizeBytes,307      stream.filename,308      stream.options.contentType,309      stream.options.aliases,310      stream.options.metadata311    );312 313    if (isAborted(stream, callback)) {314      return;315    }316 317    const remainingTimeMS = stream.timeoutContext?.remainingTimeMS;318    if (remainingTimeMS != null && remainingTimeMS <= 0) {319      return handleError(320        stream,321        new MongoOperationTimeoutError(322          `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`323        ),324        callback325      );326    }327 328    stream.files329      .insertOne(gridFSFile, { writeConcern: stream.writeConcern, timeoutMS: remainingTimeMS })330      .then(331        () => {332          stream.gridFSFile = gridFSFile;333          callback();334        },335        error => {336          return handleError(stream, error, callback);337        }338      );339    return;340  }341 342  process.nextTick(callback);343}344 345async function checkIndexes(stream: GridFSBucketWriteStream): Promise<void> {346  let remainingTimeMS = stream.timeoutContext?.getRemainingTimeMSOrThrow(347    `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`348  );349  const doc = await stream.files.findOne(350    {},351    {352      projection: { _id: 1 },353      timeoutMS: remainingTimeMS354    }355  );356  if (doc != null) {357    // If at least one document exists assume the collection has the required index358    return;359  }360 361  const index = { filename: 1, uploadDate: 1 };362 363  let indexes;364  remainingTimeMS = stream.timeoutContext?.getRemainingTimeMSOrThrow(365    `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`366  );367  const listIndexesOptions = {368    timeoutMode: remainingTimeMS != null ? CursorTimeoutMode.LIFETIME : undefined,369    timeoutMS: remainingTimeMS370  };371  try {372    indexes = await stream.files.listIndexes(listIndexesOptions).toArray();373  } catch (error) {374    if (error instanceof MongoError && error.code === MONGODB_ERROR_CODES.NamespaceNotFound) {375      indexes = [];376    } else {377      throw error;378    }379  }380 381  const hasFileIndex = !!indexes.find(index => {382    const keys = Object.keys(index.key);383    if (keys.length === 2 && index.key.filename === 1 && index.key.uploadDate === 1) {384      return true;385    }386    return false;387  });388 389  if (!hasFileIndex) {390    remainingTimeMS = stream.timeoutContext?.getRemainingTimeMSOrThrow(391      `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`392    );393 394    await stream.files.createIndex(index, { background: false, timeoutMS: remainingTimeMS });395  }396 397  await checkChunksIndex(stream);398}399 400function createFilesDoc(401  _id: ObjectId,402  length: number,403  chunkSize: number,404  filename: string,405  contentType?: string,406  aliases?: string[],407  metadata?: Document408): GridFSFile {409  const ret: GridFSFile = {410    _id,411    length,412    chunkSize,413    uploadDate: new Date(),414    filename415  };416 417  if (contentType) {418    ret.contentType = contentType;419  }420 421  if (aliases) {422    ret.aliases = aliases;423  }424 425  if (metadata) {426    ret.metadata = metadata;427  }428 429  return ret;430}431 432function doWrite(433  stream: GridFSBucketWriteStream,434  chunk: Buffer | string,435  encoding: BufferEncoding,436  callback: Callback<void>437): void {438  if (isAborted(stream, callback)) {439    return;440  }441 442  const inputBuf = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk, encoding);443 444  stream.length += inputBuf.length;445 446  // Input is small enough to fit in our buffer447  if (stream.pos + inputBuf.length < stream.chunkSizeBytes) {448    inputBuf.copy(stream.bufToStore, stream.pos);449    stream.pos += inputBuf.length;450    process.nextTick(callback);451    return;452  }453 454  // Otherwise, buffer is too big for current chunk, so we need to flush455  // to MongoDB.456  let inputBufRemaining = inputBuf.length;457  let spaceRemaining: number = stream.chunkSizeBytes - stream.pos;458  let numToCopy = Math.min(spaceRemaining, inputBuf.length);459  let outstandingRequests = 0;460  while (inputBufRemaining > 0) {461    const inputBufPos = inputBuf.length - inputBufRemaining;462    inputBuf.copy(stream.bufToStore, stream.pos, inputBufPos, inputBufPos + numToCopy);463    stream.pos += numToCopy;464    spaceRemaining -= numToCopy;465    let doc: GridFSChunk;466    if (spaceRemaining === 0) {467      doc = createChunkDoc(stream.id, stream.n, Buffer.from(stream.bufToStore));468 469      const remainingTimeMS = stream.timeoutContext?.remainingTimeMS;470      if (remainingTimeMS != null && remainingTimeMS <= 0) {471        return handleError(472          stream,473          new MongoOperationTimeoutError(474            `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`475          ),476          callback477        );478      }479 480      ++stream.state.outstandingRequests;481      ++outstandingRequests;482 483      if (isAborted(stream, callback)) {484        return;485      }486 487      stream.chunks488        .insertOne(doc, { writeConcern: stream.writeConcern, timeoutMS: remainingTimeMS })489        .then(490          () => {491            --stream.state.outstandingRequests;492            --outstandingRequests;493 494            if (!outstandingRequests) {495              checkDone(stream, callback);496            }497          },498          error => {499            return handleError(stream, error, callback);500          }501        );502 503      spaceRemaining = stream.chunkSizeBytes;504      stream.pos = 0;505      ++stream.n;506    }507    inputBufRemaining -= numToCopy;508    numToCopy = Math.min(spaceRemaining, inputBufRemaining);509  }510}511 512function writeRemnant(stream: GridFSBucketWriteStream, callback: Callback): void {513  // Buffer is empty, so don't bother to insert514  if (stream.pos === 0) {515    return checkDone(stream, callback);516  }517 518  // Create a new buffer to make sure the buffer isn't bigger than it needs519  // to be.520  const remnant = Buffer.alloc(stream.pos);521  stream.bufToStore.copy(remnant, 0, 0, stream.pos);522  const doc = createChunkDoc(stream.id, stream.n, remnant);523 524  // If the stream was aborted, do not write remnant525  if (isAborted(stream, callback)) {526    return;527  }528 529  const remainingTimeMS = stream.timeoutContext?.remainingTimeMS;530  if (remainingTimeMS != null && remainingTimeMS <= 0) {531    return handleError(532      stream,533      new MongoOperationTimeoutError(534        `Upload timed out after ${stream.timeoutContext?.timeoutMS}ms`535      ),536      callback537    );538  }539  ++stream.state.outstandingRequests;540  stream.chunks541    .insertOne(doc, { writeConcern: stream.writeConcern, timeoutMS: remainingTimeMS })542    .then(543      () => {544        --stream.state.outstandingRequests;545        checkDone(stream, callback);546      },547      error => {548        return handleError(stream, error, callback);549      }550    );551}552 553function isAborted(stream: GridFSBucketWriteStream, callback: Callback<void>): boolean {554  if (stream.state.aborted) {555    process.nextTick(callback, new MongoAPIError('Stream has been aborted'));556    return true;557  }558  return false;559}560