@dataql/mongodb-adapter/src/index.ts

MongoDB adapter for DataQL with zero API changes

import { Data } from "@dataql/core";

// Import DataOptions from the correct location
type DataOptions = {
  dbName?: string;
  env?: "dev" | "prod";
  devPrefix?: string;
  appToken?: string;
  customConnection?: any; // CustomRequestConnection
};

// MongoDB-like types and interfaces
export interface MongoClientOptions {
  // MongoDB compatibility options (ignored but supported for zero-API-change migration)
  useNewUrlParser?: boolean;
  useUnifiedTopology?: boolean;
  maxPoolSize?: number;
  serverSelectionTimeoutMS?: number;
  socketTimeoutMS?: number;
  family?: number;
  bufferMaxEntries?: number;
  readPreference?: string;
  ssl?: boolean;
  sslValidate?: boolean;
  sslCA?: string;
  sslCert?: string;
  sslKey?: string;
  sslPass?: string;
  sslCRL?: string;
  authSource?: string;
  authMechanism?: string;

  // DataQL configuration - operations go through DataQL infrastructure (Client → Worker → Lambda → MongoDB)
  dataql?: DataOptions;
}

export interface InsertOneOptions {
  bypassDocumentValidation?: boolean;
  forceServerObjectId?: boolean;
  writeConcern?: WriteConcern;
  comment?: any;
}

export interface InsertManyOptions {
  bypassDocumentValidation?: boolean;
  forceServerObjectId?: boolean;
  ordered?: boolean;
  writeConcern?: WriteConcern;
  comment?: any;
}

export interface UpdateOptions {
  arrayFilters?: any[];
  bypassDocumentValidation?: boolean;
  collation?: any;
  hint?: any;
  upsert?: boolean;
  writeConcern?: WriteConcern;
  comment?: any;
}

export interface ReplaceOptions {
  bypassDocumentValidation?: boolean;
  collation?: any;
  hint?: any;
  upsert?: boolean;
  writeConcern?: WriteConcern;
  comment?: any;
}

export interface DeleteOptions {
  collation?: any;
  hint?: any;
  writeConcern?: WriteConcern;
  comment?: any;
}

export interface FindOptions {
  allowDiskUse?: boolean;
  allowPartialResults?: boolean;
  batchSize?: number;
  collation?: any;
  comment?: any;
  cursorType?: string;
  hint?: any;
  limit?: number;
  max?: any;
  maxAwaitTimeMS?: number;
  maxTimeMS?: number;
  min?: any;
  noCursorTimeout?: boolean;
  oplogReplay?: boolean;
  projection?: any;
  readConcern?: any;
  readPreference?: any;
  returnKey?: boolean;
  showRecordId?: boolean;
  skip?: number;
  sort?: any;
  tailable?: boolean;
  awaitData?: boolean;
}

export interface WriteConcern {
  w?: number | string;
  j?: boolean;
  wtimeout?: number;
}

export interface InsertOneResult {
  acknowledged: boolean;
  insertedId: any;
}

export interface InsertManyResult {
  acknowledged: boolean;
  insertedCount: number;
  insertedIds: { [key: number]: any };
}

export interface UpdateResult {
  acknowledged: boolean;
  matchedCount: number;
  modifiedCount: number;
  upsertedId?: any;
  upsertedCount: number;
}

export interface DeleteResult {
  acknowledged: boolean;
  deletedCount: number;
}

export interface BulkWriteOptions {
  ordered?: boolean;
  bypassDocumentValidation?: boolean;
  writeConcern?: WriteConcern;
  comment?: any;
}

export interface BulkWriteResult {
  acknowledged: boolean;
  insertedCount: number;
  matchedCount: number;
  modifiedCount: number;
  deletedCount: number;
  upsertedCount: number;
  insertedIds: { [key: number]: any };
  upsertedIds: { [key: number]: any };
}

export type BulkWriteOperation =
  | { insertOne: { document: any } }
  | { updateOne: { filter: any; update: any; upsert?: boolean } }
  | { updateMany: { filter: any; update: any; upsert?: boolean } }
  | { deleteOne: { filter: any } }
  | { deleteMany: { filter: any } }
  | { replaceOne: { filter: any; replacement: any; upsert?: boolean } };

// ObjectId class for MongoDB compatibility
export class ObjectId {
  private _id: string;

