MCP CosmosDB

Instructions

Implementation Reference

• src/tools/containerAnalysis.ts:106-170 (handler)

  The primary handler function for the mcp_container_stats tool. It queries the CosmosDB container for document count, samples documents to estimate size and partition key distribution, and returns ContainerStats.

  export const mcp_container_stats = async (args: { container_id: string; sample_size?: number }): Promise<ToolResult<ContainerStats>> => {
    const { container_id, sample_size = 1000 } = args;
    console.log('Executing mcp_container_stats with:', args);
  
    try {
      const container = getContainer(container_id);
      
      // Query to count total documents
      const countQuery = 'SELECT VALUE COUNT(1) FROM c';
      const { resources: countResult } = await container.items.query(countQuery).fetchAll();
      const documentCount = countResult[0] || 0;
  
      // Get partition key path for statistics
      const { resource: containerDef } = await container.read();
      if (!containerDef || !containerDef.partitionKey || !containerDef.partitionKey.paths || containerDef.partitionKey.paths.length === 0) {
        throw new Error(`Container ${container_id} does not have a valid partition key defined`);
      }
      const partitionKeyPath = containerDef.partitionKey.paths[0];
  
      // Sample documents to estimate size and analyze partitions
      const sampleQuery = `SELECT TOP ${sample_size} * FROM c`;
      const { resources: sampleDocs } = await container.items.query(sampleQuery).fetchAll();
  
      // Calculate estimated size based on sample
      let totalSampleSize = 0;
      const partitionStats: Record<string, { count: number; size: number }> = {};
  
      sampleDocs.forEach(doc => {
        const docSize = JSON.stringify(doc).length;
        totalSampleSize += docSize;
  
        // Get partition key value
        const partitionValue = getNestedProperty(doc, partitionKeyPath.substring(1)); // Remove leading '/'
        const partitionKey = String(partitionValue || 'undefined');
  
        if (!partitionStats[partitionKey]) {
          partitionStats[partitionKey] = { count: 0, size: 0 };
        }
        partitionStats[partitionKey].count++;
        partitionStats[partitionKey].size += docSize;
      });
  
      // Estimate total size
      const avgDocSize = sampleDocs.length > 0 ? totalSampleSize / sampleDocs.length : 0;
      const estimatedSizeInKB = Math.round((documentCount * avgDocSize) / 1024);
  
      // Convert partition stats
      const partitionKeyStatistics = Object.entries(partitionStats).map(([key, stats]) => ({
        partitionKeyValue: key,
        documentCount: Math.round((stats.count / sampleDocs.length) * documentCount),
        sizeInKB: Math.round(stats.size / 1024)
      }));
  
      const containerStats: ContainerStats = {
        documentCount,
        sizeInKB: estimatedSizeInKB,
        partitionKeyStatistics
      };
  
      return { success: true, data: containerStats };
    } catch (error: any) {
      console.error(`Error in mcp_container_stats for container ${container_id}: ${error.message}`);
      return { success: false, error: error.message };
    }
  };

• src/tools.ts:51-69 (schema)

  The input schema definition for the mcp_container_stats tool as part of the MCP_COSMOSDB_TOOLS array, used for tool listing and validation.

  {
    name: "mcp_container_stats",
    description: "Get statistical information about a container including document count and partition key distribution",
    inputSchema: {
      type: "object",
      properties: {
        container_id: {
          type: "string",
          description: "The ID of the container to analyze"
        },
        sample_size: {
          type: "number",
          description: "Sample size for statistics calculation",
          default: 1000
        }
      },
      required: ["container_id"]
    }
  },

• src/server.ts:100-102 (registration)

  Dispatch/registration in the CallTool handler switch statement that routes execution to the mcp_container_stats handler.

  case 'mcp_container_stats':
      result = await toolHandlers.mcp_container_stats(input as any);
      break;

• src/tools/containerAnalysis.ts:173-177 (helper)

  Helper function used within mcp_container_stats to extract partition key values from nested document properties.

  function getNestedProperty(obj: any, path: string): any {
    return path.split('/').reduce((current, key) => {
      return current && current[key] !== undefined ? current[key] : undefined;
    }, obj);
  } 

• src/tools/index.ts:8-8 (registration)

  Re-export of the mcp_container_stats handler from containerAnalysis.ts, making it available via tools/index.ts which is imported by mcp-server.ts.

  mcp_container_stats