diff --git a/MssqlMcp/Node/README.md b/MssqlMcp/Node/README.md index 3f7e687..3779e83 100644 --- a/MssqlMcp/Node/README.md +++ b/MssqlMcp/Node/README.md @@ -24,6 +24,7 @@ This server leverages the Model Context Protocol (MCP), a versatile framework th - Run MSSQL Database queries by just asking questions in plain English - Create, read, update, and delete data +- Get table statistics including row counts and space usage - Manage database schema (tables, indexes) - Secure connection handling - Real-time data interaction diff --git a/MssqlMcp/Node/src/index.ts b/MssqlMcp/Node/src/index.ts index 8dc1f30..cd2bf9b 100644 --- a/MssqlMcp/Node/src/index.ts +++ b/MssqlMcp/Node/src/index.ts @@ -20,6 +20,7 @@ import { ListTableTool } from "./tools/ListTableTool.js"; import { DropTableTool } from "./tools/DropTableTool.js"; import { DefaultAzureCredential, InteractiveBrowserCredential } from "@azure/identity"; import { DescribeTableTool } from "./tools/DescribeTableTool.js"; +import { GetTableStatsTool } from "./tools/GetTableStatsTool.js"; // MSSQL Database connection configuration // const credential = new DefaultAzureCredential(); @@ -69,6 +70,7 @@ const createIndexTool = new CreateIndexTool(); const listTableTool = new ListTableTool(); const dropTableTool = new DropTableTool(); const describeTableTool = new DescribeTableTool(); +const getTableStatsTool = new GetTableStatsTool(); const server = new Server( { @@ -89,8 +91,8 @@ const isReadOnly = process.env.READONLY === "true"; server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: isReadOnly - ? [listTableTool, readDataTool, describeTableTool] // todo: add searchDataTool to the list of tools available in readonly mode once implemented - : [insertDataTool, readDataTool, describeTableTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool], // add all new tools here + ? [listTableTool, readDataTool, describeTableTool, getTableStatsTool] // todo: add searchDataTool to the list of tools available in readonly mode once implemented + : [insertDataTool, readDataTool, describeTableTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool, getTableStatsTool], // add all new tools here })); server.setRequestHandler(CallToolRequestSchema, async (request) => { @@ -128,6 +130,9 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => { } result = await describeTableTool.run(args as { tableName: string }); break; + case getTableStatsTool.name: + result = await getTableStatsTool.run(args || {}); + break; default: return { content: [{ type: "text", text: `Unknown tool: ${name}` }], @@ -197,4 +202,4 @@ function wrapToolRun(tool: { run: (...args: any[]) => Promise }) { }; } -[insertDataTool, readDataTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool, describeTableTool].forEach(wrapToolRun); \ No newline at end of file +[insertDataTool, readDataTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool, describeTableTool, getTableStatsTool].forEach(wrapToolRun); \ No newline at end of file diff --git a/MssqlMcp/Node/src/tools/GetTableStatsTool.ts b/MssqlMcp/Node/src/tools/GetTableStatsTool.ts new file mode 100644 index 0000000..e9e71d1 --- /dev/null +++ b/MssqlMcp/Node/src/tools/GetTableStatsTool.ts @@ -0,0 +1,76 @@ +import sql from "mssql"; +import { Tool } from "@modelcontextprotocol/sdk/types.js"; + +export class GetTableStatsTool implements Tool { + [key: string]: any; + name = "get_table_stats"; + description = + "Returns row counts and space usage statistics for tables in the database."; + inputSchema = { + type: "object", + properties: { + tableName: { + type: "string", + description: + "Table name to get stats for (supports 'schema.table' format). If omitted, returns stats for all tables.", + }, + }, + required: [], + } as any; + + async run(params: any) { + try { + const { tableName } = params || {}; + + let tableNamePart: string | null = null; + let schemaPart: string | null = null; + + if (tableName && typeof tableName === "string") { + if (tableName.includes(".")) { + const parts = tableName.split("."); + schemaPart = parts[0]; + tableNamePart = parts[1]; + } else { + tableNamePart = tableName; + } + } + + const request = new sql.Request(); + + const query = ` + SELECT + s.name AS [schema], + t.name AS [table], + SUM(p.rows) AS [rowCount], + SUM(a.total_pages) * 8 AS totalSpaceKB, + SUM(a.used_pages) * 8 AS usedSpaceKB + FROM sys.tables t + INNER JOIN sys.schemas s ON t.schema_id = s.schema_id + INNER JOIN sys.indexes i ON t.object_id = i.object_id + INNER JOIN sys.partitions p ON i.object_id = p.object_id AND i.index_id = p.index_id + INNER JOIN sys.allocation_units a ON p.partition_id = a.container_id + WHERE i.index_id <= 1 + AND (@TableName IS NULL OR t.name = @TableName) + AND (@TableSchema IS NULL OR s.name = @TableSchema) + GROUP BY s.name, t.name + ORDER BY s.name, t.name`; + + request.input("TableName", sql.NVarChar, tableNamePart); + request.input("TableSchema", sql.NVarChar, schemaPart); + + const result = await request.query(query); + + return { + success: true, + message: `Retrieved stats for ${result.recordset.length} table(s).`, + data: result.recordset, + }; + } catch (error) { + console.error("Error getting table stats:", error); + return { + success: false, + message: `Failed to get table stats: ${error}`, + }; + } + } +} diff --git a/MssqlMcp/dotnet/MssqlMcp.Tests/UnitTests.cs b/MssqlMcp/dotnet/MssqlMcp.Tests/UnitTests.cs index 8058a86..04fdd7d 100644 --- a/MssqlMcp/dotnet/MssqlMcp.Tests/UnitTests.cs +++ b/MssqlMcp/dotnet/MssqlMcp.Tests/UnitTests.cs @@ -210,5 +210,43 @@ public async Task SqlInjection_NotExecuted_When_QueryFails() Assert.NotNull(describeResult); Assert.True(describeResult.Success); } + + [Fact] + public async Task GetTableStats_ReturnsStats_ForAllTables() + { + var result = await _tools.GetTableStats() as DbOperationResult; + Assert.NotNull(result); + Assert.True(result.Success); + Assert.NotNull(result.Data); + var stats = result.Data as List; + Assert.NotNull(stats); + } + + [Fact] + public async Task GetTableStats_ReturnsStats_ForSpecificTable() + { + var createResult = await _tools.CreateTable($"CREATE TABLE {_tableName} (Id INT PRIMARY KEY)") as DbOperationResult; + Assert.NotNull(createResult); + Assert.True(createResult.Success); + + var result = await _tools.GetTableStats(_tableName) as DbOperationResult; + Assert.NotNull(result); + Assert.True(result.Success); + Assert.NotNull(result.Data); + var stats = result.Data as List; + Assert.NotNull(stats); + Assert.Single(stats); + } + + [Fact] + public async Task GetTableStats_ReturnsEmpty_ForNonExistentTable() + { + var result = await _tools.GetTableStats("NonExistentTable_xyz") as DbOperationResult; + Assert.NotNull(result); + Assert.True(result.Success); + var stats = result.Data as List; + Assert.NotNull(stats); + Assert.Empty(stats); + } } } \ No newline at end of file diff --git a/MssqlMcp/dotnet/MssqlMcp/Tools/GetTableStats.cs b/MssqlMcp/dotnet/MssqlMcp/Tools/GetTableStats.cs new file mode 100644 index 0000000..15771e4 --- /dev/null +++ b/MssqlMcp/dotnet/MssqlMcp/Tools/GetTableStats.cs @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +using System.ComponentModel; +using Microsoft.Data.SqlClient; +using Microsoft.Extensions.Logging; +using ModelContextProtocol.Server; + +namespace Mssql.McpServer; + +public partial class Tools +{ + private const string TableStatsQuery = @" + SELECT + s.name AS [schema], + t.name AS [table], + SUM(p.rows) AS [rowCount], + SUM(a.total_pages) * 8 AS totalSpaceKB, + SUM(a.used_pages) * 8 AS usedSpaceKB + FROM sys.tables t + INNER JOIN sys.schemas s ON t.schema_id = s.schema_id + INNER JOIN sys.indexes i ON t.object_id = i.object_id + INNER JOIN sys.partitions p ON i.object_id = p.object_id AND i.index_id = p.index_id + INNER JOIN sys.allocation_units a ON p.partition_id = a.container_id + WHERE i.index_id <= 1 + AND (@TableName IS NULL OR t.name = @TableName) + AND (@TableSchema IS NULL OR s.name = @TableSchema) + GROUP BY s.name, t.name + ORDER BY s.name, t.name"; + + [McpServerTool( + Title = "Get Table Stats", + ReadOnly = true, + Idempotent = true, + Destructive = false), + Description("Returns row counts and space usage statistics for tables in the database.")] + public async Task GetTableStats( + [Description("Table name to get stats for (supports 'schema.table' format). If omitted, returns stats for all tables.")] string? name = null) + { + string? schema = null; + if (name != null && name.Contains('.')) + { + var parts = name.Split('.'); + if (parts.Length > 1) + { + schema = parts[0]; + name = parts[1]; + } + } + + var conn = await _connectionFactory.GetOpenConnectionAsync(); + try + { + using (conn) + { + using var cmd = new SqlCommand(TableStatsQuery, conn); + var _ = cmd.Parameters.AddWithValue("@TableName", name == null ? DBNull.Value : name); + _ = cmd.Parameters.AddWithValue("@TableSchema", schema == null ? DBNull.Value : schema); + + using var reader = await cmd.ExecuteReaderAsync(); + var stats = new List(); + while (await reader.ReadAsync()) + { + stats.Add(new + { + schema = reader.GetString(0), + table = reader.GetString(1), + rowCount = reader.GetInt64(2), + totalSpaceKB = reader.GetInt64(3), + usedSpaceKB = reader.GetInt64(4) + }); + } + + return new DbOperationResult(success: true, data: stats); + } + } + catch (Exception ex) + { + _logger.LogError(ex, "GetTableStats failed: {Message}", ex.Message); + return new DbOperationResult(success: false, error: ex.Message); + } + } +} diff --git a/MssqlMcp/dotnet/README.md b/MssqlMcp/dotnet/README.md index 0c69f07..be7558a 100644 --- a/MssqlMcp/dotnet/README.md +++ b/MssqlMcp/dotnet/README.md @@ -10,6 +10,7 @@ This project is a .NET 8 console application implementing a Model Context Protoc - **MCP Tools Implemented**: - ListTables: List all tables in the database. - DescribeTable: Get schema/details for a table. + - GetTableStats: Get row counts and space usage statistics for tables. - CreateTable: Create new tables. - DropTable: Drop existing tables. - InsertData: Insert data into tables.