smartmongo/ts/ts_tsmdb/server/handlers/UpdateHandler.ts

import * as plugins from '../../plugins.js';
import type { ICommandHandler, IHandlerContext } from '../CommandRouter.js';
import type { IStoredDocument } from '../../types/interfaces.js';
import { QueryEngine } from '../../engine/QueryEngine.js';
import { UpdateEngine } from '../../engine/UpdateEngine.js';

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

    // Check findAndModify first since it also has an 'update' field
    if (command.findAndModify) {
      return this.handleFindAndModify(context);
    } else if (command.update && typeof command.update === 'string') {
      // 'update' command has collection name as the value
      return this.handleUpdate(context);
    }

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

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

    const collection = command.update;
    if (typeof collection !== 'string') {
      return {
        ok: 0,
        errmsg: 'update command requires a collection name',
        code: 2,
        codeName: 'BadValue',
      };
    }

    // Get updates from command or document sequences
    let updates: plugins.bson.Document[] = command.updates || [];

    // Check for OP_MSG document sequences
    if (documentSequences && documentSequences.has('updates')) {
      updates = documentSequences.get('updates')!;
    }

    if (!Array.isArray(updates) || updates.length === 0) {
      return {
        ok: 0,
        errmsg: 'update command requires updates array',
        code: 2,
        codeName: 'BadValue',
      };
    }

    const ordered = command.ordered !== false;
    const writeErrors: plugins.bson.Document[] = [];
    let totalMatched = 0;
    let totalModified = 0;
    let totalUpserted = 0;
    const upserted: plugins.bson.Document[] = [];

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

    const indexEngine = context.getIndexEngine(collection);

    for (let i = 0; i < updates.length; i++) {
      const updateSpec = updates[i];
      const filter = updateSpec.q || updateSpec.filter || {};
      const update = updateSpec.u || updateSpec.update || {};
      const multi = updateSpec.multi || false;
      const upsert = updateSpec.upsert || false;
      const arrayFilters = updateSpec.arrayFilters;

      try {
        // Try to use index-accelerated query
        const candidateIds = await indexEngine.findCandidateIds(filter);

        let documents: IStoredDocument[];
        if (candidateIds !== null) {
          documents = await storage.findByIds(database, collection, candidateIds);
        } else {
          documents = await storage.findAll(database, collection);
        }

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

        if (matchingDocs.length === 0 && upsert) {
          // Upsert: create new document
          const newDoc: plugins.bson.Document = { _id: new plugins.bson.ObjectId() };

          // Apply filter fields to the new document
          this.applyFilterToDoc(newDoc, filter);

          // Apply update
          const updatedDoc = UpdateEngine.applyUpdate(newDoc as any, update, arrayFilters);

          // Handle $setOnInsert
          if (update.$setOnInsert) {
            Object.assign(updatedDoc, update.$setOnInsert);
          }

          // Update index for the new document
          await indexEngine.onInsert(updatedDoc);
          await storage.insertOne(database, collection, updatedDoc);
          totalUpserted++;
          upserted.push({ index: i, _id: updatedDoc._id });
        } else {
          // Update existing documents
          const docsToUpdate = multi ? matchingDocs : matchingDocs.slice(0, 1);
          totalMatched += docsToUpdate.length;

          for (const doc of docsToUpdate) {
            const updatedDoc = UpdateEngine.applyUpdate(doc, update, arrayFilters);

            // Check if document actually changed
            const changed = JSON.stringify(doc) !== JSON.stringify(updatedDoc);
            if (changed) {
              // Update index
              await indexEngine.onUpdate(doc as any, updatedDoc);
              await storage.updateById(database, collection, doc._id, updatedDoc);
              totalModified++;
            }
          }
        }
      } catch (error: any) {
        writeErrors.push({
          index: i,
          code: error.code || 1,
          errmsg: error.message || 'Update failed',
        });

        if (ordered) {
          break;
        }
      }
    }

    const response: plugins.bson.Document = {
      ok: 1,
      n: totalMatched + totalUpserted,
      nModified: totalModified,
    };

    if (upserted.length > 0) {
      response.upserted = upserted;
    }

    if (writeErrors.length > 0) {
      response.writeErrors = writeErrors;
    }

    return response;
  }

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

    const collection = command.findAndModify;
    const query = command.query || {};
    const update = command.update;
    const remove = command.remove || false;
    const returnNew = command.new || false;
    const upsert = command.upsert || false;
    const sort = command.sort;
    const fields = command.fields;
    const arrayFilters = command.arrayFilters;

    // Validate - either update or remove, not both
    if (update && remove) {
      return {
        ok: 0,
        errmsg: 'cannot specify both update and remove',
        code: 2,
        codeName: 'BadValue',
      };
    }

    if (!update && !remove) {
      return {
        ok: 0,
        errmsg: 'either update or remove is required',
        code: 2,
        codeName: 'BadValue',
      };
    }

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

    // Try to use index-accelerated query
    const indexEngine = context.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);
    }

    let matchingDocs = QueryEngine.filter(documents, query);

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

    const doc = matchingDocs[0];

    if (remove) {
      // Delete operation
      if (!doc) {
        return { ok: 1, value: null };
      }

      // Update index for delete
      await indexEngine.onDelete(doc as any);
      await storage.deleteById(database, collection, doc._id);

      let result = doc;
      if (fields) {
        result = QueryEngine.project([doc], fields)[0] as any;
      }

      return {
        ok: 1,
        value: result,
        lastErrorObject: {
          n: 1,
        },
      };
    } else {
      // Update operation
      if (!doc && !upsert) {
        return { ok: 1, value: null };
      }

      let resultDoc: plugins.bson.Document;
      let originalDoc: plugins.bson.Document | null = null;
      let isUpsert = false;

      if (doc) {
        // Update existing
        originalDoc = { ...doc };
        resultDoc = UpdateEngine.applyUpdate(doc, update, arrayFilters);
        // Update index
        await indexEngine.onUpdate(doc as any, resultDoc as any);
        await storage.updateById(database, collection, doc._id, resultDoc as any);
      } else {
        // Upsert
        isUpsert = true;
        const newDoc: plugins.bson.Document = { _id: new plugins.bson.ObjectId() };
        this.applyFilterToDoc(newDoc, query);
        resultDoc = UpdateEngine.applyUpdate(newDoc as any, update, arrayFilters);

        if (update.$setOnInsert) {
          Object.assign(resultDoc, update.$setOnInsert);
        }

        // Update index for insert
        await indexEngine.onInsert(resultDoc as any);
        await storage.insertOne(database, collection, resultDoc);
      }

      // Apply projection
      let returnValue = returnNew ? resultDoc : (originalDoc || null);
      if (returnValue && fields) {
        returnValue = QueryEngine.project([returnValue as any], fields)[0];
      }

      const response: plugins.bson.Document = {
        ok: 1,
        value: returnValue,
        lastErrorObject: {
          n: 1,
          updatedExisting: !isUpsert && doc !== undefined,
        },
      };

      if (isUpsert) {
        response.lastErrorObject.upserted = resultDoc._id;
      }

      return response;
    }
  }

  /**
   * Apply filter equality conditions to a new document (for upsert)
   */
  private applyFilterToDoc(doc: plugins.bson.Document, filter: plugins.bson.Document): void {
    for (const [key, value] of Object.entries(filter)) {
      // Skip operators
      if (key.startsWith('$')) continue;

      // Handle nested paths
      if (typeof value === 'object' && value !== null) {
        // Check if it's an operator
        const valueKeys = Object.keys(value);
        if (valueKeys.some(k => k.startsWith('$'))) {
          // Extract equality value from $eq if present
          if ('$eq' in value) {
            this.setNestedValue(doc, key, value.$eq);
          }
          continue;
        }
      }

      // Direct value assignment
      this.setNestedValue(doc, key, value);
    }
  }

  /**
   * Set a nested value using dot notation
   */
  private setNestedValue(obj: plugins.bson.Document, path: string, value: any): void {
    const parts = path.split('.');
    let current = obj;

    for (let i = 0; i < parts.length - 1; i++) {
      const part = parts[i];
      if (!(part in current)) {
        current[part] = {};
      }
      current = current[part];
    }

    current[parts[parts.length - 1]] = value;
  }
}