opusdev/vector-similarity-api

1import { Readable } from 'stream';2 3import type { Document, ObjectId } from '../bson';4import type { Collection } from '../collection';5import { CursorTimeoutMode } from '../cursor/abstract_cursor';6import type { FindCursor } from '../cursor/find_cursor';7import {8  MongoGridFSChunkError,9  MongoGridFSStreamError,10  MongoInvalidArgumentError,11  MongoRuntimeError12} from '../error';13import type { FindOptions } from '../operations/find';14import type { ReadPreference } from '../read_preference';15import type { Sort } from '../sort';16import { CSOTTimeoutContext } from '../timeout';17import type { Callback } from '../utils';18import type { GridFSChunk } from './upload';19 20/** @public */21export interface GridFSBucketReadStreamOptions {22  sort?: Sort;23  skip?: number;24  /**25   * 0-indexed non-negative byte offset from the beginning of the file26   */27  start?: number;28  /**29   * 0-indexed non-negative byte offset to the end of the file contents30   * to be returned by the stream. `end` is non-inclusive31   */32  end?: number;33  /**34   * @experimental35   * Specifies the time an operation will run until it throws a timeout error36   */37  timeoutMS?: number;38}39 40/** @public */41export interface GridFSBucketReadStreamOptionsWithRevision extends GridFSBucketReadStreamOptions {42  /** The revision number relative to the oldest file with the given filename. 043   * gets you the oldest file, 1 gets you the 2nd oldest, -1 gets you the44   * newest. */45  revision?: number;46}47 48/** @public */49export interface GridFSFile {50  _id: ObjectId;51  length: number;52  chunkSize: number;53  filename: string;54  metadata?: Document;55  uploadDate: Date;56  /** @deprecated Will be removed in the next major version. */57  contentType?: string;58  /** @deprecated Will be removed in the next major version. */59  aliases?: string[];60}61 62/** @internal */63export interface GridFSBucketReadStreamPrivate {64  /**65   * The running total number of bytes read from the chunks collection.66   */67  bytesRead: number;68  /**69   * The number of bytes to remove from the last chunk read in the file.  This is non-zero70   * if `end` is not equal to the length of the document and `end` is not a multiple71   * of the chunkSize.72   */73  bytesToTrim: number;74 75  /**76   * The number of bytes to remove from the first chunk read in the file.  This is non-zero77   * if `start` is not equal to the 0  and `start` is not a multiple78   * of the chunkSize.79   */80  bytesToSkip: number;81 82  files: Collection<GridFSFile>;83  chunks: Collection<GridFSChunk>;84  cursor?: FindCursor<GridFSChunk>;85 86  /** The running total number of chunks read from the chunks collection. */87  expected: number;88 89  /**90   * The filter used to search in the _files_ collection (i.e., `{ _id: <> }`)91   * This is not the same filter used when reading chunks from the chunks collection.92   */93  filter: Document;94 95  /** Indicates whether or not download has started. */96  init: boolean;97 98  /** The expected number of chunks to read, calculated from start, end, chunkSize and file length. */99  expectedEnd: number;100  file?: GridFSFile;101  options: {102    sort?: Sort;103    skip?: number;104    start: number;105    end: number;106    timeoutMS?: number;107  };108  readPreference?: ReadPreference;109  timeoutContext?: CSOTTimeoutContext;110}111 112/**113 * A readable stream that enables you to read buffers from GridFS.114 *115 * Do not instantiate this class directly. Use `openDownloadStream()` instead.116 * @public117 */118export class GridFSBucketReadStream extends Readable {119  /** @internal */120  s: GridFSBucketReadStreamPrivate;121 122  /**123   * Fires when the stream loaded the file document corresponding to the provided id.124   * @event125   */126  static readonly FILE = 'file' as const;127 128  /**129   * @param chunks - Handle for chunks collection130   * @param files - Handle for files collection131   * @param readPreference - The read preference to use132   * @param filter - The filter to use to find the file document133   * @internal134   */135  constructor(136    chunks: Collection<GridFSChunk>,137    files: Collection<GridFSFile>,138    readPreference: ReadPreference | undefined,139    filter: Document,140    options?: GridFSBucketReadStreamOptions141  ) {142    super({ emitClose: true });143    this.s = {144      bytesToTrim: 0,145      bytesToSkip: 0,146      bytesRead: 0,147      chunks,148      expected: 0,149      files,150      filter,151      init: false,152      expectedEnd: 0,153      options: {154        start: 0,155        end: 0,156        ...options157      },158      readPreference,159      timeoutContext:160        options?.timeoutMS != null161          ? new CSOTTimeoutContext({ timeoutMS: options.timeoutMS, serverSelectionTimeoutMS: 0 })162          : undefined163    };164  }165 166  /**167   * Reads from the cursor and pushes to the stream.168   * Private Impl, do not call directly169   * @internal170   */171  override _read(): void {172    if (this.destroyed) return;173    waitForFile(this, () => doRead(this));174  }175 176  /**177   * Sets the 0-based offset in bytes to start streaming from. Throws178   * an error if this stream has entered flowing mode179   * (e.g. if you've already called `on('data')`)180   *181   * @param start - 0-based offset in bytes to start streaming from182   */183  start(start = 0): this {184    throwIfInitialized(this);185    this.s.options.start = start;186    return this;187  }188 189  /**190   * Sets the 0-based offset in bytes to start streaming from. Throws191   * an error if this stream has entered flowing mode192   * (e.g. if you've already called `on('data')`)193   *194   * @param end - Offset in bytes to stop reading at195   */196  end(end = 0): this {197    throwIfInitialized(this);198    this.s.options.end = end;199    return this;200  }201 202  /**203   * Marks this stream as aborted (will never push another `data` event)204   * and kills the underlying cursor. Will emit the 'end' event, and then205   * the 'close' event once the cursor is successfully killed.206   */207  async abort(): Promise<void> {208    this.push(null);209    this.destroy();210    const remainingTimeMS = this.s.timeoutContext?.getRemainingTimeMSOrThrow();211    await this.s.cursor?.close({ timeoutMS: remainingTimeMS });212  }213}214 215function throwIfInitialized(stream: GridFSBucketReadStream): void {216  if (stream.s.init) {217    throw new MongoGridFSStreamError('Options cannot be changed after the stream is initialized');218  }219}220 221function doRead(stream: GridFSBucketReadStream): void {222  if (stream.destroyed) return;223  if (!stream.s.cursor) return;224  if (!stream.s.file) return;225 226  const handleReadResult = (doc: Document | null) => {227    if (stream.destroyed) return;228 229    if (!doc) {230      stream.push(null);231 232      stream.s.cursor?.close().then(undefined, error => stream.destroy(error));233      return;234    }235 236    if (!stream.s.file) return;237 238    const bytesRemaining = stream.s.file.length - stream.s.bytesRead;239    const expectedN = stream.s.expected++;240    const expectedLength = Math.min(stream.s.file.chunkSize, bytesRemaining);241    if (doc.n > expectedN) {242      return stream.destroy(243        new MongoGridFSChunkError(244          `ChunkIsMissing: Got unexpected n: ${doc.n}, expected: ${expectedN}`245        )246      );247    }248 249    if (doc.n < expectedN) {250      return stream.destroy(251        new MongoGridFSChunkError(`ExtraChunk: Got unexpected n: ${doc.n}, expected: ${expectedN}`)252      );253    }254 255    let buf = Buffer.isBuffer(doc.data) ? doc.data : doc.data.buffer;256 257    if (buf.byteLength !== expectedLength) {258      if (bytesRemaining <= 0) {259        return stream.destroy(260          new MongoGridFSChunkError(261            `ExtraChunk: Got unexpected n: ${doc.n}, expected file length ${stream.s.file.length} bytes but already read ${stream.s.bytesRead} bytes`262          )263        );264      }265 266      return stream.destroy(267        new MongoGridFSChunkError(268          `ChunkIsWrongSize: Got unexpected length: ${buf.byteLength}, expected: ${expectedLength}`269        )270      );271    }272 273    stream.s.bytesRead += buf.byteLength;274 275    if (buf.byteLength === 0) {276      return stream.push(null);277    }278 279    let sliceStart = null;280    let sliceEnd = null;281 282    if (stream.s.bytesToSkip != null) {283      sliceStart = stream.s.bytesToSkip;284      stream.s.bytesToSkip = 0;285    }286 287    const atEndOfStream = expectedN === stream.s.expectedEnd - 1;288    const bytesLeftToRead = stream.s.options.end - stream.s.bytesToSkip;289    if (atEndOfStream && stream.s.bytesToTrim != null) {290      sliceEnd = stream.s.file.chunkSize - stream.s.bytesToTrim;291    } else if (stream.s.options.end && bytesLeftToRead < doc.data.byteLength) {292      sliceEnd = bytesLeftToRead;293    }294 295    if (sliceStart != null || sliceEnd != null) {296      buf = buf.slice(sliceStart || 0, sliceEnd || buf.byteLength);297    }298 299    stream.push(buf);300    return;301  };302 303  stream.s.cursor.next().then(handleReadResult, error => {304    if (stream.destroyed) return;305    stream.destroy(error);306  });307}308 309function init(stream: GridFSBucketReadStream): void {310  const findOneOptions: FindOptions = {};311  if (stream.s.readPreference) {312    findOneOptions.readPreference = stream.s.readPreference;313  }314  if (stream.s.options && stream.s.options.sort) {315    findOneOptions.sort = stream.s.options.sort;316  }317  if (stream.s.options && stream.s.options.skip) {318    findOneOptions.skip = stream.s.options.skip;319  }320 321  const handleReadResult = (doc: Document | null) => {322    if (stream.destroyed) return;323 324    if (!doc) {325      const identifier = stream.s.filter._id326        ? stream.s.filter._id.toString()327        : stream.s.filter.filename;328      const errmsg = `FileNotFound: file ${identifier} was not found`;329      // TODO(NODE-3483)330      const err = new MongoRuntimeError(errmsg);331      err.code = 'ENOENT'; // TODO: NODE-3338 set property as part of constructor332      return stream.destroy(err);333    }334 335    // If document is empty, kill the stream immediately and don't336    // execute any reads337    if (doc.length <= 0) {338      stream.push(null);339      return;340    }341 342    if (stream.destroyed) {343      // If user destroys the stream before we have a cursor, wait344      // until the query is done to say we're 'closed' because we can't345      // cancel a query.346      stream.destroy();347      return;348    }349 350    try {351      stream.s.bytesToSkip = handleStartOption(stream, doc, stream.s.options);352    } catch (error) {353      return stream.destroy(error);354    }355 356    const filter: Document = { files_id: doc._id };357 358    // Currently (MongoDB 3.4.4) skip function does not support the index,359    // it needs to retrieve all the documents first and then skip them. (CS-25811)360    // As work around we use $gte on the "n" field.361    if (stream.s.options && stream.s.options.start != null) {362      const skip = Math.floor(stream.s.options.start / doc.chunkSize);363      if (skip > 0) {364        filter['n'] = { $gte: skip };365      }366    }367 368    let remainingTimeMS: number | undefined;369    try {370      remainingTimeMS = stream.s.timeoutContext?.getRemainingTimeMSOrThrow(371        `Download timed out after ${stream.s.timeoutContext?.timeoutMS}ms`372      );373    } catch (error) {374      return stream.destroy(error);375    }376 377    stream.s.cursor = stream.s.chunks378      .find(filter, {379        timeoutMode: stream.s.options.timeoutMS != null ? CursorTimeoutMode.LIFETIME : undefined,380        timeoutMS: remainingTimeMS381      })382      .sort({ n: 1 });383 384    if (stream.s.readPreference) {385      stream.s.cursor.withReadPreference(stream.s.readPreference);386    }387 388    stream.s.expectedEnd = Math.ceil(doc.length / doc.chunkSize);389    stream.s.file = doc as GridFSFile;390 391    try {392      stream.s.bytesToTrim = handleEndOption(stream, doc, stream.s.cursor, stream.s.options);393    } catch (error) {394      return stream.destroy(error);395    }396 397    stream.emit(GridFSBucketReadStream.FILE, doc);398    return;399  };400 401  let remainingTimeMS: number | undefined;402  try {403    remainingTimeMS = stream.s.timeoutContext?.getRemainingTimeMSOrThrow(404      `Download timed out after ${stream.s.timeoutContext?.timeoutMS}ms`405    );406  } catch (error) {407    if (!stream.destroyed) stream.destroy(error);408    return;409  }410 411  findOneOptions.timeoutMS = remainingTimeMS;412 413  stream.s.files.findOne(stream.s.filter, findOneOptions).then(handleReadResult, error => {414    if (stream.destroyed) return;415    stream.destroy(error);416  });417}418 419function waitForFile(stream: GridFSBucketReadStream, callback: Callback): void {420  if (stream.s.file) {421    return callback();422  }423 424  if (!stream.s.init) {425    init(stream);426    stream.s.init = true;427  }428 429  stream.once('file', () => {430    callback();431  });432}433 434function handleStartOption(435  stream: GridFSBucketReadStream,436  doc: Document,437  options: GridFSBucketReadStreamOptions438): number {439  if (options && options.start != null) {440    if (options.start > doc.length) {441      throw new MongoInvalidArgumentError(442        `Stream start (${options.start}) must not be more than the length of the file (${doc.length})`443      );444    }445    if (options.start < 0) {446      throw new MongoInvalidArgumentError(`Stream start (${options.start}) must not be negative`);447    }448    if (options.end != null && options.end < options.start) {449      throw new MongoInvalidArgumentError(450        `Stream start (${options.start}) must not be greater than stream end (${options.end})`451      );452    }453 454    stream.s.bytesRead = Math.floor(options.start / doc.chunkSize) * doc.chunkSize;455    stream.s.expected = Math.floor(options.start / doc.chunkSize);456 457    return options.start - stream.s.bytesRead;458  }459  throw new MongoInvalidArgumentError('Start option must be defined');460}461 462function handleEndOption(463  stream: GridFSBucketReadStream,464  doc: Document,465  cursor: FindCursor<GridFSChunk>,466  options: GridFSBucketReadStreamOptions467) {468  if (options && options.end != null) {469    if (options.end > doc.length) {470      throw new MongoInvalidArgumentError(471        `Stream end (${options.end}) must not be more than the length of the file (${doc.length})`472      );473    }474    if (options.start == null || options.start < 0) {475      throw new MongoInvalidArgumentError(`Stream end (${options.end}) must not be negative`);476    }477 478    const start = options.start != null ? Math.floor(options.start / doc.chunkSize) : 0;479 480    cursor.limit(Math.ceil(options.end / doc.chunkSize) - start);481 482    stream.s.expectedEnd = Math.ceil(options.end / doc.chunkSize);483 484    return Math.ceil(options.end / doc.chunkSize) * doc.chunkSize - options.end;485  }486  throw new MongoInvalidArgumentError('End option must be defined');487}488