smartmongo/ts/tsmdb/server/handlers/FindHandler.ts

import * as plugins from '../../tsmdb.plugins.js';
import type { ICommandHandler, IHandlerContext, ICursorState } from '../CommandRouter.js';
import { QueryEngine } from '../../engine/QueryEngine.js';

/**
 * FindHandler - Handles find, getMore, killCursors, count, distinct commands
 */
export class FindHandler implements ICommandHandler {
  private cursors: Map<bigint, ICursorState>;
  private nextCursorId: () => bigint;

  constructor(
    cursors: Map<bigint, ICursorState>,
    nextCursorId: () => bigint
  ) {
    this.cursors = cursors;
    this.nextCursorId = nextCursorId;
  }

  async handle(context: IHandlerContext): Promise<plugins.bson.Document> {
    const { command } = context;

    // Determine which operation to perform
    if (command.find) {
      return this.handleFind(context);
    } else if (command.getMore !== undefined) {
      return this.handleGetMore(context);
    } else if (command.killCursors) {
      return this.handleKillCursors(context);
    } else if (command.count) {
      return this.handleCount(context);
    } else if (command.distinct) {
      return this.handleDistinct(context);
    }

    return {
      ok: 0,
      errmsg: 'Unknown find-related command',
      code: 59,
      codeName: 'CommandNotFound',
    };
  }

  /**
   * Handle find command
   */
  private async handleFind(context: IHandlerContext): Promise<plugins.bson.Document> {
    const { storage, database, command } = context;

    const collection = command.find;
    const filter = command.filter || {};
    const projection = command.projection;
    const sort = command.sort;
    const skip = command.skip || 0;
    const limit = command.limit || 0;
    const batchSize = command.batchSize || 101;
    const singleBatch = command.singleBatch || false;

    // Ensure collection exists
    const exists = await storage.collectionExists(database, collection);
    if (!exists) {
      // Return empty cursor for non-existent collection
      return {
        ok: 1,
        cursor: {
          id: plugins.bson.Long.fromNumber(0),
          ns: `${database}.${collection}`,
          firstBatch: [],
        },
      };
    }

    // Get all documents
    let documents = await storage.findAll(database, collection);

    // Apply filter
    documents = QueryEngine.filter(documents, filter);

    // Apply sort
    if (sort) {
      documents = QueryEngine.sort(documents, sort);
    }

    // Apply skip
    if (skip > 0) {
      documents = documents.slice(skip);
    }

    // Apply limit
    if (limit > 0) {
      documents = documents.slice(0, limit);
    }

    // Apply projection
    if (projection) {
      documents = QueryEngine.project(documents, projection) as any[];
    }

    // Determine how many documents to return in first batch
    const effectiveBatchSize = Math.min(batchSize, documents.length);
    const firstBatch = documents.slice(0, effectiveBatchSize);
    const remaining = documents.slice(effectiveBatchSize);

    // Create cursor if there are more documents
    let cursorId = BigInt(0);
    if (remaining.length > 0 && !singleBatch) {
      cursorId = this.nextCursorId();
      this.cursors.set(cursorId, {
        id: cursorId,
        database,
        collection,
        documents: remaining,
        position: 0,
        batchSize,
        createdAt: new Date(),
      });
    }

    return {
      ok: 1,
      cursor: {
        id: plugins.bson.Long.fromBigInt(cursorId),
        ns: `${database}.${collection}`,
        firstBatch,
      },
    };
  }

