import * as plugins from '../../tsmdb.plugins.js'; import type { ICommandHandler, IHandlerContext, ICursorState } from '../CommandRouter.js'; import type { IStoredDocument } from '../../types/interfaces.js'; import { QueryEngine } from '../../engine/QueryEngine.js'; /** * FindHandler - Handles find, getMore, killCursors, count, distinct commands */ export class FindHandler implements ICommandHandler { private cursors: Map; private nextCursorId: () => bigint; constructor( cursors: Map, nextCursorId: () => bigint ) { this.cursors = cursors; this.nextCursorId = nextCursorId; } async handle(context: IHandlerContext): Promise { 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 { const { storage, database, command, getIndexEngine } = 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: [], }, }; } // Try to use index-accelerated query const indexEngine = getIndexEngine(collection); const candidateIds = await indexEngine.findCandidateIds(filter); let documents: IStoredDocument[]; if (candidateIds !== null) { // Index hit - fetch only candidate documents documents = await storage.findByIds(database, collection, candidateIds); // Still apply filter for any conditions the index couldn't fully satisfy documents = QueryEngine.filter(documents, filter); } else { // No suitable index - full collection scan 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 { 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 { 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 { const { storage, database, command, getIndexEngine } = 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 }; } // Try to use index-accelerated query const indexEngine = getIndexEngine(collection); const candidateIds = await indexEngine.findCandidateIds(query); let documents: IStoredDocument[]; if (candidateIds !== null) { // Index hit - fetch only candidate documents documents = await storage.findByIds(database, collection, candidateIds); documents = QueryEngine.filter(documents, query); } else { // No suitable index - full collection scan documents = await storage.findAll(database, collection); 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 { const { storage, database, command, getIndexEngine } = 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: [] }; } // Try to use index-accelerated query const indexEngine = getIndexEngine(collection); const candidateIds = await indexEngine.findCandidateIds(query); let documents: IStoredDocument[]; if (candidateIds !== null) { documents = await storage.findByIds(database, collection, candidateIds); } else { documents = await storage.findAll(database, collection); } // Get distinct values const values = QueryEngine.distinct(documents, key, query); return { ok: 1, values }; } }