  constructor(id?: string | ObjectId) {
    if (id instanceof ObjectId) {
      this._id = id.toString();
    } else if (typeof id === "string") {
      this._id = id;
    } else {
      // Generate a random ObjectId-like string
      this._id = this._generateObjectId();
    }
  }

  private _generateObjectId(): string {
    const timestamp = Math.floor(Date.now() / 1000).toString(16);
    const randomBytes = Math.random().toString(16).substr(2, 16);
    return (timestamp + randomBytes).substr(0, 24);
  }

  toString(): string {
    return this._id;
  }

  toHexString(): string {
    return this._id;
  }

  equals(other: ObjectId | string): boolean {
    if (other instanceof ObjectId) {
      return this._id === other._id;
    }
    return this._id === other;
  }

  getTimestamp(): Date {
    const timestamp = parseInt(this._id.substr(0, 8), 16);
    return new Date(timestamp * 1000);
  }

  static isValid(id: any): boolean {
    if (typeof id === "string") {
      return /^[0-9a-fA-F]{24}$/.test(id);
    }
    return id instanceof ObjectId;
  }

  static createFromHexString(hexString: string): ObjectId {
    return new ObjectId(hexString);
  }
}

// Cursor class for query results
export class FindCursor<T = any> implements AsyncIterable<T> {
  private _data: Data;
  private _collectionName: string;
  private _filter: any;
  private _options: FindOptions;
  private _results?: T[];

  constructor(
    data: Data,
    collectionName: string,
    filter: any = {},
    options: FindOptions = {}
  ) {
    this._data = data;
    this._collectionName = collectionName;
    this._filter = filter;
    this._options = options;
  }

  private get _collection() {
    return this._data.collection(
      this._collectionName,
      this._getDefaultSchema()
    );
  }

  private _getDefaultSchema() {
    return {
      _id: { type: "ID", required: true },
      createdAt: { type: "Date", default: "now" },
      updatedAt: { type: "Date", default: "now" },
    };
  }

  async toArray(): Promise<T[]> {
    if (this._results) {
      return this._results;
    }

    let results = await this._collection.find(this._filter);

    // Apply sorting
    if (this._options.sort) {
      results = this._applySorting(results, this._options.sort);
    }

    // Apply skip
    if (this._options.skip) {
      results = results.slice(this._options.skip);
    }

    // Apply limit
    if (this._options.limit) {
      results = results.slice(0, this._options.limit);
    }

    // Apply projection
    if (this._options.projection) {
      results = this._applyProjection(results, this._options.projection);
    }

    this._results = results;
    return results;
  }

  private _applySorting(results: any[], sort: any): any[] {
    return results.sort((a, b) => {
      for (const [field, direction] of Object.entries(sort)) {
        const aVal = a[field];
        const bVal = b[field];
        const sortDir = direction === -1 ? -1 : 1;

        if (aVal < bVal) return -1 * sortDir;
        if (aVal > bVal) return 1 * sortDir;
      }
      return 0;
    });
  }

  private _applyProjection(results: any[], projection: any): any[] {
    const isInclusion = Object.values(projection).some((val) => val === 1);

    return results.map((doc) => {
      if (isInclusion) {
        const projected: any = {};
        for (const [field, include] of Object.entries(projection)) {
          if (include === 1) {
            projected[field] = doc[field];
          }
        }
        return projected;
      } else {
        const projected = { ...doc };
        for (const [field, exclude] of Object.entries(projection)) {
          if (exclude === 0) {
            delete projected[field];
          }
        }
        return projected;
      }
    });
  }

  async next(): Promise<T | null> {
    const results = await this.toArray();
    return results.length > 0 ? results.shift()! : null;
  }

  async hasNext(): Promise<boolean> {
    const results = await this.toArray();
    return results.length > 0;
  }

  async forEach(fn: (doc: T) => void): Promise<void> {
    const results = await this.toArray();
    results.forEach(fn);
  }

  map<U>(fn: (doc: T) => U): FindCursor<U> {
    // Return a new cursor with mapped results
    const newCursor = new FindCursor<U>(
      this._data,
      this._collectionName,
      this._filter,
      this._options
    );
    return newCursor;
  }

  filter(fn: (doc: T) => boolean): FindCursor<T> {
    // Return a new cursor with filtered results
    return new FindCursor<T>(
      this._data,
      this._collectionName,
      this._filter,
      this._options
    );
  }

  limit(count: number): FindCursor<T> {
    const newCursor = new FindCursor<T>(
      this._data,
      this._collectionName,
      this._filter,
      { ...this._options, limit: count }
    );
    return newCursor;
  }

