Skip to content

Commit 69b6b43

Browse files
Alan GuoAlan Guo
authored andcommitted
feat(MssqlMcp): Add GetTableStats tool for row counts and space usage
1 parent b2eda9c commit 69b6b43

File tree

6 files changed

+207
-3
lines changed

6 files changed

+207
-3
lines changed

MssqlMcp/Node/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ This server leverages the Model Context Protocol (MCP), a versatile framework th
2424

2525
- Run MSSQL Database queries by just asking questions in plain English
2626
- Create, read, update, and delete data
27+
- Get table statistics including row counts and space usage
2728
- Manage database schema (tables, indexes)
2829
- Secure connection handling
2930
- Real-time data interaction

MssqlMcp/Node/src/index.ts

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import { ListTableTool } from "./tools/ListTableTool.js";
2020
import { DropTableTool } from "./tools/DropTableTool.js";
2121
import { DefaultAzureCredential, InteractiveBrowserCredential } from "@azure/identity";
2222
import { DescribeTableTool } from "./tools/DescribeTableTool.js";
23+
import { GetTableStatsTool } from "./tools/GetTableStatsTool.js";
2324

2425
// MSSQL Database connection configuration
2526
// const credential = new DefaultAzureCredential();
@@ -69,6 +70,7 @@ const createIndexTool = new CreateIndexTool();
6970
const listTableTool = new ListTableTool();
7071
const dropTableTool = new DropTableTool();
7172
const describeTableTool = new DescribeTableTool();
73+
const getTableStatsTool = new GetTableStatsTool();
7274

