refactor(NODE-6055): implement OnDemandDocument

package.json

@@ -26,7 +26,7 @@
   },
   "dependencies": {
     "@mongodb-js/saslprep": "^1.1.5",
-    "bson": "^6.5.0",
+    "bson": "^6.6.0",
     "mongodb-connection-string-url": "^3.0.0"
   },
   "peerDependencies": {

src/bson.ts

@@ -1,8 +1,10 @@
 import type { DeserializeOptions, SerializeOptions } from 'bson';
+import { BSON } from 'bson';
 
 export {
   Binary,
   BSON,
+  BSONError,
   BSONRegExp,
   BSONSymbol,
   BSONType,
@@ -25,6 +27,17 @@ export {
   UUID
 } from 'bson';
 
+export type BSONElement = BSON.OnDemand['BSONElement'];
+
+export function parseToElementsToArray(bytes: Uint8Array, offset?: number): BSONElement[] {
+  const res = BSON.onDemand.parseToElements(bytes, offset);
+  return Array.isArray(res) ? res : [...res];
+}
+export const getInt32LE = BSON.onDemand.NumberUtils.getInt32LE;
+export const getFloat64LE = BSON.onDemand.NumberUtils.getFloat64LE;
+export const getBigInt64LE = BSON.onDemand.NumberUtils.getBigInt64LE;
+export const toUTF8 = BSON.onDemand.ByteUtils.toUTF8;
+
 /**
  * BSON Serialization options.
  * @public

src/cmap/wire_protocol/on_demand/document.ts

@@ -0,0 +1,322 @@
+import {
+  Binary,
+  BSON,
+  type BSONElement,
+  BSONError,
+  type BSONSerializeOptions,
+  BSONType,
+  getBigInt64LE,
+  getFloat64LE,
+  getInt32LE,
+  ObjectId,
+  parseToElementsToArray,
+  Timestamp,
+  toUTF8
+} from '../../../bson';
+
+// eslint-disable-next-line no-restricted-syntax
+const enum BSONElementOffset {
+  type = 0,
+  nameOffset = 1,
+  nameLength = 2,
+  offset = 3,
+  length = 4
+}
+
+export type JSTypeOf = {
+  [BSONType.null]: null;
+  [BSONType.undefined]: null;
+  [BSONType.double]: number;
+  [BSONType.int]: number;
+  [BSONType.long]: bigint;
+  [BSONType.timestamp]: Timestamp;
+  [BSONType.binData]: Binary;
+  [BSONType.bool]: boolean;
+  [BSONType.objectId]: ObjectId;
+  [BSONType.string]: string;
+  [BSONType.date]: Date;
+  [BSONType.object]: OnDemandDocument;
+  [BSONType.array]: OnDemandDocument;
+};
+
+/** @internal */
+type CachedBSONElement = { element: BSONElement; value: any | undefined };
+
+/** @internal */
+export class OnDemandDocument {
+  /**
+   * Maps JS strings to elements and jsValues for speeding up subsequent lookups.
+   * - If `false` then name does not exist in the BSON document
+   * - If `CachedBSONElement` instance name exists
+   * - If `cache[name].value == null` jsValue has not yet been parsed
+   *   - Null/Undefined values do not get cached because they are zero-length values.
+   */
+  private readonly cache: Record<string, CachedBSONElement | false | undefined> =
+    Object.create(null);
+  /** Caches the index of elements that have been named */
+  private readonly indexFound: Record<number, boolean> = Object.create(null);
+
+  /** All bson elements in this document */
+  private readonly elements: BSONElement[];
+
+  constructor(
+    /** BSON bytes, this document begins at offset */
+    protected readonly bson: Uint8Array,
+    /** The start of the document */
+    private readonly offset = 0,
+    /** If this is an embedded document, indicates if this was a BSON array */
+    public readonly isArray = false
+  ) {
+    this.elements = parseToElementsToArray(this.bson, offset);
+  }
+
+  /** Only supports basic latin strings */
+  private isElementName(name: string, element: BSONElement): boolean {
+    const nameLength = element[BSONElementOffset.nameLength];
+    const nameOffset = element[BSONElementOffset.nameOffset];
+
+    if (name.length !== nameLength) return false;
+
+    for (let i = 0; i < name.length; i++) {
+      if (this.bson[nameOffset + i] !== name.charCodeAt(i)) return false;
+    }
+
+    return true;
+  }
+
+  /**
+   * Seeks into the elements array for an element matching the given name.
+   *
+   * @remarks
+   * Caching:
+   * - Caches the existence of a property making subsequent look ups for non-existent properties return immediately
+   * - Caches names mapped to elements to avoid reiterating the array and comparing the name again
+   * - Caches the index at which an element has been found to prevent rechecking against elements already determined to belong to another name
+   *
+   * @param name - a basic latin string name of a BSON element
+   * @returns
+   */
+  private getElement(name: string): CachedBSONElement | null {
+    const cachedElement = this.cache[name];
+    if (cachedElement === false) return null;
+
+    if (cachedElement != null) {
+      return cachedElement;
+    }
+
+    for (let index = 0; index < this.elements.length; index++) {
+      const element = this.elements[index];
+
+      // skip this element if it has already been associated with a name
+      if (!this.indexFound[index] && this.isElementName(name, element)) {
+        const cachedElement = { element, value: undefined };
+        this.cache[name] = cachedElement;
+        this.indexFound[index] = true;
+        return cachedElement;
+      }
+    }
+
+    this.cache[name] = false;
+    return null;
+  }
+
+  /**
+   * Translates BSON bytes into a javascript value. Checking `as` against the BSON element's type
+   * this methods returns the small subset of BSON types that the driver needs to function.
+   *
+   * @remarks
+   * - BSONType.null and BSONType.undefined always return null
+   * - If the type requested does not match this returns null
+   *
+   * @param element - The element to revive to a javascript value
+   * @param as - A type byte expected to be returned
+   */
+  private toJSValue<T extends keyof JSTypeOf>(element: BSONElement, as: T): JSTypeOf[T];
+  private toJSValue(element: BSONElement, as: keyof JSTypeOf): any {
+    const type = element[BSONElementOffset.type];
+    const offset = element[BSONElementOffset.offset];
+    const length = element[BSONElementOffset.length];
+
+    if (as !== type) {
+      return null;
+    }
+
+    switch (as) {
+      case BSONType.null:
+      case BSONType.undefined:
+        return null;
+      case BSONType.double:
+        return getFloat64LE(this.bson, offset);
+      case BSONType.int:
+        return getInt32LE(this.bson, offset);
+      case BSONType.long:
+        return getBigInt64LE(this.bson, offset);
+      case BSONType.bool:
+        return Boolean(this.bson[offset]);
+      case BSONType.objectId:
+        return new ObjectId(this.bson.subarray(offset, offset + 12));
+      case BSONType.timestamp:
+        return new Timestamp(getBigInt64LE(this.bson, offset));
+      case BSONType.string:
+        return toUTF8(this.bson, offset + 4, offset + length - 1, false);
+      case BSONType.binData: {
+        const totalBinarySize = getInt32LE(this.bson, offset);
+        const subType = this.bson[offset + 4];
+
+        if (subType === 2) {
+          const subType2BinarySize = getInt32LE(this.bson, offset + 1 + 4);
+          if (subType2BinarySize < 0)
+            throw new BSONError('Negative binary type element size found for subtype 0x02');
+          if (subType2BinarySize > totalBinarySize - 4)
+            throw new BSONError('Binary type with subtype 0x02 contains too long binary size');
+          if (subType2BinarySize < totalBinarySize - 4)
+            throw new BSONError('Binary type with subtype 0x02 contains too short binary size');
+          return new Binary(
+            this.bson.subarray(offset + 1 + 4 + 4, offset + 1 + 4 + 4 + subType2BinarySize),
+            2
+          );
+        }
+
+        return new Binary(
+          this.bson.subarray(offset + 1 + 4, offset + 1 + 4 + totalBinarySize),
+          subType
+        );
+      }
+      case BSONType.date:
+        // Pretend this is correct.
+        return new Date(Number(getBigInt64LE(this.bson, offset)));
+
+      case BSONType.object:
+        return new OnDemandDocument(this.bson, offset);
+      case BSONType.array:
+        return new OnDemandDocument(this.bson, offset, true);
+
+      default:
+        throw new BSONError(`Unsupported BSON type: ${as}`);
+    }
+  }
+
+  /**
+   * Checks for the existence of an element by name.
+   *
+   * @remarks
+   * Uses `getElement` with the expectation that will populate caches such that a `has` call
+   * followed by a `getElement` call will not repeat the cost paid by the first look up.
+   *
+   * @param name - element name
+   */
+  public has(name: string): boolean {
+    const cachedElement = this.cache[name];
+    if (cachedElement === false) return false;
+    if (cachedElement != null) return true;
+    return this.getElement(name) != null;
+  }
+
+  /**
+   * Turns BSON element with `name` into a javascript value.
+   *
+   * @typeParam T - must be one of the supported BSON types determined by `JSTypeOf` this will determine the return type of this function.
+   * @param name - the element name
+   * @param as - the bson type expected
+   * @param required - whether or not the element is expected to exist, if true this function will throw if it is not present
+   */
+  public get<const T extends keyof JSTypeOf>(
+    name: string,
+    as: T,
+    required?: false | undefined
+  ): JSTypeOf[T] | null;
+
+  /** `required` will make `get` throw if name does not exist or is null/undefined */
+  public get<const T extends keyof JSTypeOf>(name: string, as: T, required: true): JSTypeOf[T];
+
+  public get<const T extends keyof JSTypeOf>(
+    name: string,
+    as: T,
+    required?: boolean
+  ): JSTypeOf[T] | null {
+    const element = this.getElement(name);
+    if (element == null) {
+      if (required === true) {
+        throw new BSONError(`BSON element "${name}" is missing`);
+      } else {
+        return null;
+      }
+    }
+
+    if (element.value == null) {
+      const value = this.toJSValue(element.element, as);
+      if (value == null) {
+        if (required === true) {
+          throw new BSONError(`BSON element "${name}" is missing`);
+        } else {
+          return null;
+        }
+      }
+      // It is important to never store null
+      element.value = value;
+    }
+
+    return element.value;
+  }
+
+  /**
+   * Supports returning int, double, long, and bool as javascript numbers
+   *
+   * @remarks
+   * **NOTE:**
+   * - Use this _only_ when you believe the potential precision loss of an int64 is acceptable
+   * - This method does not cache the result as Longs or booleans would be stored incorrectly
+   *
+   * @param name - element name
+   * @param required - throws if name does not exist
+   */
+  public getNumber<const Req extends boolean = false>(
+    name: string,
+    required?: Req
+  ): Req extends true ? number : number | null;
+  public getNumber(name: string, required: boolean): number | null {
+    const maybeBool = this.get(name, BSONType.bool);
+    const bool = maybeBool == null ? null : maybeBool ? 1 : 0;
+
+    const maybeLong = this.get(name, BSONType.long);
+    const long = maybeLong == null ? null : Number(maybeLong);
+
+    const result = bool ?? long ?? this.get(name, BSONType.int) ?? this.get(name, BSONType.double);
+
+    if (required === true && result == null) {
+      throw new BSONError(`BSON element "${name}" is missing`);
+    }
+
+    return result;
+  }
+
+  /**
+   * Deserialize this object, DOES NOT cache result so avoid multiple invocations
+   * @param options - BSON deserialization options
+   */
+  public toObject(options?: BSONSerializeOptions): Record<string, any> {
+    return BSON.deserialize(this.bson, {
+      ...options,
+      index: this.offset,
+      allowObjectSmallerThanBufferSize: true
+    });
+  }
+
+  /**
+   * Iterates through the elements of a document reviving them using the `as` BSONType.
+   *
+   * @param as - The type to revive all elements as
+   */
+  public *valuesAs<const T extends keyof JSTypeOf>(as: T): Generator<JSTypeOf[T]> {
+    if (!this.isArray) {
+      throw new BSONError('Unexpected conversion of non-array value to array');
+    }
+    let counter = 0;
+    for (const element of this.elements) {
+      const value = this.toJSValue<T>(element, as);
+      this.cache[counter] = { element, value };
+      yield value;
+      counter += 1;
+    }
+  }
+}

test/mongodb.ts

@@ -130,6 +130,7 @@ export * from '../src/cmap/metrics';
 export * from '../src/cmap/stream_description';
 export * from '../src/cmap/wire_protocol/compression';
 export * from '../src/cmap/wire_protocol/constants';
+export * from '../src/cmap/wire_protocol/on_demand/document';
 export * from '../src/cmap/wire_protocol/shared';
 export * from '../src/collection';
 export * from '../src/connection_string';