  skip(count: number): FindCursor<T> {
    const newCursor = new FindCursor<T>(
      this._data,
      this._collectionName,
      this._filter,
      { ...this._options, skip: count }
    );
    return newCursor;
  }

  sort(sort: any): FindCursor<T> {
    const newCursor = new FindCursor<T>(
      this._data,
      this._collectionName,
      this._filter,
      { ...this._options, sort }
    );
    return newCursor;
  }

  project(projection: any): FindCursor<T> {
    const newCursor = new FindCursor<T>(
      this._data,
      this._collectionName,
      this._filter,
      { ...this._options, projection }
    );
    return newCursor;
  }

  async count(): Promise<number> {
    const results = await this._collection.find(this._filter);
    return results.length;
  }

  async *[Symbol.asyncIterator](): AsyncIterator<T> {
    const results = await this.toArray();
    for (const result of results) {
      yield result;
    }
  }
}

// Collection class
export class Collection<T = any> {
  constructor(
    private _data: Data,
    private _db: Db,
    public collectionName: string
  ) {}

  private get _collection() {
    return this._data.collection(this.collectionName, this._getDefaultSchema());
  }

  private _getDefaultSchema() {
    return {
      _id: { type: "ID", required: true },
      createdAt: { type: "Date", default: "now" },
      updatedAt: { type: "Date", default: "now" },
    };
  }

  // Insert operations
  async insertOne(
    doc: T,
    options?: InsertOneOptions
  ): Promise<InsertOneResult> {
    const docWithId = {
      _id: new ObjectId(),
      ...doc,
    };

    const result = await this._collection.create(docWithId);
    return {
      acknowledged: true,
      insertedId: docWithId._id,
    };
  }

  async insertMany(
    docs: T[],
    options?: InsertManyOptions
  ): Promise<InsertManyResult> {
    const docsWithIds = docs.map((doc, index) => ({
      _id: new ObjectId(),
      ...doc,
    }));

    const results = await this._collection.create(docsWithIds);
    const insertedIds: { [key: number]: any } = {};
    docsWithIds.forEach((doc, index) => {
      insertedIds[index] = doc._id;
    });

    return {
      acknowledged: true,
      insertedCount: docsWithIds.length,
      insertedIds,
    };
  }

  // Find operations
  find(filter: any = {}, options?: FindOptions): FindCursor<T> {
    return new FindCursor<T>(this._data, this.collectionName, filter, options);
  }

  async findOne(filter: any = {}, options?: FindOptions): Promise<T | null> {
    const cursor = this.find(filter, { ...options, limit: 1 });
    const results = await cursor.toArray();
    return results.length > 0 ? results[0] : null;
  }

  async findOneAndUpdate(
    filter: any,
    update: any,
    options?: UpdateOptions & { returnDocument?: "before" | "after" }
  ): Promise<{ value: T | null }> {
    const existing = await this.findOne(filter);

    if (!existing && !options?.upsert) {
      return { value: null };
    }

    const updateResult = await this._collection.update(
      filter,
      update,
      options?.upsert
    );

    if (options?.returnDocument === "before") {
      return { value: existing };
    } else {
      const updated = await this.findOne(filter);
      return { value: updated };
    }
  }

  async findOneAndDelete(filter: any): Promise<{ value: T | null }> {
    const existing = await this.findOne(filter);
    if (existing) {
      await this._collection.delete(filter);
    }
    return { value: existing };
  }

  async findOneAndReplace(
    filter: any,
    replacement: any,
    options?: ReplaceOptions & { returnDocument?: "before" | "after" }
  ): Promise<{ value: T | null }> {
    const existing = await this.findOne(filter);

    if (!existing && !options?.upsert) {
      return { value: null };
    }

    await this._collection.update(filter, replacement, options?.upsert);

    if (options?.returnDocument === "before") {
      return { value: existing };
    } else {
      const updated = await this.findOne(filter);
      return { value: updated };
    }
  }

  // Update operations
  async updateOne(
    filter: any,
    update: any,
    options?: UpdateOptions
  ): Promise<UpdateResult> {
    const result = await this._collection.update(
      filter,
      update,
      options?.upsert
    );

    return {
      acknowledged: true,
      matchedCount: result ? 1 : 0,
      modifiedCount: result ? 1 : 0,
      upsertedId: options?.upsert && !result ? new ObjectId() : undefined,
      upsertedCount: options?.upsert && !result ? 1 : 0,
    };
  }