7375
const server = new Server(
7476
{
@@ -89,8 +91,8 @@ const isReadOnly = process.env.READONLY === "true";
8991

9092
server.setRequestHandler(ListToolsRequestSchema, async () => ({
9193
tools: isReadOnly
92-
? [listTableTool, readDataTool, describeTableTool] // todo: add searchDataTool to the list of tools available in readonly mode once implemented
93-
: [insertDataTool, readDataTool, describeTableTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool], // add all new tools here
94+
? [listTableTool, readDataTool, describeTableTool, getTableStatsTool] // todo: add searchDataTool to the list of tools available in readonly mode once implemented
95+
: [insertDataTool, readDataTool, describeTableTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool, getTableStatsTool], // add all new tools here
9496
}));
9597

9698
server.setRequestHandler(CallToolRequestSchema, async (request) => {
@@ -128,6 +130,9 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
128130
}
129131
result = await describeTableTool.run(args as { tableName: string });
130132
break;
133+
case getTableStatsTool.name:
134+
result = await getTableStatsTool.run(args || {});
135+
break;
131136
default:
132137
return {
133138
content: [{ type: "text", text: `Unknown tool: ${name}` }],
@@ -197,4 +202,4 @@ function wrapToolRun(tool: { run: (...args: any[]) => Promise<any> }) {
197202
};
198203
}
199204

200-
[insertDataTool, readDataTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool, describeTableTool].forEach(wrapToolRun);
205+
[insertDataTool, readDataTool, updateDataTool, createTableTool, createIndexTool, dropTableTool, listTableTool, describeTableTool, getTableStatsTool].forEach(wrapToolRun);
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import sql from "mssql";
2+
import { Tool } from "@modelcontextprotocol/sdk/types.js";
3+
4+
export class GetTableStatsTool implements Tool {
5+
[key: string]: any;
6+
name = "get_table_stats";
7+
description =
8+
"Returns row counts and space usage statistics for tables in the database.";
9+
inputSchema = {
10+
type: "object",
11+
properties: {
12+
tableName: {
13+
type: "string",
14+
description:
15+
"Table name to get stats for (supports 'schema.table' format). If omitted, returns stats for all tables.",
16+
},
17+
},
18+
required: [],
19+
} as any;
20+
21+
async run(params: any) {
22+
try {
23+
const { tableName } = params || {};
24+
25+
let tableNamePart: string | null = null;
26+
let schemaPart: string | null = null;
27+
28+
if (tableName && typeof tableName === "string") {
29+
if (tableName.includes(".")) {
30+
const parts = tableName.split(".");
31+
schemaPart = parts[0];
32+
tableNamePart = parts[1];
33+
} else {
34+
tableNamePart = tableName;
35+
}
36+
}
37+
38+
const request = new sql.Request();
39+
40+
const query = `
41+
SELECT
42+
s.name AS [schema],
43+
t.name AS [table],
44+
SUM(p.rows) AS [rowCount],
45+
SUM(a.total_pages) * 8 AS totalSpaceKB,
46+
SUM(a.used_pages) * 8 AS usedSpaceKB
47+
FROM sys.tables t
48+
INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
49+
INNER JOIN sys.indexes i ON t.object_id = i.object_id
50+
INNER JOIN sys.partitions p ON i.object_id = p.object_id AND i.index_id = p.index_id
51+
INNER JOIN sys.allocation_units a ON p.partition_id = a.container_id
52+
WHERE i.index_id <= 1
53+
AND (@TableName IS NULL OR t.name = @TableName)
54+
AND (@TableSchema IS NULL OR s.name = @TableSchema)
55+
GROUP BY s.name, t.name
56+
ORDER BY s.name, t.name`;
57+
58+
request.input("TableName", sql.NVarChar, tableNamePart);
59+
request.input("TableSchema", sql.NVarChar, schemaPart);
60+
61+
const result = await request.query(query);
62+
63+
return {
64+
success: true,
65+
message: `Retrieved stats for ${result.recordset.length} table(s).`,
66+
data: result.recordset,
67+
};
68+
} catch (error) {
69+
console.error("Error getting table stats:", error);
70+
return {
71+
success: false,
72+
message: `Failed to get table stats: ${error}`,
73+
};
74+
}
75+
}
76+
}

MssqlMcp/dotnet/MssqlMcp.Tests/UnitTests.cs

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,5 +210,43 @@ public async Task SqlInjection_NotExecuted_When_QueryFails()
210210
Assert.NotNull(describeResult);
211211
Assert.True(describeResult.Success);
212212
}
213+
214+
[Fact]
215+
public async Task GetTableStats_ReturnsStats_ForAllTables()
216+
{
217+
var result = await _tools.GetTableStats() as DbOperationResult;
218+
Assert.NotNull(result);
219+
Assert.True(result.Success);
220+
Assert.NotNull(result.Data);
221+
var stats = result.Data as List<object>;
222+
Assert.NotNull(stats);
223+
}
224+
225+
[Fact]
226+
public async Task GetTableStats_ReturnsStats_ForSpecificTable()
227+
{
228+
var createResult = await _tools.CreateTable($"CREATE TABLE {_tableName} (Id INT PRIMARY KEY)") as DbOperationResult;
229+
Assert.NotNull(createResult);
230+
Assert.True(createResult.Success);
231+
232+
var result = await _tools.GetTableStats(_tableName) as DbOperationResult;
233+
Assert.NotNull(result);
234+
Assert.True(result.Success);
235+
Assert.NotNull(result.Data);
236+
var stats = result.Data as List<object>;
237+
Assert.NotNull(stats);
238+
Assert.Single(stats);
239+
}
240+
241+
[Fact]
242+
public async Task GetTableStats_ReturnsEmpty_ForNonExistentTable()
243+
{
244+
var result = await _tools.GetTableStats("NonExistentTable_xyz") as DbOperationResult;
245+
Assert.NotNull(result);
246+
Assert.True(result.Success);
247+
var stats = result.Data as List<object>;
248+
Assert.NotNull(stats);
249+
Assert.Empty(stats);
250+
}
213251
}
214252
}
Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT license.
3+
4+
using System.ComponentModel;
5+
using Microsoft.Data.SqlClient;
6+
using Microsoft.Extensions.Logging;
7+
using ModelContextProtocol.Server;
8+
9+
namespace Mssql.McpServer;
10+
11+
public partial class Tools
12+
{
13+
private const string TableStatsQuery = @"
14+
SELECT
15+
s.name AS [schema],
16+
t.name AS [table],
17+
SUM(p.rows) AS [rowCount],
18+
SUM(a.total_pages) * 8 AS totalSpaceKB,
19+
SUM(a.used_pages) * 8 AS usedSpaceKB
20+
FROM sys.tables t
21+
INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
22+
INNER JOIN sys.indexes i ON t.object_id = i.object_id
23+
INNER JOIN sys.partitions p ON i.object_id = p.object_id AND i.index_id = p.index_id
24+
INNER JOIN sys.allocation_units a ON p.partition_id = a.container_id
25+
WHERE i.index_id <= 1
26+
AND (@TableName IS NULL OR t.name = @TableName)
27+
AND (@TableSchema IS NULL OR s.name = @TableSchema)
28+
GROUP BY s.name, t.name
29+
ORDER BY s.name, t.name";
30+
31+
[McpServerTool(
32+
Title = "Get Table Stats",
33+
ReadOnly = true,
34+
Idempotent = true,
35+
Destructive = false),
36+
Description("Returns row counts and space usage statistics for tables in the database.")]
37+
public async Task<DbOperationResult> GetTableStats(
38+
[Description("Table name to get stats for (supports 'schema.table' format). If omitted, returns stats for all tables.")] string? name = null)
39+
{
40+
string? schema = null;
41+
if (name != null && name.Contains('.'))
42+
{
43+
var parts = name.Split('.');
44+
if (parts.Length > 1)
45+
{
46+
schema = parts[0];
47+
name = parts[1];
48+
}
49+
}
50+
51+
var conn = await _connectionFactory.GetOpenConnectionAsync();
52+
try
53+
{
54+
using (conn)
55+
{
56+
using var cmd = new SqlCommand(TableStatsQuery, conn);
57+
var _ = cmd.Parameters.AddWithValue("@TableName", name == null ? DBNull.Value : name);
58+
_ = cmd.Parameters.AddWithValue("@TableSchema", schema == null ? DBNull.Value : schema);
59+
60+
using var reader = await cmd.ExecuteReaderAsync();
61+
var stats = new List<object>();
62+
while (await reader.ReadAsync())
63+
{
64+
stats.Add(new
65+
{
66+
schema = reader.GetString(0),
67+
table = reader.GetString(1),
68+
rowCount = reader.GetInt64(2),
69+
totalSpaceKB = reader.GetInt64(3),
70+
usedSpaceKB = reader.GetInt64(4)
71+
});
72+
}
73+
74+
return new DbOperationResult(success: true, data: stats);
75+
}
76+
}
77+
catch (Exception ex)
78+
{
79+
_logger.LogError(ex, "GetTableStats failed: {Message}", ex.Message);
80+
return new DbOperationResult(success: false, error: ex.Message);
81+
}
82+
}
83+
}

MssqlMcp/dotnet/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ This project is a .NET 8 console application implementing a Model Context Protoc
1010
- **MCP Tools Implemented**:
1111
- ListTables: List all tables in the database.
1212
- DescribeTable: Get schema/details for a table.
13+
- GetTableStats: Get row counts and space usage statistics for tables.
1314
- CreateTable: Create new tables.
1415
- DropTable: Drop existing tables.
1516
- InsertData: Insert data into tables.

0 commit comments

Comments
 (0)