  /**
   * Handle getMore command
   */
  private async handleGetMore(context: IHandlerContext): Promise<plugins.bson.Document> {
    const { database, command } = context;

    const cursorIdInput = command.getMore;
    const collection = command.collection;
    const batchSize = command.batchSize || 101;

    // Convert cursorId to bigint
    let cursorId: bigint;
    if (typeof cursorIdInput === 'bigint') {
      cursorId = cursorIdInput;
    } else if (cursorIdInput instanceof plugins.bson.Long) {
      cursorId = cursorIdInput.toBigInt();
    } else {
      cursorId = BigInt(cursorIdInput);
    }

    const cursor = this.cursors.get(cursorId);
    if (!cursor) {
      return {
        ok: 0,
        errmsg: `cursor id ${cursorId} not found`,
        code: 43,
        codeName: 'CursorNotFound',
      };
    }

    // Verify namespace
    if (cursor.database !== database || cursor.collection !== collection) {
      return {
        ok: 0,
        errmsg: 'cursor namespace mismatch',
        code: 43,
        codeName: 'CursorNotFound',
      };
    }

    // Get next batch
    const start = cursor.position;
    const end = Math.min(start + batchSize, cursor.documents.length);
    const nextBatch = cursor.documents.slice(start, end);
    cursor.position = end;

    // Check if cursor is exhausted
    let returnCursorId = cursorId;
    if (cursor.position >= cursor.documents.length) {
      this.cursors.delete(cursorId);
      returnCursorId = BigInt(0);
    }

    return {
      ok: 1,
      cursor: {
        id: plugins.bson.Long.fromBigInt(returnCursorId),
        ns: `${database}.${collection}`,
        nextBatch,
      },
    };
  }

  /**
   * Handle killCursors command
   */
  private async handleKillCursors(context: IHandlerContext): Promise<plugins.bson.Document> {
    const { command } = context;

    const collection = command.killCursors;
    const cursorIds = command.cursors || [];

    const cursorsKilled: plugins.bson.Long[] = [];
    const cursorsNotFound: plugins.bson.Long[] = [];
    const cursorsUnknown: plugins.bson.Long[] = [];

    for (const idInput of cursorIds) {
      let cursorId: bigint;
      if (typeof idInput === 'bigint') {
        cursorId = idInput;
      } else if (idInput instanceof plugins.bson.Long) {
        cursorId = idInput.toBigInt();
      } else {
        cursorId = BigInt(idInput);
      }

      if (this.cursors.has(cursorId)) {
        this.cursors.delete(cursorId);
        cursorsKilled.push(plugins.bson.Long.fromBigInt(cursorId));
      } else {
        cursorsNotFound.push(plugins.bson.Long.fromBigInt(cursorId));
      }
    }

    return {
      ok: 1,
      cursorsKilled,
      cursorsNotFound,
      cursorsUnknown,
      cursorsAlive: [],
    };
  }

  /**
   * Handle count command
   */
  private async handleCount(context: IHandlerContext): Promise<plugins.bson.Document> {
    const { storage, database, command } = context;

    const collection = command.count;
    const query = command.query || {};
    const skip = command.skip || 0;
    const limit = command.limit || 0;

    // Check if collection exists
    const exists = await storage.collectionExists(database, collection);
    if (!exists) {
      return { ok: 1, n: 0 };
    }

    // Get all documents
    let documents = await storage.findAll(database, collection);

    // Apply filter
    documents = QueryEngine.filter(documents, query);

    // Apply skip
    if (skip > 0) {
      documents = documents.slice(skip);
    }

    // Apply limit
    if (limit > 0) {
      documents = documents.slice(0, limit);
    }

    return { ok: 1, n: documents.length };
  }

  /**
   * Handle distinct command
   */
  private async handleDistinct(context: IHandlerContext): Promise<plugins.bson.Document> {
    const { storage, database, command } = context;

    const collection = command.distinct;
    const key = command.key;
    const query = command.query || {};

    if (!key) {
      return {
        ok: 0,
        errmsg: 'distinct requires a key',
        code: 2,
        codeName: 'BadValue',
      };
    }

    // Check if collection exists
    const exists = await storage.collectionExists(database, collection);
    if (!exists) {
      return { ok: 1, values: [] };
    }

    // Get all documents
    const documents = await storage.findAll(database, collection);

    // Get distinct values
    const values = QueryEngine.distinct(documents, key, query);

    return { ok: 1, values };
  }
}