  async updateMany(
    filter: any,
    update: any,
    options?: UpdateOptions
  ): Promise<UpdateResult> {
    // DataQL doesn't have updateMany, so we'll find and update individually
    const existing = await this._collection.find(filter);
    let modifiedCount = 0;

    for (const doc of existing) {
      const result = await this._collection.update({ _id: doc._id }, update);
      if (result) modifiedCount++;
    }

    return {
      acknowledged: true,
      matchedCount: existing.length,
      modifiedCount,
      upsertedCount: 0,
    };
  }

  async replaceOne(
    filter: any,
    replacement: any,
    options?: ReplaceOptions
  ): Promise<UpdateResult> {
    const result = await this._collection.update(
      filter,
      replacement,
      options?.upsert
    );

    return {
      acknowledged: true,
      matchedCount: result ? 1 : 0,
      modifiedCount: result ? 1 : 0,
      upsertedId: options?.upsert && !result ? new ObjectId() : undefined,
      upsertedCount: options?.upsert && !result ? 1 : 0,
    };
  }

  // Delete operations
  async deleteOne(filter: any, options?: DeleteOptions): Promise<DeleteResult> {
    const existing = await this.findOne(filter);
    if (existing) {
      await this._collection.delete(filter);
      return {
        acknowledged: true,
        deletedCount: 1,
      };
    }
    return {
      acknowledged: true,
      deletedCount: 0,
    };
  }

  async deleteMany(
    filter: any,
    options?: DeleteOptions
  ): Promise<DeleteResult> {
    const existing = await this._collection.find(filter);
    let deletedCount = 0;

    for (const doc of existing) {
      await this._collection.delete({ _id: doc._id });
      deletedCount++;
    }

    return {
      acknowledged: true,
      deletedCount,
    };
  }

  // Bulk operations
  async bulkWrite(
    operations: BulkWriteOperation[],
    options?: BulkWriteOptions
  ): Promise<BulkWriteResult> {
    let insertedCount = 0;
    let matchedCount = 0;
    let modifiedCount = 0;
    let deletedCount = 0;
    let upsertedCount = 0;
    const insertedIds: { [key: number]: any } = {};
    const upsertedIds: { [key: number]: any } = {};

    for (let i = 0; i < operations.length; i++) {
      const operation = operations[i];

      try {
        if ("insertOne" in operation) {
          const result = await this.insertOne(operation.insertOne.document);
          insertedCount++;
          insertedIds[i] = result.insertedId;
        } else if ("updateOne" in operation) {
          const result = await this.updateOne(
            operation.updateOne.filter,
            operation.updateOne.update,
            { upsert: operation.updateOne.upsert }
          );
          matchedCount += result.matchedCount;
          modifiedCount += result.modifiedCount;
          if (result.upsertedId) {
            upsertedCount++;
            upsertedIds[i] = result.upsertedId;
          }
        } else if ("updateMany" in operation) {
          const result = await this.updateMany(
            operation.updateMany.filter,
            operation.updateMany.update,
            { upsert: operation.updateMany.upsert }
          );
          matchedCount += result.matchedCount;
          modifiedCount += result.modifiedCount;
          upsertedCount += result.upsertedCount;
        } else if ("deleteOne" in operation) {
          const result = await this.deleteOne(operation.deleteOne.filter);
          deletedCount += result.deletedCount;
        } else if ("deleteMany" in operation) {
          const result = await this.deleteMany(operation.deleteMany.filter);
          deletedCount += result.deletedCount;
        } else if ("replaceOne" in operation) {
          const result = await this.replaceOne(
            operation.replaceOne.filter,
            operation.replaceOne.replacement,
            { upsert: operation.replaceOne.upsert }
          );
          matchedCount += result.matchedCount;
          modifiedCount += result.modifiedCount;
          if (result.upsertedId) {
            upsertedCount++;
            upsertedIds[i] = result.upsertedId;
          }
        }
      } catch (error) {
        if (!options?.ordered) {
          continue; // Continue with next operation if not ordered
        }
        throw error; // Stop on first error if ordered
      }
    }

    return {
      acknowledged: true,
      insertedCount,
      matchedCount,
      modifiedCount,
      deletedCount,
      upsertedCount,
      insertedIds,
      upsertedIds,
    };
  }

  // Count operations
  async countDocuments(filter: any = {}): Promise<number> {
    const results = await this._collection.find(filter);
    return results.length;
  }

  async estimatedDocumentCount(): Promise<number> {
    return this.countDocuments();
  }

