smartmongo/ts/ts_tsmdb/server/handlers/IndexHandler.ts

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

// Cache of index engines per collection
const indexEngines: Map<string, IndexEngine> = new Map();

/**
 * Get or create an IndexEngine for a collection
 */
function getIndexEngine(storage: any, database: string, collection: string): IndexEngine {
  const key = `${database}.${collection}`;
  let engine = indexEngines.get(key);

  if (!engine) {
    engine = new IndexEngine(database, collection, storage);
    indexEngines.set(key, engine);
  }

  return engine;
}

/**
 * IndexHandler - Handles createIndexes, dropIndexes, listIndexes commands
 */
export class IndexHandler implements ICommandHandler {
  async handle(context: IHandlerContext): Promise<plugins.bson.Document> {
    const { command } = context;

    if (command.createIndexes) {
      return this.handleCreateIndexes(context);
    } else if (command.dropIndexes) {
      return this.handleDropIndexes(context);
    } else if (command.listIndexes) {
      return this.handleListIndexes(context);
    }

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

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

    const collection = command.createIndexes;
    const indexes = command.indexes || [];

    if (!Array.isArray(indexes)) {
      return {
        ok: 0,
        errmsg: 'indexes must be an array',
        code: 2,
        codeName: 'BadValue',
      };
    }

    // Ensure collection exists
    await storage.createCollection(database, collection);

    const indexEngine = getIndexEngine(storage, database, collection);
    const createdNames: string[] = [];
    let numIndexesBefore = 0;
    let numIndexesAfter = 0;

    try {
      const existingIndexes = await indexEngine.listIndexes();
      numIndexesBefore = existingIndexes.length;

      for (const indexSpec of indexes) {
        const key = indexSpec.key;
        const options = {
          name: indexSpec.name,
          unique: indexSpec.unique,
          sparse: indexSpec.sparse,
          expireAfterSeconds: indexSpec.expireAfterSeconds,
          background: indexSpec.background,
          partialFilterExpression: indexSpec.partialFilterExpression,
        };

        const name = await indexEngine.createIndex(key, options);
        createdNames.push(name);
      }

      const finalIndexes = await indexEngine.listIndexes();
      numIndexesAfter = finalIndexes.length;
    } catch (error: any) {
      return {
        ok: 0,
        errmsg: error.message || 'Failed to create index',
        code: error.code || 1,
        codeName: error.codeName || 'InternalError',
      };
    }

    return {
      ok: 1,
      numIndexesBefore,
      numIndexesAfter,
      createdCollectionAutomatically: false,
      commitQuorum: 'votingMembers',
    };
  }

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

    const collection = command.dropIndexes;
    const indexName = command.index;

    // Check if collection exists
    const exists = await storage.collectionExists(database, collection);
    if (!exists) {
      return {
        ok: 0,
        errmsg: `ns not found ${database}.${collection}`,
        code: 26,
        codeName: 'NamespaceNotFound',
      };
    }

    const indexEngine = getIndexEngine(storage, database, collection);

    try {
      if (indexName === '*') {
        // Drop all indexes except _id
        await indexEngine.dropAllIndexes();
      } else if (typeof indexName === 'string') {
        // Drop specific index by name
        await indexEngine.dropIndex(indexName);
      } else if (typeof indexName === 'object') {
        // Drop index by key specification
        const indexes = await indexEngine.listIndexes();
        const keyStr = JSON.stringify(indexName);

        for (const idx of indexes) {
          if (JSON.stringify(idx.key) === keyStr) {
            await indexEngine.dropIndex(idx.name);
            break;
          }
        }
      }

      return { ok: 1, nIndexesWas: 1 };
    } catch (error: any) {
      return {
        ok: 0,
        errmsg: error.message || 'Failed to drop index',
        code: error.code || 27,
        codeName: error.codeName || 'IndexNotFound',
      };
    }
  }

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

    const collection = command.listIndexes;
    const cursor = command.cursor || {};
    const batchSize = cursor.batchSize || 101;

    // Check if collection exists
    const exists = await storage.collectionExists(database, collection);
    if (!exists) {
      return {
        ok: 0,
        errmsg: `ns not found ${database}.${collection}`,
        code: 26,
        codeName: 'NamespaceNotFound',
      };
    }

    const indexEngine = getIndexEngine(storage, database, collection);
    const indexes = await indexEngine.listIndexes();

    // Format indexes for response
    const indexDocs = indexes.map(idx => ({
      v: idx.v || 2,
      key: idx.key,
      name: idx.name,
      ...(idx.unique ? { unique: idx.unique } : {}),
      ...(idx.sparse ? { sparse: idx.sparse } : {}),
      ...(idx.expireAfterSeconds !== undefined ? { expireAfterSeconds: idx.expireAfterSeconds } : {}),
    }));

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