From ec160d3a4a940c57a644c83580fa5517de7b3544 Mon Sep 17 00:00:00 2001 From: Victor Colin Amador Date: Sun, 21 Jun 2026 23:24:06 -0600 Subject: [PATCH 1/4] Fix unbounded enumeration in azmcp postgres list (#472) Cap both 'azmcp postgres list' database and table results at 10,000 entries to prevent OOM/perf issues on servers with very large catalogs, mirroring the existing MySQL behavior: - Add MaxRowCount = 10_000 constant in PostgresService (parity with MySQL's MaxRowCount). - ListDatabasesAsync: ORDER BY datname LIMIT @maxResults; append the same '... (output limited to 10,000 databases ...)' sentinel row that MySQL appends when the cap is reached. - ListTablesAsync: ORDER BY table_name LIMIT @maxResults (cap+1) and return new TableListResult(Tables, IsTruncated) so the command can surface truncation via an explicit flag. Truncation is detected via Count > cap rather than the MySQL read-then-probe pattern, which silently swallows the truncation signal when combined with a SQL LIMIT cap+1. - IPostgresService updated to return TableListResult from ListTablesAsync. - PostgresListCommandResult gains an optional 'tablesTruncated' bool. - Add unit tests for the new cap/truncation behavior and update existing PostgresListCommand tests for the new return type. - Document the cap in azmcp-commands.md and add a Bugs Fixed changelog entry. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../changelog-entries/1782105637872.yaml | 3 + .../Azure.Mcp.Server/docs/azmcp-commands.md | 3 + .../src/Commands/PostgresListCommand.cs | 6 +- .../src/Services/IPostgresService.cs | 2 +- .../src/Services/PostgresService.cs | 35 ++++- .../src/Services/TableListResult.cs | 6 + .../PostgresListCommandTests.cs | 37 +++++- .../Services/PostgresServiceRowLimitTests.cs | 120 ++++++++++++++++++ 8 files changed, 201 insertions(+), 11 deletions(-) create mode 100644 servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml create mode 100644 tools/Azure.Mcp.Tools.Postgres/src/Services/TableListResult.cs create mode 100644 tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs diff --git a/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml b/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml new file mode 100644 index 0000000000..a03b96af5f --- /dev/null +++ b/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml @@ -0,0 +1,3 @@ +changes: + - section: "Bugs Fixed" + description: "Cap 'azmcp postgres list' database and table results at 10,000 entries to prevent unbounded enumeration on large servers, mirroring the existing 'azmcp mysql list' behavior. When the database list is truncated, a sentinel row is appended; when the table list is truncated, the response includes a 'tablesTruncated: true' flag." diff --git a/servers/Azure.Mcp.Server/docs/azmcp-commands.md b/servers/Azure.Mcp.Server/docs/azmcp-commands.md index 236f542d8d..9af84a7c6e 100644 --- a/servers/Azure.Mcp.Server/docs/azmcp-commands.md +++ b/servers/Azure.Mcp.Server/docs/azmcp-commands.md @@ -2298,6 +2298,9 @@ azmcp mysql server param set --subscription \ # Without parameters: lists all PostgreSQL servers in the resource group # With --server: lists all databases on that server # With --server and --database: lists all tables in that database +# Database and table results are capped at 10,000 entries. When the database list is +# truncated a "... (output limited to 10,000 databases ...)" sentinel row is appended; +# when the table list is truncated, the response also includes "tablesTruncated": true. # ❌ Destructive | ✅ Idempotent | ❌ OpenWorld | ✅ ReadOnly | ❌ Secret | ❌ LocalRequired azmcp postgres list --subscription \ --resource-group \ diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs index ec6306aa47..62f743d67f 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs @@ -80,7 +80,7 @@ public override async Task ExecuteAsync(CommandContext context, if (!string.IsNullOrEmpty(options.Database)) { // List tables in specified database - List tables = await _postgresService.ListTablesAsync( + TableListResult tableResult = await _postgresService.ListTablesAsync( options.Subscription!, options.ResourceGroup!, options.AuthType!, @@ -91,7 +91,7 @@ public override async Task ExecuteAsync(CommandContext context, cancellationToken); context.Response.Results = ResponseResult.Create( - new(null, null, tables ?? []), + new(null, null, tableResult.Tables ?? [], tableResult.IsTruncated ? true : null), PostgresJsonContext.Default.PostgresListCommandResult); } else if (!string.IsNullOrEmpty(options.Server)) @@ -132,5 +132,5 @@ public override async Task ExecuteAsync(CommandContext context, return context.Response; } - public record PostgresListCommandResult(List? Servers, List? Databases, List? Tables); + public record PostgresListCommandResult(List? Servers, List? Databases, List? Tables, bool? TablesTruncated = null); } diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs index 0f7f7e68e8..c009d5dbf2 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs @@ -27,7 +27,7 @@ Task> ExecuteQueryAsync( string query, CancellationToken cancellationToken); - Task> ListTablesAsync( + Task ListTablesAsync( string subscriptionId, string resourceGroup, string authType, diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs index 6203be93dd..f7a54f01f4 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs @@ -36,6 +36,8 @@ public class PostgresService( private readonly IEntraTokenProvider _entraTokenAuth = entraTokenAuth; private readonly IDbProvider _dbProvider = dbProvider; + private const int MaxRowCount = 10_000; + private async Task GetEntraIdAccessTokenAsync(CancellationToken cancellationToken) { var tokenCredential = await GetCredential(cancellationToken); @@ -91,15 +93,27 @@ public async Task> ListDatabasesAsync( var host = NormalizeServerName(server); var connectionString = BuildConnectionString(host, "postgres", user, passwordToUse); - var query = "SELECT datname FROM pg_database WHERE datistemplate = false;"; + var query = "SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname LIMIT @maxResults;"; await using IPostgresResource resource = await _dbProvider.GetPostgresResource(connectionString, authType, cancellationToken); await using NpgsqlCommand command = _dbProvider.GetCommand(query, resource); + // Cap at exactly MaxRowCount: truncation is signaled by appending a sentinel row to the returned list + // (mirrors MySqlService.ListDatabasesAsync). Tables use cap+1 instead because they signal truncation + // via a structured IsTruncated flag and need to observe an N+1th row to set it. + command.Parameters.AddWithValue("maxResults", MaxRowCount); await using DbDataReader reader = await _dbProvider.ExecuteReaderAsync(command, cancellationToken); var dbs = new List(); - while (await reader.ReadAsync(cancellationToken)) + var dbCount = 0; + while (await reader.ReadAsync(cancellationToken) && dbCount < MaxRowCount) { dbs.Add(reader.GetString(0)); + dbCount++; + } + + if (dbCount >= MaxRowCount) + { + dbs.Add($"... (output limited to {MaxRowCount:N0} databases for security and performance reasons)"); } + return dbs; } @@ -161,7 +175,7 @@ public async Task> ExecuteQueryAsync( return rows; } - public async Task> ListTablesAsync( + public async Task ListTablesAsync( string subscriptionId, string resourceGroup, string authType, @@ -175,16 +189,27 @@ public async Task> ListTablesAsync( var host = NormalizeServerName(server); var connectionString = BuildConnectionString(host, database, user, passwordToUse); - var query = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';"; + var query = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' ORDER BY table_name LIMIT @maxResults;"; await using IPostgresResource resource = await _dbProvider.GetPostgresResource(connectionString, authType, cancellationToken); await using NpgsqlCommand command = _dbProvider.GetCommand(query, resource); + // Fetch cap+1 rows so we can detect truncation by observing whether an extra row exists, then trim it. + // Unlike ListDatabasesAsync (which signals truncation via a sentinel row appended to the list), + // tables return a structured TableListResult with an IsTruncated flag, so we need the extra row to set it. + command.Parameters.AddWithValue("maxResults", MaxRowCount + 1); await using DbDataReader reader = await _dbProvider.ExecuteReaderAsync(command, cancellationToken); var tables = new List(); while (await reader.ReadAsync(cancellationToken)) { tables.Add(reader.GetString(0)); } - return tables; + + var isTruncated = tables.Count > MaxRowCount; + if (isTruncated) + { + tables.RemoveRange(MaxRowCount, tables.Count - MaxRowCount); + } + + return new TableListResult(tables, isTruncated); } public async Task> GetTableSchemaAsync( diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/TableListResult.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/TableListResult.cs new file mode 100644 index 0000000000..729ad30107 --- /dev/null +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/TableListResult.cs @@ -0,0 +1,6 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +namespace Azure.Mcp.Tools.Postgres.Services; + +public sealed record TableListResult(List Tables, bool IsTruncated); diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs index 18efb26a29..11ccbffe30 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs @@ -107,7 +107,7 @@ public async Task ExecuteAsync_ListsTables_WhenServerAndDatabaseProvided() "server1", "db1", Arg.Any()) - .Returns(expectedTables); + .Returns(new TableListResult(expectedTables, false)); var response = await ExecuteCommandAsync( "--subscription", "sub123", @@ -122,6 +122,38 @@ public async Task ExecuteAsync_ListsTables_WhenServerAndDatabaseProvided() Assert.Null(result.Servers); Assert.Null(result.Databases); Assert.Equal(expectedTables, result.Tables); + Assert.Null(result.TablesTruncated); + } + + [Fact] + public async Task ExecuteAsync_SetsTablesTruncated_WhenServiceReportsTruncation() + { + var expectedTables = new List { "users", "products", "orders" }; + Service.ListTablesAsync( + "sub123", + "rg1", + AuthTypes.MicrosoftEntra, + "user1", + null, + "server1", + "db1", + Arg.Any()) + .Returns(new TableListResult(expectedTables, true)); + + var response = await ExecuteCommandAsync( + "--subscription", "sub123", + "--resource-group", "rg1", + "--user", "user1", + $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, + "--server", "server1", + "--database", "db1"); + + var result = ValidateAndDeserializeResponse(response, PostgresJsonContext.Default.PostgresListCommandResult); + + Assert.Null(result.Servers); + Assert.Null(result.Databases); + Assert.Equal(expectedTables, result.Tables); + Assert.True(result.TablesTruncated); } [Fact] @@ -181,7 +213,7 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoTablesExist() "server1", "db1", Arg.Any()) - .Returns([]); + .Returns(new TableListResult([], false)); var response = await ExecuteCommandAsync( "--subscription", "sub123", @@ -197,6 +229,7 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoTablesExist() Assert.Null(result.Databases); Assert.NotNull(result.Tables); Assert.Empty(result.Tables); + Assert.Null(result.TablesTruncated); } [Fact] diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs new file mode 100644 index 0000000000..c085b4244f --- /dev/null +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Data.Common; +using Azure.Mcp.Core.Services.Azure.ResourceGroup; +using Azure.Mcp.Core.Services.Azure.Subscription; +using Azure.Mcp.Core.Services.Azure.Tenant; +using Azure.Mcp.Tools.Postgres.Auth; +using Azure.Mcp.Tools.Postgres.Providers; +using Azure.Mcp.Tools.Postgres.Services; +using Azure.Mcp.Tools.Postgres.Tests.Services.Support; +using Npgsql; +using NSubstitute; +using Xunit; + +namespace Azure.Mcp.Tools.Postgres.Tests.Services; + +public class PostgresServiceRowLimitTests +{ + private const int MaxRowCount = 10_000; + + private readonly IResourceGroupService _resourceGroupService = Substitute.For(); + private readonly ISubscriptionService _subscriptionService = Substitute.For(); + private readonly ITenantService _tenantService = Substitute.For(); + private readonly IEntraTokenProvider _entraTokenAuth = Substitute.For(); + private readonly IDbProvider _dbProvider = Substitute.For(); + private readonly PostgresService _postgresService; + + private const string SubscriptionId = "test-sub"; + private const string ResourceGroup = "test-rg"; + private const string User = "test-user"; + private const string Server = "test-server"; + private const string Database = "test-db"; + private const string AuthType = "MicrosoftEntra"; + + public PostgresServiceRowLimitTests() + { + _entraTokenAuth.GetEntraToken(Arg.Any(), Arg.Any()) + .Returns(new Azure.Core.AccessToken("fake-token", DateTime.UtcNow.AddHours(1))); + + _dbProvider.GetPostgresResource(Arg.Any(), Arg.Any(), Arg.Any()) + .Returns(Substitute.For()); + _dbProvider.GetCommand(Arg.Any(), Arg.Any()) + .Returns(Substitute.For()); + + _postgresService = new PostgresService(_resourceGroupService, _subscriptionService, _tenantService, _entraTokenAuth, _dbProvider); + } + + private void StubReader(int rowCount, string columnName) + { + var rows = Enumerable.Range(0, rowCount).Select(i => new[] { $"item{i:D5}" }).ToArray(); + _dbProvider.ExecuteReaderAsync(Arg.Any(), Arg.Any()) + .Returns(new FakeDbDataReader(rows, [columnName])); + } + + [Fact] + public async Task ListDatabasesAsync_UnderCap_DoesNotAppendSentinel() + { + StubReader(rowCount: 3, columnName: "datname"); + + var result = await _postgresService.ListDatabasesAsync( + SubscriptionId, ResourceGroup, AuthType, User, null, Server, TestContext.Current.CancellationToken); + + Assert.Equal(3, result.Count); + Assert.DoesNotContain(result, r => r.StartsWith("... (output limited", StringComparison.Ordinal)); + } + + [Fact] + public async Task ListDatabasesAsync_AtCap_AppendsSentinelRow() + { + // FakeDbDataReader ignores SQL LIMIT, so simulate the worst case by returning exactly MaxRowCount rows. + StubReader(rowCount: MaxRowCount, columnName: "datname"); + + var result = await _postgresService.ListDatabasesAsync( + SubscriptionId, ResourceGroup, AuthType, User, null, Server, TestContext.Current.CancellationToken); + + // MaxRowCount real rows + 1 sentinel row (mirrors MySQL behavior). + Assert.Equal(MaxRowCount + 1, result.Count); + Assert.StartsWith("... (output limited", result[^1]); + Assert.Contains("10,000", result[^1]); + } + + [Fact] + public async Task ListTablesAsync_UnderCap_ReturnsAllAndNotTruncated() + { + StubReader(rowCount: 5, columnName: "table_name"); + + var result = await _postgresService.ListTablesAsync( + SubscriptionId, ResourceGroup, AuthType, User, null, Server, Database, TestContext.Current.CancellationToken); + + Assert.Equal(5, result.Tables.Count); + Assert.False(result.IsTruncated); + } + + [Fact] + public async Task ListTablesAsync_AtCap_ReturnsCapRowsAndNotTruncatedWhenNoExtra() + { + // Reader returns exactly MaxRowCount rows — boundary case, nothing beyond the cap. + StubReader(rowCount: MaxRowCount, columnName: "table_name"); + + var result = await _postgresService.ListTablesAsync( + SubscriptionId, ResourceGroup, AuthType, User, null, Server, Database, TestContext.Current.CancellationToken); + + Assert.Equal(MaxRowCount, result.Tables.Count); + Assert.False(result.IsTruncated); + } + + [Fact] + public async Task ListTablesAsync_OverCap_ReturnsCapRowsAndIsTruncated() + { + // Reader returns MaxRowCount + 1 (the cap+1 LIMIT in production); detect truncation via the extra read. + StubReader(rowCount: MaxRowCount + 1, columnName: "table_name"); + + var result = await _postgresService.ListTablesAsync( + SubscriptionId, ResourceGroup, AuthType, User, null, Server, Database, TestContext.Current.CancellationToken); + + Assert.Equal(MaxRowCount, result.Tables.Count); + Assert.True(result.IsTruncated); + } +} From 095d7786949285db9358b0b0527ac39ead6c067e Mon Sep 17 00:00:00 2001 From: Victor Colin Amador Date: Mon, 22 Jun 2026 16:30:32 -0600 Subject: [PATCH 2/4] Use truncation flag for database listing per review feedback (#472) Address alzimmermsft's PR review: - ListDatabasesAsync now returns DatabaseListResult with an IsTruncated flag instead of appending a sentinel row, and uses LIMIT cap+1 so truncation is actually detectable. - Surface DatabasesTruncated on PostgresListCommandResult. - Make PostgresService.MaxRowCount internal so tests reference it instead of duplicating the constant. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../changelog-entries/1782105637872.yaml | 2 +- .../Azure.Mcp.Server/docs/azmcp-commands.md | 5 ++- .../src/Commands/PostgresListCommand.cs | 8 ++--- .../src/Services/DatabaseListResult.cs | 6 ++++ .../src/Services/IPostgresService.cs | 2 +- .../src/Services/PostgresService.cs | 23 +++++-------- .../PostgresListCommandTests.cs | 34 +++++++++++++++++-- .../Services/PostgresServiceRowLimitTests.cs | 31 +++++++++++------ 8 files changed, 76 insertions(+), 35 deletions(-) create mode 100644 tools/Azure.Mcp.Tools.Postgres/src/Services/DatabaseListResult.cs diff --git a/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml b/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml index a03b96af5f..bfa4730b32 100644 --- a/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml +++ b/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml @@ -1,3 +1,3 @@ changes: - section: "Bugs Fixed" - description: "Cap 'azmcp postgres list' database and table results at 10,000 entries to prevent unbounded enumeration on large servers, mirroring the existing 'azmcp mysql list' behavior. When the database list is truncated, a sentinel row is appended; when the table list is truncated, the response includes a 'tablesTruncated: true' flag." + description: "Cap 'azmcp postgres list' database and table results at 10,000 entries to prevent unbounded enumeration on large servers. When the results are truncated, the response includes a 'databasesTruncated: true' or 'tablesTruncated: true' flag." diff --git a/servers/Azure.Mcp.Server/docs/azmcp-commands.md b/servers/Azure.Mcp.Server/docs/azmcp-commands.md index 9af84a7c6e..7f3d11cacf 100644 --- a/servers/Azure.Mcp.Server/docs/azmcp-commands.md +++ b/servers/Azure.Mcp.Server/docs/azmcp-commands.md @@ -2298,9 +2298,8 @@ azmcp mysql server param set --subscription \ # Without parameters: lists all PostgreSQL servers in the resource group # With --server: lists all databases on that server # With --server and --database: lists all tables in that database -# Database and table results are capped at 10,000 entries. When the database list is -# truncated a "... (output limited to 10,000 databases ...)" sentinel row is appended; -# when the table list is truncated, the response also includes "tablesTruncated": true. +# Database and table results are capped at 10,000 entries. When the results are truncated, +# the response includes "databasesTruncated": true or "tablesTruncated": true respectively. # ❌ Destructive | ✅ Idempotent | ❌ OpenWorld | ✅ ReadOnly | ❌ Secret | ❌ LocalRequired azmcp postgres list --subscription \ --resource-group \ diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs index 62f743d67f..2344c1a79e 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs @@ -91,13 +91,13 @@ public override async Task ExecuteAsync(CommandContext context, cancellationToken); context.Response.Results = ResponseResult.Create( - new(null, null, tableResult.Tables ?? [], tableResult.IsTruncated ? true : null), + new(null, null, tableResult.Tables ?? [], null, tableResult.IsTruncated ? true : null), PostgresJsonContext.Default.PostgresListCommandResult); } else if (!string.IsNullOrEmpty(options.Server)) { // List databases on specified server - List databases = await _postgresService.ListDatabasesAsync( + DatabaseListResult databaseResult = await _postgresService.ListDatabasesAsync( options.Subscription!, options.ResourceGroup!, options.AuthType!, @@ -107,7 +107,7 @@ public override async Task ExecuteAsync(CommandContext context, cancellationToken); context.Response.Results = ResponseResult.Create( - new(null, databases ?? [], null), + new(null, databaseResult.Databases ?? [], null, databaseResult.IsTruncated ? true : null, null), PostgresJsonContext.Default.PostgresListCommandResult); } else @@ -132,5 +132,5 @@ public override async Task ExecuteAsync(CommandContext context, return context.Response; } - public record PostgresListCommandResult(List? Servers, List? Databases, List? Tables, bool? TablesTruncated = null); + public record PostgresListCommandResult(List? Servers, List? Databases, List? Tables, bool? DatabasesTruncated = null, bool? TablesTruncated = null); } diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/DatabaseListResult.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/DatabaseListResult.cs new file mode 100644 index 0000000000..74d197258f --- /dev/null +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/DatabaseListResult.cs @@ -0,0 +1,6 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +namespace Azure.Mcp.Tools.Postgres.Services; + +public sealed record DatabaseListResult(List Databases, bool IsTruncated); diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs index c009d5dbf2..fd1602f747 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs @@ -7,7 +7,7 @@ namespace Azure.Mcp.Tools.Postgres.Services; public interface IPostgresService { - Task> ListDatabasesAsync( + Task ListDatabasesAsync( string subscriptionId, string resourceGroup, string authType, diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs index f7a54f01f4..17d6c18217 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs @@ -36,7 +36,7 @@ public class PostgresService( private readonly IEntraTokenProvider _entraTokenAuth = entraTokenAuth; private readonly IDbProvider _dbProvider = dbProvider; - private const int MaxRowCount = 10_000; + internal const int MaxRowCount = 10_000; private async Task GetEntraIdAccessTokenAsync(CancellationToken cancellationToken) { @@ -80,7 +80,7 @@ private string NormalizeServerName(string server) return server; } - public async Task> ListDatabasesAsync( + public async Task ListDatabasesAsync( string subscriptionId, string resourceGroup, string authType, @@ -96,25 +96,22 @@ public async Task> ListDatabasesAsync( var query = "SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname LIMIT @maxResults;"; await using IPostgresResource resource = await _dbProvider.GetPostgresResource(connectionString, authType, cancellationToken); await using NpgsqlCommand command = _dbProvider.GetCommand(query, resource); - // Cap at exactly MaxRowCount: truncation is signaled by appending a sentinel row to the returned list - // (mirrors MySqlService.ListDatabasesAsync). Tables use cap+1 instead because they signal truncation - // via a structured IsTruncated flag and need to observe an N+1th row to set it. - command.Parameters.AddWithValue("maxResults", MaxRowCount); + // Fetch cap+1 rows so we can detect truncation by observing whether an extra row exists, then trim it. + command.Parameters.AddWithValue("maxResults", MaxRowCount + 1); await using DbDataReader reader = await _dbProvider.ExecuteReaderAsync(command, cancellationToken); var dbs = new List(); - var dbCount = 0; - while (await reader.ReadAsync(cancellationToken) && dbCount < MaxRowCount) + while (await reader.ReadAsync(cancellationToken)) { dbs.Add(reader.GetString(0)); - dbCount++; } - if (dbCount >= MaxRowCount) + var isTruncated = dbs.Count > MaxRowCount; + if (isTruncated) { - dbs.Add($"... (output limited to {MaxRowCount:N0} databases for security and performance reasons)"); + dbs.RemoveRange(MaxRowCount, dbs.Count - MaxRowCount); } - return dbs; + return new DatabaseListResult(dbs, isTruncated); } public async Task> ExecuteQueryAsync( @@ -193,8 +190,6 @@ public async Task ListTablesAsync( await using IPostgresResource resource = await _dbProvider.GetPostgresResource(connectionString, authType, cancellationToken); await using NpgsqlCommand command = _dbProvider.GetCommand(query, resource); // Fetch cap+1 rows so we can detect truncation by observing whether an extra row exists, then trim it. - // Unlike ListDatabasesAsync (which signals truncation via a sentinel row appended to the list), - // tables return a structured TableListResult with an IsTruncated flag, so we need the extra row to set it. command.Parameters.AddWithValue("maxResults", MaxRowCount + 1); await using DbDataReader reader = await _dbProvider.ExecuteReaderAsync(command, cancellationToken); var tables = new List(); diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs index 11ccbffe30..d91487d43a 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs @@ -78,7 +78,7 @@ public async Task ExecuteAsync_ListsDatabases_WhenServerProvided() null, "server1", Arg.Any()) - .Returns(expectedDatabases); + .Returns(new DatabaseListResult(expectedDatabases, false)); var response = await ExecuteCommandAsync( "--subscription", "sub123", @@ -92,6 +92,7 @@ public async Task ExecuteAsync_ListsDatabases_WhenServerProvided() Assert.Null(result.Servers); Assert.Equal(expectedDatabases, result.Databases); Assert.Null(result.Tables); + Assert.Null(result.DatabasesTruncated); } [Fact] @@ -156,6 +157,35 @@ public async Task ExecuteAsync_SetsTablesTruncated_WhenServiceReportsTruncation( Assert.True(result.TablesTruncated); } + [Fact] + public async Task ExecuteAsync_SetsDatabasesTruncated_WhenServiceReportsTruncation() + { + var expectedDatabases = new List { "db1", "db2", "db3" }; + Service.ListDatabasesAsync( + "sub123", + "rg1", + AuthTypes.MicrosoftEntra, + "user1", + null, + "server1", + Arg.Any()) + .Returns(new DatabaseListResult(expectedDatabases, true)); + + var response = await ExecuteCommandAsync( + "--subscription", "sub123", + "--resource-group", "rg1", + "--user", "user1", + $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, + "--server", "server1"); + + var result = ValidateAndDeserializeResponse(response, PostgresJsonContext.Default.PostgresListCommandResult); + + Assert.Null(result.Servers); + Assert.Equal(expectedDatabases, result.Databases); + Assert.Null(result.Tables); + Assert.True(result.DatabasesTruncated); + } + [Fact] public async Task ExecuteAsync_ReturnsNull_WhenNoServersExist() { @@ -184,7 +214,7 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoDatabasesExist() null, "server1", Arg.Any()) - .Returns([]); + .Returns(new DatabaseListResult([], false)); var response = await ExecuteCommandAsync( "--subscription", "sub123", diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs index c085b4244f..e71df9736d 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs @@ -17,7 +17,7 @@ namespace Azure.Mcp.Tools.Postgres.Tests.Services; public class PostgresServiceRowLimitTests { - private const int MaxRowCount = 10_000; + private const int MaxRowCount = PostgresService.MaxRowCount; private readonly IResourceGroupService _resourceGroupService = Substitute.For(); private readonly ISubscriptionService _subscriptionService = Substitute.For(); @@ -54,30 +54,41 @@ private void StubReader(int rowCount, string columnName) } [Fact] - public async Task ListDatabasesAsync_UnderCap_DoesNotAppendSentinel() + public async Task ListDatabasesAsync_UnderCap_ReturnsAllAndNotTruncated() { StubReader(rowCount: 3, columnName: "datname"); var result = await _postgresService.ListDatabasesAsync( SubscriptionId, ResourceGroup, AuthType, User, null, Server, TestContext.Current.CancellationToken); - Assert.Equal(3, result.Count); - Assert.DoesNotContain(result, r => r.StartsWith("... (output limited", StringComparison.Ordinal)); + Assert.Equal(3, result.Databases.Count); + Assert.False(result.IsTruncated); } [Fact] - public async Task ListDatabasesAsync_AtCap_AppendsSentinelRow() + public async Task ListDatabasesAsync_AtCap_ReturnsCapRowsAndNotTruncatedWhenNoExtra() { - // FakeDbDataReader ignores SQL LIMIT, so simulate the worst case by returning exactly MaxRowCount rows. + // Reader returns exactly MaxRowCount rows — boundary case, nothing beyond the cap. StubReader(rowCount: MaxRowCount, columnName: "datname"); var result = await _postgresService.ListDatabasesAsync( SubscriptionId, ResourceGroup, AuthType, User, null, Server, TestContext.Current.CancellationToken); - // MaxRowCount real rows + 1 sentinel row (mirrors MySQL behavior). - Assert.Equal(MaxRowCount + 1, result.Count); - Assert.StartsWith("... (output limited", result[^1]); - Assert.Contains("10,000", result[^1]); + Assert.Equal(MaxRowCount, result.Databases.Count); + Assert.False(result.IsTruncated); + } + + [Fact] + public async Task ListDatabasesAsync_OverCap_ReturnsCapRowsAndIsTruncated() + { + // Reader returns MaxRowCount + 1 (the cap+1 LIMIT in production); detect truncation via the extra read. + StubReader(rowCount: MaxRowCount + 1, columnName: "datname"); + + var result = await _postgresService.ListDatabasesAsync( + SubscriptionId, ResourceGroup, AuthType, User, null, Server, TestContext.Current.CancellationToken); + + Assert.Equal(MaxRowCount, result.Databases.Count); + Assert.True(result.IsTruncated); } [Fact] From 576a1876f2fa1a91b323cf4c7718f65f9a8c492e Mon Sep 17 00:00:00 2001 From: Victor Colin Amador Date: Mon, 22 Jun 2026 16:49:46 -0600 Subject: [PATCH 3/4] Consolidate truncation flag to single ResultsTruncated (#472) A postgres list call returns servers, databases, or tables - never more than one - so a single ResultsTruncated flag is sufficient instead of separate DatabasesTruncated/TablesTruncated fields. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../changelog-entries/1782105637872.yaml | 2 +- servers/Azure.Mcp.Server/docs/azmcp-commands.md | 2 +- .../src/Commands/PostgresListCommand.cs | 6 +++--- .../PostgresListCommandTests.cs | 14 +++++++------- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml b/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml index bfa4730b32..6f0d6c24a6 100644 --- a/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml +++ b/servers/Azure.Mcp.Server/changelog-entries/1782105637872.yaml @@ -1,3 +1,3 @@ changes: - section: "Bugs Fixed" - description: "Cap 'azmcp postgres list' database and table results at 10,000 entries to prevent unbounded enumeration on large servers. When the results are truncated, the response includes a 'databasesTruncated: true' or 'tablesTruncated: true' flag." + description: "Cap 'azmcp postgres list' database and table results at 10,000 entries to prevent unbounded enumeration on large servers. When the results are truncated, the response includes a 'resultsTruncated: true' flag." diff --git a/servers/Azure.Mcp.Server/docs/azmcp-commands.md b/servers/Azure.Mcp.Server/docs/azmcp-commands.md index 7f3d11cacf..bc6a17b051 100644 --- a/servers/Azure.Mcp.Server/docs/azmcp-commands.md +++ b/servers/Azure.Mcp.Server/docs/azmcp-commands.md @@ -2299,7 +2299,7 @@ azmcp mysql server param set --subscription \ # With --server: lists all databases on that server # With --server and --database: lists all tables in that database # Database and table results are capped at 10,000 entries. When the results are truncated, -# the response includes "databasesTruncated": true or "tablesTruncated": true respectively. +# the response includes "resultsTruncated": true. # ❌ Destructive | ✅ Idempotent | ❌ OpenWorld | ✅ ReadOnly | ❌ Secret | ❌ LocalRequired azmcp postgres list --subscription \ --resource-group \ diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs index 2344c1a79e..fa28e132f4 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs @@ -91,7 +91,7 @@ public override async Task ExecuteAsync(CommandContext context, cancellationToken); context.Response.Results = ResponseResult.Create( - new(null, null, tableResult.Tables ?? [], null, tableResult.IsTruncated ? true : null), + new(null, null, tableResult.Tables ?? [], tableResult.IsTruncated ? true : null), PostgresJsonContext.Default.PostgresListCommandResult); } else if (!string.IsNullOrEmpty(options.Server)) @@ -107,7 +107,7 @@ public override async Task ExecuteAsync(CommandContext context, cancellationToken); context.Response.Results = ResponseResult.Create( - new(null, databaseResult.Databases ?? [], null, databaseResult.IsTruncated ? true : null, null), + new(null, databaseResult.Databases ?? [], null, databaseResult.IsTruncated ? true : null), PostgresJsonContext.Default.PostgresListCommandResult); } else @@ -132,5 +132,5 @@ public override async Task ExecuteAsync(CommandContext context, return context.Response; } - public record PostgresListCommandResult(List? Servers, List? Databases, List? Tables, bool? DatabasesTruncated = null, bool? TablesTruncated = null); + public record PostgresListCommandResult(List? Servers, List? Databases, List? Tables, bool? ResultsTruncated = null); } diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs index d91487d43a..77fcfd18a1 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresListCommandTests.cs @@ -92,7 +92,7 @@ public async Task ExecuteAsync_ListsDatabases_WhenServerProvided() Assert.Null(result.Servers); Assert.Equal(expectedDatabases, result.Databases); Assert.Null(result.Tables); - Assert.Null(result.DatabasesTruncated); + Assert.Null(result.ResultsTruncated); } [Fact] @@ -123,11 +123,11 @@ public async Task ExecuteAsync_ListsTables_WhenServerAndDatabaseProvided() Assert.Null(result.Servers); Assert.Null(result.Databases); Assert.Equal(expectedTables, result.Tables); - Assert.Null(result.TablesTruncated); + Assert.Null(result.ResultsTruncated); } [Fact] - public async Task ExecuteAsync_SetsTablesTruncated_WhenServiceReportsTruncation() + public async Task ExecuteAsync_SetsResultsTruncated_WhenTableResultsAreTruncated() { var expectedTables = new List { "users", "products", "orders" }; Service.ListTablesAsync( @@ -154,11 +154,11 @@ public async Task ExecuteAsync_SetsTablesTruncated_WhenServiceReportsTruncation( Assert.Null(result.Servers); Assert.Null(result.Databases); Assert.Equal(expectedTables, result.Tables); - Assert.True(result.TablesTruncated); + Assert.True(result.ResultsTruncated); } [Fact] - public async Task ExecuteAsync_SetsDatabasesTruncated_WhenServiceReportsTruncation() + public async Task ExecuteAsync_SetsResultsTruncated_WhenDatabaseResultsAreTruncated() { var expectedDatabases = new List { "db1", "db2", "db3" }; Service.ListDatabasesAsync( @@ -183,7 +183,7 @@ public async Task ExecuteAsync_SetsDatabasesTruncated_WhenServiceReportsTruncati Assert.Null(result.Servers); Assert.Equal(expectedDatabases, result.Databases); Assert.Null(result.Tables); - Assert.True(result.DatabasesTruncated); + Assert.True(result.ResultsTruncated); } [Fact] @@ -259,7 +259,7 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoTablesExist() Assert.Null(result.Databases); Assert.NotNull(result.Tables); Assert.Empty(result.Tables); - Assert.Null(result.TablesTruncated); + Assert.Null(result.ResultsTruncated); } [Fact] From 97dc8c5abfbcfb8c493c597ed55190117e32d13e Mon Sep 17 00:00:00 2001 From: vcolin7 Date: Mon, 22 Jun 2026 15:52:03 -0700 Subject: [PATCH 4/4] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../Services/PostgresServiceRowLimitTests.cs | 1 - 1 file changed, 1 deletion(-) diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs index e71df9736d..a23a0d8141 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceRowLimitTests.cs @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -using System.Data.Common; using Azure.Mcp.Core.Services.Azure.ResourceGroup; using Azure.Mcp.Core.Services.Azure.Subscription; using Azure.Mcp.Core.Services.Azure.Tenant;