  // Aggregation
  aggregate(pipeline: any[]): any {
    // Basic aggregation support - would need more sophisticated implementation
    return {
      toArray: async () => {
        // For now, just return find results
        const results = await this._collection.find({});
        return results;
      },
    };
  }

  // Index operations (no-op for DataQL)
  async createIndex(fieldOrSpec: any, options?: any): Promise<string> {
    return "index_created";
  }

  async createIndexes(indexSpecs: any[]): Promise<string[]> {
    return indexSpecs.map((_, i) => `index_${i}_created`);
  }

  async dropIndex(indexName: string): Promise<any> {
    return { ok: 1 };
  }

  async dropIndexes(): Promise<any> {
    return { ok: 1 };
  }

  async listIndexes(): Promise<any[]> {
    return [];
  }

  // Utility methods
  async distinct(field: string, filter: any = {}): Promise<any[]> {
    const results = await this._collection.find(filter);
    const values = results.map((doc) => doc[field]);
    return [...new Set(values)];
  }

  // Drop collection
  async drop(): Promise<boolean> {
    // DataQL doesn't have explicit drop - return true for compatibility
    return true;
  }
}

// Database class
export class Db {
  constructor(
    private _data: Data,
    public databaseName: string
  ) {}

  collection<T = any>(name: string): Collection<T> {
    return new Collection<T>(this._data, this, name);
  }

  async listCollections(): Promise<any[]> {
    // DataQL doesn't track collections separately
    return [];
  }

  async dropCollection(name: string): Promise<boolean> {
    // DataQL doesn't have explicit drop - return true for compatibility
    return true;
  }

  async dropDatabase(): Promise<any> {
    // DataQL doesn't have explicit drop - return ok for compatibility
    return { ok: 1 };
  }

  async stats(): Promise<any> {
    return {
      db: this.databaseName,
      collections: 0,
      views: 0,
      objects: 0,
      avgObjSize: 0,
      dataSize: 0,
      storageSize: 0,
      totalSize: 0,
      indexes: 0,
      indexSize: 0,
      scaleFactor: 1,
    };
  }
}

// MongoClient class
export class MongoClient {
  private _data: Data;
  private _connected = false;
  private _databaseName?: string;

  constructor(
    private _url: string,
    private _options?: MongoClientOptions
  ) {
    // Extract database name from MongoDB URL for DataQL database isolation
    this._databaseName = this._extractDatabaseName(_url);

    // Configure DataQL to work through its infrastructure (Client → Worker → Lambda → MongoDB)
    const dataqlOptions = {
      // Pass database name for DataQL's per-client database isolation
      dbName: this._databaseName || "mongodb_app",
      // MongoDB URL will be handled by DataQL's infrastructure routing
      appToken: _options?.dataql?.appToken || "mongodb_adapter",
      env: _options?.dataql?.env || "prod",
      devPrefix: _options?.dataql?.devPrefix || "mongodb_",
      customConnection: _options?.dataql?.customConnection,
    };

    this._data = new Data(dataqlOptions);
  }

  private _extractDatabaseName(url: string): string | undefined {
    try {
      // Extract database name from MongoDB connection string
      // mongodb://host:port/database or mongodb+srv://host/database
      const match = url.match(/\/([^/?]+)(?:\?|$)/);
      return match ? match[1] : undefined;
    } catch {
      return undefined;
    }
  }

  async connect(): Promise<this> {
    // Connection is handled by DataQL infrastructure
    this._connected = true;
    return this;
  }

  async close(): Promise<void> {
    // Cleanup handled by DataQL
    this._connected = false;
  }

  db(name?: string): Db {
    if (!this._connected) {
      throw new Error(
        "MongoClient must be connected before running operations"
      );
    }
    // Database operations go through DataQL infrastructure
    const dbName = name || this._databaseName || "default";
    return new Db(this._data, dbName);
  }

  isConnected(): boolean {
    return this._connected;
  }

  // Static connect method
  static async connect(
    url: string,
    options?: MongoClientOptions
  ): Promise<MongoClient> {
    const client = new MongoClient(url, options);
    await client.connect();
    return client;
  }
}

// Default export
export default MongoClient;

// Additional MongoDB utilities
export const BSON = {
  ObjectId,
  ObjectID: ObjectId, // Legacy alias
};

// Connection helper
export async function connect(
  url: string,
  options?: MongoClientOptions
): Promise<MongoClient> {
  return MongoClient.connect(url, options);
}