From 09c921cf7b3d2587608100bb5c057cf9d9f1e659 Mon Sep 17 00:00:00 2001 From: Victor Colin Amador Date: Thu, 11 Jun 2026 18:47:52 -0700 Subject: [PATCH 1/4] Fix Postgres data-plane params and add multi-schema table listing Resolves #471 by removing the unused subscriptionId and resourceGroup parameters from the ListDatabasesAsync, ExecuteQueryAsync, ListTablesAsync and GetTableSchemaAsync data-plane methods (these connect directly via Npgsql) and updating all callers. ARM server-config methods and the --subscription/--resource-group CLI options are unchanged. Resolves #469 by adding an optional --schema parameter (default 'public') to the table-listing path. ListTablesAsync now accepts a schema argument and uses a parameterized query: SELECT table_name FROM information_schema.tables WHERE table_schema = @schema ORDER BY table_name; Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../vcolin7-postgres-schema-listing.yaml | 3 + .../Azure.Mcp.Server/docs/azmcp-commands.md | 5 +- .../Azure.Mcp.Server/docs/e2eTestPrompts.md | 1 + .../Commands/Database/DatabaseQueryCommand.cs | 2 - .../src/Commands/PostgresListCommand.cs | 7 +- .../Commands/Table/TableSchemaGetCommand.cs | 2 - .../src/Options/BasePostgresOptions.cs | 3 + .../src/Options/PostgresOptionDefinitions.cs | 10 +++ .../src/Services/IPostgresService.cs | 9 +-- .../src/Services/PostgresService.cs | 12 +-- .../Database/DatabaseQueryCommandTests.cs | 9 ++- .../PostgresListCommandTests.cs | 74 ++++++++++++++++--- ...esServiceConnectionStringInjectionTests.cs | 15 ++-- .../PostgresServiceParameterizedQueryTests.cs | 37 +++------- ...ostgresServiceServerNameValidationTests.cs | 15 ++-- .../Services/PostgresServiceTests.cs | 34 +++++++-- .../Table/TableSchemaGetCommandTests.cs | 4 +- 17 files changed, 149 insertions(+), 93 deletions(-) create mode 100644 servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml diff --git a/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml b/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml new file mode 100644 index 0000000000..b82a60ddfc --- /dev/null +++ b/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml @@ -0,0 +1,3 @@ +changes: + - section: "Features Added" + description: "Added an optional `--schema` parameter to `azmcp postgres list` for listing tables in non-public PostgreSQL schemas (defaults to `public`), and removed unused `subscription`/`resource-group` parameters from the PostgreSQL data-plane service methods." diff --git a/servers/Azure.Mcp.Server/docs/azmcp-commands.md b/servers/Azure.Mcp.Server/docs/azmcp-commands.md index c5c204ce33..c49056c1c4 100644 --- a/servers/Azure.Mcp.Server/docs/azmcp-commands.md +++ b/servers/Azure.Mcp.Server/docs/azmcp-commands.md @@ -2264,13 +2264,14 @@ azmcp mysql server param set --subscription \ # Hierarchical list command for PostgreSQL resources # 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 +# With --server and --database: lists all tables in that database (optionally scoped to a --schema, defaults to 'public') # ❌ Destructive | ✅ Idempotent | ❌ OpenWorld | ✅ ReadOnly | ❌ Secret | ❌ LocalRequired azmcp postgres list --subscription \ --resource-group \ --user \ [--server ] \ - [--database ] + [--database ] \ + [--schema ] # Execute a query on a PostgreSQL database # ❌ Destructive | ✅ Idempotent | ❌ OpenWorld | ✅ ReadOnly | ❌ Secret | ❌ LocalRequired diff --git a/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md b/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md index 08bcf795bf..6361931c75 100644 --- a/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md +++ b/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md @@ -418,6 +418,7 @@ This file contains prompts used for end-to-end testing to ensure each tool is in | postgres_list | Show me the PostgreSQL databases in server \ | | postgres_list | List all tables in the PostgreSQL database \ in server \ | | postgres_list | Show me the tables in the PostgreSQL database \ in server \ | +| postgres_list | List all tables in the \ schema of the PostgreSQL database \ in server \ | | postgres_database_query | Show me all items that contain the word \ in the PostgreSQL database \ in server \ | | postgres_server_config_get | Show me the configuration of PostgreSQL server \ | | postgres_server_param_get | Show me if the parameter my PostgreSQL server \ has replication enabled | diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Commands/Database/DatabaseQueryCommand.cs b/tools/Azure.Mcp.Tools.Postgres/src/Commands/Database/DatabaseQueryCommand.cs index 9497ace9a0..baa2e4f2ce 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Commands/Database/DatabaseQueryCommand.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Commands/Database/DatabaseQueryCommand.cs @@ -54,8 +54,6 @@ public override async Task ExecuteAsync(CommandContext context, // Validate the query early to avoid sending unsafe SQL to the server. SqlQueryValidator.EnsureReadOnlySelect(options.Query); List queryResult = await _postgresService.ExecuteQueryAsync( - options.Subscription!, - options.ResourceGroup!, options.AuthType!, options.User!, options.Password, diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs index ec6306aa47..bd60f2fd10 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs @@ -33,6 +33,7 @@ protected override void RegisterOptions(Command command) command.Options.Add(PostgresOptionDefinitions.User.AsOptional()); command.Options.Add(PostgresOptionDefinitions.ServerOptional); command.Options.Add(PostgresOptionDefinitions.DatabaseOptional); + command.Options.Add(PostgresOptionDefinitions.Schema); command.Options.Add(PostgresOptionDefinitions.AuthType); command.Options.Add(PostgresOptionDefinitions.Password); command.Validators.Add(result => @@ -58,6 +59,7 @@ protected override BasePostgresOptions BindOptions(ParseResult parseResult) var options = base.BindOptions(parseResult); options.Server = parseResult.GetValueOrDefault(PostgresOptionDefinitions.ServerOptional.Name); options.Database = parseResult.GetValueOrDefault(PostgresOptionDefinitions.DatabaseOptional.Name); + options.Schema = parseResult.GetValueOrDefault(PostgresOptionDefinitions.Schema.Name); options.AuthType = parseResult.GetValueOrDefault(PostgresOptionDefinitions.AuthType.Name); options.Password = parseResult.GetValueOrDefault(PostgresOptionDefinitions.Password.Name); return options; @@ -81,13 +83,12 @@ public override async Task ExecuteAsync(CommandContext context, { // List tables in specified database List tables = await _postgresService.ListTablesAsync( - options.Subscription!, - options.ResourceGroup!, options.AuthType!, options.User!, options.Password, options.Server!, options.Database!, + string.IsNullOrEmpty(options.Schema) ? "public" : options.Schema, cancellationToken); context.Response.Results = ResponseResult.Create( @@ -98,8 +99,6 @@ public override async Task ExecuteAsync(CommandContext context, { // List databases on specified server List databases = await _postgresService.ListDatabasesAsync( - options.Subscription!, - options.ResourceGroup!, options.AuthType!, options.User!, options.Password, diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Commands/Table/TableSchemaGetCommand.cs b/tools/Azure.Mcp.Tools.Postgres/src/Commands/Table/TableSchemaGetCommand.cs index 95d3dbe187..e2579bd2cf 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Commands/Table/TableSchemaGetCommand.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Commands/Table/TableSchemaGetCommand.cs @@ -52,8 +52,6 @@ public override async Task ExecuteAsync(CommandContext context, { List schema = await _postgresService.GetTableSchemaAsync( - options.Subscription!, - options.ResourceGroup!, options.AuthType!, options.User!, options.Password, diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Options/BasePostgresOptions.cs b/tools/Azure.Mcp.Tools.Postgres/src/Options/BasePostgresOptions.cs index a6e0ee30c3..bc3f757c9f 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Options/BasePostgresOptions.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Options/BasePostgresOptions.cs @@ -22,4 +22,7 @@ public class BasePostgresOptions : SubscriptionOptions [JsonPropertyName(PostgresOptionDefinitions.DatabaseName)] public string? Database { get; set; } + + [JsonPropertyName(PostgresOptionDefinitions.SchemaName)] + public string? Schema { get; set; } } diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Options/PostgresOptionDefinitions.cs b/tools/Azure.Mcp.Tools.Postgres/src/Options/PostgresOptionDefinitions.cs index 7588499067..b35166b356 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Options/PostgresOptionDefinitions.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Options/PostgresOptionDefinitions.cs @@ -10,6 +10,7 @@ public static class PostgresOptionDefinitions public const string PasswordText = "password"; public const string ServerName = "server"; public const string DatabaseName = "database"; + public const string SchemaName = "schema"; public const string TableName = "table"; public const string QueryText = "query"; public const string ParamName = "param"; @@ -72,6 +73,15 @@ public static class PostgresOptionDefinitions Description = "The PostgreSQL database to list tables from (optional, requires --server)." }; + public static readonly Option Schema = new( + $"--{SchemaName}" + ) + { + Description = "The PostgreSQL schema to list tables from when listing tables (optional, defaults to 'public').", + Arity = ArgumentArity.ZeroOrOne, + Required = false + }; + public static readonly Option Table = new( $"--{TableName}" ) diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs index 398326f949..93003a2517 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/IPostgresService.cs @@ -6,8 +6,6 @@ namespace Azure.Mcp.Tools.Postgres.Services; public interface IPostgresService { Task> ListDatabasesAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, @@ -15,8 +13,6 @@ Task> ListDatabasesAsync( CancellationToken cancellationToken); Task> ExecuteQueryAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, @@ -26,18 +22,15 @@ Task> ExecuteQueryAsync( CancellationToken cancellationToken); Task> ListTablesAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, string server, string database, + string schema, CancellationToken cancellationToken); Task> GetTableSchemaAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs b/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs index a846e59e67..66f76bd809 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs @@ -78,8 +78,6 @@ private string NormalizeServerName(string server) } public async Task> ListDatabasesAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, @@ -103,8 +101,6 @@ public async Task> ListDatabasesAsync( } public async Task> ExecuteQueryAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, @@ -161,22 +157,22 @@ public async Task> ExecuteQueryAsync( } public async Task> ListTablesAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, string server, string database, + string schema, CancellationToken cancellationToken) { string? passwordToUse = await GetPassword(authType, password, cancellationToken); 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 = @schema ORDER BY table_name;"; await using IPostgresResource resource = await _dbProvider.GetPostgresResource(connectionString, authType, cancellationToken); await using NpgsqlCommand command = _dbProvider.GetCommand(query, resource); + command.Parameters.AddWithValue("schema", schema); await using DbDataReader reader = await _dbProvider.ExecuteReaderAsync(command, cancellationToken); var tables = new List(); while (await reader.ReadAsync(cancellationToken)) @@ -187,8 +183,6 @@ public async Task> ListTablesAsync( } public async Task> GetTableSchemaAsync( - string subscriptionId, - string resourceGroup, string authType, string user, string? password, diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs index f4e549f71b..6a0772efc8 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs @@ -22,7 +22,7 @@ public async Task ExecuteAsync_ReturnsQueryResults_WhenQueryIsValid() { var expectedResults = new List { "result1", "result2" }; - Service.ExecuteQueryAsync("sub123", "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "SELECT * FROM test;", Arg.Any()) + Service.ExecuteQueryAsync(AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "SELECT * FROM test;", Arg.Any()) .Returns(expectedResults); var response = await ExecuteCommandAsync( @@ -41,7 +41,7 @@ public async Task ExecuteAsync_ReturnsQueryResults_WhenQueryIsValid() [Fact] public async Task ExecuteAsync_ReturnsEmpty_WhenQueryFails() { - Service.ExecuteQueryAsync("sub123", "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "SELECT * FROM test;", Arg.Any()) + Service.ExecuteQueryAsync(AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "SELECT * FROM test;", Arg.Any()) .Returns([]); var response = await ExecuteCommandAsync( @@ -137,7 +137,7 @@ public async Task ExecuteAsync_InvalidQuery_ValidationError(string badQuery) Assert.NotNull(response); Assert.Equal(HttpStatusCode.BadRequest, response.Status); // CommandValidationException => 400 // Service should never be called for invalid queries. - await Service.DidNotReceive().ExecuteQueryAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()); + await Service.DidNotReceive().ExecuteQueryAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()); } [Fact] @@ -155,6 +155,7 @@ public async Task ExecuteAsync_LongQuery_ValidationError() Assert.NotNull(response); Assert.Equal(HttpStatusCode.BadRequest, response.Status); - await Service.DidNotReceive().ExecuteQueryAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()); + await Service.DidNotReceive().ExecuteQueryAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()); } } + 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..19140803aa 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 @@ -71,8 +71,6 @@ public async Task ExecuteAsync_ListsDatabases_WhenServerProvided() { var expectedDatabases = new List { "db1", "db2", "db3" }; Service.ListDatabasesAsync( - "sub123", - "rg1", AuthTypes.MicrosoftEntra, "user1", null, @@ -99,13 +97,12 @@ public async Task ExecuteAsync_ListsTables_WhenServerAndDatabaseProvided() { var expectedTables = new List { "users", "products", "orders" }; Service.ListTablesAsync( - "sub123", - "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db1", + "public", Arg.Any()) .Returns(expectedTables); @@ -145,8 +142,6 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoServersExist() public async Task ExecuteAsync_ReturnsNull_WhenNoDatabasesExist() { Service.ListDatabasesAsync( - "sub123", - "rg1", AuthTypes.MicrosoftEntra, "user1", null, @@ -173,13 +168,12 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoDatabasesExist() public async Task ExecuteAsync_ReturnsNull_WhenNoTablesExist() { Service.ListTablesAsync( - "sub123", - "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db1", + "public", Arg.Any()) .Returns([]); @@ -199,6 +193,65 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoTablesExist() Assert.Empty(result.Tables); } + [Fact] + public async Task ExecuteAsync_ListsTablesWithSpecifiedSchema_WhenSchemaProvided() + { + var expectedTables = new List { "audit_log", "events" }; + Service.ListTablesAsync( + AuthTypes.MicrosoftEntra, + "user1", + null, + "server1", + "db1", + "analytics", + Arg.Any()) + .Returns(expectedTables); + + var response = await ExecuteCommandAsync( + "--subscription", "sub123", + "--resource-group", "rg1", + "--user", "user1", + $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, + "--server", "server1", + "--database", "db1", + $"--{PostgresOptionDefinitions.SchemaName}", "analytics"); + + var result = ValidateAndDeserializeResponse(response, PostgresJsonContext.Default.PostgresListCommandResult); + + Assert.Null(result.Servers); + Assert.Null(result.Databases); + Assert.Equal(expectedTables, result.Tables); + await Service.Received(1).ListTablesAsync( + AuthTypes.MicrosoftEntra, "user1", null, "server1", "db1", "analytics", Arg.Any()); + } + + [Fact] + public async Task ExecuteAsync_ListsTablesWithPublicSchema_WhenSchemaOmitted() + { + Service.ListTablesAsync( + AuthTypes.MicrosoftEntra, + "user1", + null, + "server1", + "db1", + "public", + Arg.Any()) + .Returns(["users"]); + + var response = await ExecuteCommandAsync( + "--subscription", "sub123", + "--resource-group", "rg1", + "--user", "user1", + $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, + "--server", "server1", + "--database", "db1"); + + ValidateAndDeserializeResponse(response, PostgresJsonContext.Default.PostgresListCommandResult); + + await Service.Received(1).ListTablesAsync( + AuthTypes.MicrosoftEntra, "user1", null, "server1", "db1", "public", Arg.Any()); + } + [Fact] public async Task ExecuteAsync_ReturnsError_WhenListServersThrows() { @@ -220,8 +273,6 @@ public async Task ExecuteAsync_ReturnsError_WhenListDatabasesThrows() { var expectedError = "Test error. To mitigate this issue, please refer to the troubleshooting guidelines here at https://aka.ms/azmcp/troubleshooting."; Service.ListDatabasesAsync( - "sub123", - "rg1", AuthTypes.MicrosoftEntra, "user1", null, @@ -246,13 +297,12 @@ public async Task ExecuteAsync_ReturnsError_WhenListTablesThrows() { var expectedError = "Test error. To mitigate this issue, please refer to the troubleshooting guidelines here at https://aka.ms/azmcp/troubleshooting."; Service.ListTablesAsync( - "sub123", - "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db1", + "public", Arg.Any()) .ThrowsAsync(new Exception("Test error")); diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceConnectionStringInjectionTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceConnectionStringInjectionTests.cs index b7ccc8b832..d19e54cfcf 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceConnectionStringInjectionTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceConnectionStringInjectionTests.cs @@ -68,7 +68,7 @@ public async Task ExecuteQueryAsync_WithInjectedDatabase_DoesNotOverrideHost(str { // Act await _postgresService.ExecuteQueryAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, "legitimate-server", maliciousDatabase, "SELECT 1", TestContext.Current.CancellationToken); @@ -85,8 +85,8 @@ public async Task ListTablesAsync_WithInjectedDatabase_DoesNotOverrideHost(strin { // Act await _postgresService.ListTablesAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, - "legitimate-server", maliciousDatabase, + AuthTypes.MicrosoftEntra, "test-user", null, + "legitimate-server", maliciousDatabase, "public", TestContext.Current.CancellationToken); // Assert @@ -102,7 +102,7 @@ public async Task GetTableSchemaAsync_WithInjectedDatabase_DoesNotOverrideHost(s { // Act await _postgresService.GetTableSchemaAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, "legitimate-server", maliciousDatabase, "some_table", TestContext.Current.CancellationToken); @@ -121,7 +121,7 @@ public async Task ExecuteQueryAsync_WithSemicolonInDatabase_PreservesOriginalHos // Act await _postgresService.ExecuteQueryAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, legitimateServer, maliciousDatabase, "SELECT 1", TestContext.Current.CancellationToken); @@ -142,7 +142,7 @@ public async Task ExecuteQueryAsync_WithSslDowngradeInDatabase_EnforcesSslRequir // Act await _postgresService.ExecuteQueryAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, "safe-server", maliciousDatabase, "SELECT 1", TestContext.Current.CancellationToken); @@ -160,7 +160,7 @@ public async Task ExecuteQueryAsync_NormalInputs_EnforcesSslRequire() { // Act — exercise the normal (non-injection) code path await _postgresService.ExecuteQueryAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, "safe-server", "mydb", "SELECT 1", TestContext.Current.CancellationToken); @@ -171,3 +171,4 @@ await _postgresService.ExecuteQueryAsync( Assert.Equal(SslMode.Require, parsed.SslMode); } } + diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceParameterizedQueryTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceParameterizedQueryTests.cs index 2723a0426b..202a44c55d 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceParameterizedQueryTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceParameterizedQueryTests.cs @@ -61,8 +61,6 @@ public void GetTableSchemaAsync_ParameterizedQuery_ShouldHandleTableNamesCorrect // but we can verify the method signature and that it doesn't throw for valid inputs // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -71,7 +69,7 @@ public void GetTableSchemaAsync_ParameterizedQuery_ShouldHandleTableNamesCorrect // Act & Assert - Method should accept these parameters without throwing // The actual parameterization is tested through integration tests - var task = _postgresService.GetTableSchemaAsync(subscriptionId, resourceGroup, authType, user, password, server, database, tableName, TestContext.Current.CancellationToken); + var task = _postgresService.GetTableSchemaAsync(authType, user, password, server, database, tableName, TestContext.Current.CancellationToken); // The method will fail at the connection stage, but that's expected in unit tests // What we're testing is that the method signature accepts these parameters correctly @@ -89,8 +87,6 @@ public void GetTableSchemaAsync_WithSQLInjectionAttempts_ShouldNotCauseSecurityV // With parameterized queries, these should be treated as literal table names // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -100,7 +96,7 @@ public void GetTableSchemaAsync_WithSQLInjectionAttempts_ShouldNotCauseSecurityV // Act & Assert // The method should not throw due to SQL injection attempts // With proper parameterization, malicious input is treated as a literal table name - var task = _postgresService.GetTableSchemaAsync(subscriptionId, resourceGroup, authType, user, password, server, database, maliciousTableName, TestContext.Current.CancellationToken); + var task = _postgresService.GetTableSchemaAsync(authType, user, password, server, database, maliciousTableName, TestContext.Current.CancellationToken); // The method will fail at the connection stage, but importantly, // it won't fail due to SQL parsing errors caused by injection attempts @@ -112,8 +108,6 @@ public void GetTableSchemaAsync_WithSpecialCharacters_ShouldHandleSafely() { // Arrange string tableName = "table_with_special_chars_123!@#$%^&*()"; - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -122,7 +116,7 @@ public void GetTableSchemaAsync_WithSpecialCharacters_ShouldHandleSafely() // Act & Assert // Should handle special characters safely through parameterization - var task = _postgresService.GetTableSchemaAsync(subscriptionId, resourceGroup, authType, user, password, server, database, tableName, TestContext.Current.CancellationToken); + var task = _postgresService.GetTableSchemaAsync(authType, user, password, server, database, tableName, TestContext.Current.CancellationToken); Assert.NotNull(task); } @@ -132,8 +126,6 @@ public void GetTableSchemaAsync_WithSpecialCharacters_ShouldHandleSafely() public void GetTableSchemaAsync_WithEmptyTableName_ShouldHandleGracefully(string tableName) { // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -142,7 +134,7 @@ public void GetTableSchemaAsync_WithEmptyTableName_ShouldHandleGracefully(string // Act & Assert // Should handle empty/whitespace table names without security issues - var task = _postgresService.GetTableSchemaAsync(subscriptionId, resourceGroup, authType, user, password, server, database, tableName, TestContext.Current.CancellationToken); + var task = _postgresService.GetTableSchemaAsync(authType, user, password, server, database, tableName, TestContext.Current.CancellationToken); Assert.NotNull(task); } @@ -150,8 +142,6 @@ public void GetTableSchemaAsync_WithEmptyTableName_ShouldHandleGracefully(string public void GetTableSchemaAsync_WithNullTableName_ShouldHandleGracefully() { // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -160,7 +150,7 @@ public void GetTableSchemaAsync_WithNullTableName_ShouldHandleGracefully() // Act & Assert // Should handle null table name without security issues - var task = _postgresService.GetTableSchemaAsync(subscriptionId, resourceGroup, authType, user, password, server, database, null!, TestContext.Current.CancellationToken); + var task = _postgresService.GetTableSchemaAsync(authType, user, password, server, database, null!, TestContext.Current.CancellationToken); Assert.NotNull(task); } @@ -169,8 +159,6 @@ public void ExecuteQueryAsync_CallsValidationBeforeExecution() { // This test verifies that query validation is called before execution // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -180,7 +168,7 @@ public void ExecuteQueryAsync_CallsValidationBeforeExecution() // Act & Assert // The method should fail validation before attempting to connect to database - var task = _postgresService.ExecuteQueryAsync(subscriptionId, resourceGroup, authType, user, password, server, database, maliciousQuery, TestContext.Current.CancellationToken); + var task = _postgresService.ExecuteQueryAsync(authType, user, password, server, database, maliciousQuery, TestContext.Current.CancellationToken); // We expect this to eventually throw due to validation, not due to database connection // The validation should catch dangerous queries before any database interaction @@ -194,8 +182,6 @@ public void ExecuteQueryAsync_CallsValidationBeforeExecution() public void ExecuteQueryAsync_WithValidQueries_ShouldPassValidation(string validQuery) { // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -204,7 +190,7 @@ public void ExecuteQueryAsync_WithValidQueries_ShouldPassValidation(string valid // Act & Assert // Valid queries should pass validation and proceed to connection attempt - var task = _postgresService.ExecuteQueryAsync(subscriptionId, resourceGroup, authType, user, password, server, database, validQuery, TestContext.Current.CancellationToken); + var task = _postgresService.ExecuteQueryAsync(authType, user, password, server, database, validQuery, TestContext.Current.CancellationToken); Assert.NotNull(task); } @@ -215,8 +201,6 @@ public void ExecuteQueryAsync_WithVectorSimilaritySearchQuery_ShouldHandleComple // and Azure OpenAI embeddings are properly handled // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -234,7 +218,7 @@ ORDER BY similarity // Act & Assert // The method should accept complex queries with vector operations and Azure OpenAI functions - var task = _postgresService.ExecuteQueryAsync(subscriptionId, resourceGroup, authType, user, password, server, database, vectorQuery, TestContext.Current.CancellationToken); + var task = _postgresService.ExecuteQueryAsync(authType, user, password, server, database, vectorQuery, TestContext.Current.CancellationToken); // Verify the task is created successfully (will fail at connection stage in unit test) Assert.NotNull(task); @@ -250,8 +234,6 @@ public void ExecuteQueryAsync_WithVariousVectorOperators_ShouldHandleCorrectly(s // PostgreSQL pgvector supports multiple operators: <=> (cosine), <-> (L2), <#> (inner product) // Arrange - string subscriptionId = "test-sub"; - string resourceGroup = "test-rg"; string authType = AuthTypes.MicrosoftEntra; string user = "test-user"; string? password = null; @@ -260,7 +242,8 @@ public void ExecuteQueryAsync_WithVariousVectorOperators_ShouldHandleCorrectly(s // Act & Assert // Should accept queries with various vector similarity operators - var task = _postgresService.ExecuteQueryAsync(subscriptionId, resourceGroup, authType, user, password, server, database, vectorQuery, TestContext.Current.CancellationToken); + var task = _postgresService.ExecuteQueryAsync(authType, user, password, server, database, vectorQuery, TestContext.Current.CancellationToken); Assert.NotNull(task); } } + diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceServerNameValidationTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceServerNameValidationTests.cs index 453918cace..26732ce9a3 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceServerNameValidationTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceServerNameValidationTests.cs @@ -58,7 +58,7 @@ public async Task ExecuteQueryAsync_WithNonAzureServerFQDN_ThrowsArgumentExcepti { var ex = await Assert.ThrowsAsync(() => _postgresService.ExecuteQueryAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, maliciousServer, "testdb", "SELECT 1", TestContext.Current.CancellationToken)); @@ -72,7 +72,7 @@ public async Task ListDatabasesAsync_WithNonAzureServerFQDN_ThrowsArgumentExcept { var ex = await Assert.ThrowsAsync(() => _postgresService.ListDatabasesAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, maliciousServer, TestContext.Current.CancellationToken)); @@ -86,8 +86,8 @@ public async Task ListTablesAsync_WithNonAzureServerFQDN_ThrowsArgumentException { var ex = await Assert.ThrowsAsync(() => _postgresService.ListTablesAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, - maliciousServer, "testdb", + AuthTypes.MicrosoftEntra, "test-user", null, + maliciousServer, "testdb", "public", TestContext.Current.CancellationToken)); Assert.Contains("not a valid Azure Database for PostgreSQL hostname", ex.Message); @@ -100,7 +100,7 @@ public async Task GetTableSchemaAsync_WithNonAzureServerFQDN_ThrowsArgumentExcep { var ex = await Assert.ThrowsAsync(() => _postgresService.GetTableSchemaAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, maliciousServer, "testdb", "test_table", TestContext.Current.CancellationToken)); @@ -116,7 +116,7 @@ public async Task ExecuteQueryAsync_WithValidAzureServerFQDN_DoesNotThrow(string { // Should not throw - valid Azure PostgreSQL FQDNs are accepted await _postgresService.ExecuteQueryAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, validServer, "testdb", "SELECT 1", TestContext.Current.CancellationToken); @@ -130,7 +130,7 @@ public async Task ExecuteQueryAsync_WithShortServerName_AppendsSuffix() { // Short names (no dot) get the suffix appended automatically await _postgresService.ExecuteQueryAsync( - "test-sub", "test-rg", AuthTypes.MicrosoftEntra, "test-user", null, + AuthTypes.MicrosoftEntra, "test-user", null, "myserver", "testdb", "SELECT 1", TestContext.Current.CancellationToken); @@ -139,3 +139,4 @@ await _postgresService.ExecuteQueryAsync( Assert.Equal("myserver.postgres.database.azure.com", parsed.Host); } } + diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs index 1cb5b49425..0053620a99 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft Corporation. +// Copyright (c) Microsoft Corporation. // Licensed under the MIT License. using System.Data.Common; @@ -85,7 +85,7 @@ public async Task ExecuteQueryAsync_InvalidCastException_Test() // Act CommandValidationException exception = await Assert.ThrowsAsync(async () => { - await _postgresService.ExecuteQueryAsync(subscriptionId, resourceGroup, authType, user, null, server, database, query, TestContext.Current.CancellationToken); + await _postgresService.ExecuteQueryAsync(authType, user, null, server, database, query, TestContext.Current.CancellationToken); }); // Assert @@ -109,7 +109,7 @@ public async Task ExecuteQueryAsync_MixedDataTypes_Test() [typeof(string), typeof(int), typeof(InvalidCastItem)])); // Act - List rows = await _postgresService.ExecuteQueryAsync(subscriptionId, resourceGroup, authType, user, null, server, database, query, TestContext.Current.CancellationToken); + List rows = await _postgresService.ExecuteQueryAsync(authType, user, null, server, database, query, TestContext.Current.CancellationToken); // Assert Assert.Equal(4, rows.Count); @@ -132,7 +132,7 @@ public async Task ExecuteQueryAsync_NoRows_Test() [typeof(string), typeof(int), typeof(InvalidCastItem)])); // Act - List rows = await _postgresService.ExecuteQueryAsync(subscriptionId, resourceGroup, authType, user, null, server, database, query, TestContext.Current.CancellationToken); + List rows = await _postgresService.ExecuteQueryAsync(authType, user, null, server, database, query, TestContext.Current.CancellationToken); // Assert Assert.Single(rows); @@ -202,10 +202,29 @@ public async Task ListServersAsync_ResourceGroupScope_ThrowsWhenRgNotFound() () => sut.ListServersAsync(subscriptionId, resourceGroup, TestContext.Current.CancellationToken)); Assert.Contains(resourceGroup, ex.Message); } + [Fact] + public async Task ListTablesAsync_UsesParameterizedSchemaQuery_AndDefaultsAreNotInterpolated() + { + // Verify the table listing query is parameterized on @schema (no string interpolation), + // orders results deterministically, and actually binds the schema parameter value. + string? capturedQuery = null; + var command = Substitute.For(); + _dbProvider.GetCommand(Arg.Do(q => capturedQuery = q), Arg.Any()) + .Returns(command); + + await _postgresService.ListTablesAsync(authType, user, null, server, database, "analytics", TestContext.Current.CancellationToken); + + Assert.NotNull(capturedQuery); + Assert.Contains("@schema", capturedQuery); + Assert.Contains("ORDER BY table_name", capturedQuery); + Assert.DoesNotContain("'public'", capturedQuery); + Assert.DoesNotContain("analytics", capturedQuery); + + Assert.Single(command.Parameters); + Assert.Equal("schema", command.Parameters[0].ParameterName); + Assert.Equal("analytics", command.Parameters[0].Value); + } } - - /// - /// Subclass that replaces the un-mockable ARM SDK extension-method calls with /// in-memory sequences, isolating logic. /// internal sealed class TestablePostgresService( @@ -237,3 +256,4 @@ protected override async IAsyncEnumerable ListResourceGroupServerNamesAs } } } + diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs index 66183bf73c..77de8d43a5 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs @@ -19,7 +19,7 @@ public class TableSchemaGetCommandTests : CommandUnitTestsBase(["CREATE TABLE test (id INT);"]); - Service.GetTableSchemaAsync("sub123", "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "table123", Arg.Any()).Returns(expectedSchema); + Service.GetTableSchemaAsync(AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "table123", Arg.Any()).Returns(expectedSchema); var response = await ExecuteCommandAsync( "--subscription", "sub123", @@ -37,7 +37,7 @@ public async Task ExecuteAsync_ReturnsSchema_WhenSchemaExists() [Fact] public async Task ExecuteAsync_ReturnsEmpty_WhenSchemaDoesNotExist() { - Service.GetTableSchemaAsync("sub123", "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "table123", Arg.Any()).Returns([]); + Service.GetTableSchemaAsync(AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "table123", Arg.Any()).Returns([]); var response = await ExecuteCommandAsync( "--subscription", "sub123", From feb18fa1eb99a4a83f216f64d49365273c78d0ee Mon Sep 17 00:00:00 2001 From: Victor Colin Amador Date: Mon, 15 Jun 2026 12:25:03 -0600 Subject: [PATCH 2/4] Remove subscription/resource-group from Postgres data-plane input schema Address maintainer review feedback on the Postgres data-plane commands. The query and table-schema-get commands connect directly via Npgsql and never use ARM scoping, but --subscription/--resource-group were still exposed in their MCP input schema. Re-parent BaseDatabaseCommand to GlobalCommand so those options are no longer registered, and add regression tests asserting the options are absent. Control-plane commands (list, server config/param) keep them. Relates to #469 and #471. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../vcolin7-postgres-schema-listing.yaml | 2 +- .../Azure.Mcp.Server/docs/azmcp-commands.md | 8 ++----- .../src/Commands/BaseDatabaseCommand.cs | 14 +++++++++-- .../Database/DatabaseQueryCommandTests.cs | 24 +++++++++---------- .../PostgresCommandTests.cs | 12 ---------- .../Table/TableSchemaGetCommandTests.cs | 20 +++++++++------- 6 files changed, 39 insertions(+), 41 deletions(-) diff --git a/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml b/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml index b82a60ddfc..1fa8dfbcc4 100644 --- a/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml +++ b/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml @@ -1,3 +1,3 @@ changes: - section: "Features Added" - description: "Added an optional `--schema` parameter to `azmcp postgres list` for listing tables in non-public PostgreSQL schemas (defaults to `public`), and removed unused `subscription`/`resource-group` parameters from the PostgreSQL data-plane service methods." + description: "Added an optional `--schema` parameter to `azmcp postgres list` for listing tables in non-public PostgreSQL schemas (defaults to `public`), and removed unused `--subscription`/`--resource-group` parameters from the PostgreSQL data-plane commands (`azmcp postgres database query`, `azmcp postgres table schema get`) and service methods." diff --git a/servers/Azure.Mcp.Server/docs/azmcp-commands.md b/servers/Azure.Mcp.Server/docs/azmcp-commands.md index c49056c1c4..ae53d88364 100644 --- a/servers/Azure.Mcp.Server/docs/azmcp-commands.md +++ b/servers/Azure.Mcp.Server/docs/azmcp-commands.md @@ -2275,18 +2275,14 @@ azmcp postgres list --subscription \ # Execute a query on a PostgreSQL database # ❌ Destructive | ✅ Idempotent | ❌ OpenWorld | ✅ ReadOnly | ❌ Secret | ❌ LocalRequired -azmcp postgres database query --subscription \ - --resource-group \ - --user \ +azmcp postgres database query --user \ --server \ --database \ --query # Get the schema of a specific table in a PostgreSQL database # ❌ Destructive | ✅ Idempotent | ❌ OpenWorld | ✅ ReadOnly | ❌ Secret | ❌ LocalRequired -azmcp postgres table schema get --subscription \ - --resource-group \ - --user \ +azmcp postgres table schema get --user \ --server \ --database \ --table diff --git a/tools/Azure.Mcp.Tools.Postgres/src/Commands/BaseDatabaseCommand.cs b/tools/Azure.Mcp.Tools.Postgres/src/Commands/BaseDatabaseCommand.cs index c09e32a6d3..600a676c24 100644 --- a/tools/Azure.Mcp.Tools.Postgres/src/Commands/BaseDatabaseCommand.cs +++ b/tools/Azure.Mcp.Tools.Postgres/src/Commands/BaseDatabaseCommand.cs @@ -9,13 +9,21 @@ namespace Azure.Mcp.Tools.Postgres.Commands; +// Data-plane commands connect directly to PostgreSQL via Npgsql and therefore do not +// require ARM-scoping options (--subscription / --resource-group). They derive from +// GlobalCommand rather than the subscription-based hierarchy so those options are not +// part of the MCP input schema. public abstract class BaseDatabaseCommand< - [DynamicallyAccessedMembers(TrimAnnotations.CommandAnnotations)] TOptions>(ILogger> logger) - : BaseServerCommand(logger) where TOptions : BasePostgresOptions, new() + [DynamicallyAccessedMembers(TrimAnnotations.CommandAnnotations)] TOptions>(ILogger> logger) + : GlobalCommand where TOptions : BasePostgresOptions, new() { + protected readonly ILogger> _logger = logger; + protected override void RegisterOptions(Command command) { base.RegisterOptions(command); + command.Options.Add(PostgresOptionDefinitions.Server); + command.Options.Add(PostgresOptionDefinitions.User); command.Options.Add(PostgresOptionDefinitions.Database); command.Options.Add(PostgresOptionDefinitions.AuthType); command.Options.Add(PostgresOptionDefinitions.Password); @@ -24,6 +32,8 @@ protected override void RegisterOptions(Command command) protected override TOptions BindOptions(ParseResult parseResult) { var options = base.BindOptions(parseResult); + options.Server = parseResult.GetValueOrDefault(PostgresOptionDefinitions.Server.Name); + options.User = parseResult.GetValueOrDefault(PostgresOptionDefinitions.User.Name); options.Database = parseResult.GetValueOrDefault(PostgresOptionDefinitions.Database.Name); options.AuthType = parseResult.GetValueOrDefault(PostgresOptionDefinitions.AuthType.Name); options.Password = parseResult.GetValueOrDefault(PostgresOptionDefinitions.Password.Name); diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs index 6a0772efc8..ec719ce450 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Database/DatabaseQueryCommandTests.cs @@ -26,8 +26,6 @@ public async Task ExecuteAsync_ReturnsQueryResults_WhenQueryIsValid() .Returns(expectedResults); var response = await ExecuteCommandAsync( - "--subscription", "sub123", - "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", @@ -45,8 +43,6 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenQueryFails() .Returns([]); var response = await ExecuteCommandAsync( - "--subscription", "sub123", - "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", @@ -58,8 +54,6 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenQueryFails() } [Theory] - [InlineData("--subscription")] - [InlineData("--resource-group")] [InlineData("--user")] [InlineData("--server")] [InlineData("--database")] @@ -67,8 +61,6 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenQueryFails() public async Task ExecuteAsync_ReturnsError_WhenParameterIsMissing(string missingParameter) { var response = await ExecuteCommandAsync(ArgBuilder.BuildArgs(missingParameter, - ("--subscription", "sub123"), - ("--resource-group", "rg1"), ($"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra), ("--user", "user1"), ("--server", "server123"), @@ -81,6 +73,18 @@ public async Task ExecuteAsync_ReturnsError_WhenParameterIsMissing(string missin Assert.Equal($"Missing Required options: {missingParameter}", response.Message); } + [Fact] + public void Command_DoesNotExposeArmScopingOptions() + { + var optionNames = CommandDefinition.Options.Select(o => o.Name.TrimStart('-')).ToList(); + + Assert.DoesNotContain("subscription", optionNames); + Assert.DoesNotContain("resource-group", optionNames); + Assert.Contains(PostgresOptionDefinitions.UserName, optionNames); + Assert.Contains(PostgresOptionDefinitions.ServerName, optionNames); + Assert.Contains(PostgresOptionDefinitions.DatabaseName, optionNames); + } + [Theory] [InlineData("DELETE FROM users;")] [InlineData("SELECT * FROM users; DROP TABLE users;")] @@ -126,8 +130,6 @@ public async Task ExecuteAsync_ReturnsError_WhenParameterIsMissing(string missin public async Task ExecuteAsync_InvalidQuery_ValidationError(string badQuery) { var response = await ExecuteCommandAsync( - "--subscription", "sub123", - "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", @@ -145,8 +147,6 @@ public async Task ExecuteAsync_LongQuery_ValidationError() { var longSelect = "SELECT " + new string('a', 6000) + " FROM test"; // exceeds max length var response = await ExecuteCommandAsync( - "--subscription", "sub123", - "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresCommandTests.cs index cec5e7ed32..80fe7b17ed 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/PostgresCommandTests.cs @@ -208,8 +208,6 @@ public async Task Should_GetTableSchema_Successfully() "postgres_table_schema_get", new() { - { "subscription", Settings.SubscriptionId }, - { "resource-group", Settings.ResourceGroupName }, { "server", ServerName }, { "database", TestDatabaseName }, { PostgresOptionDefinitions.AuthTypeText, AuthTypes.MicrosoftEntra }, @@ -246,8 +244,6 @@ public async Task Should_ExecuteQuery_Successfully() "postgres_database_query", new() { - { "subscription", Settings.SubscriptionId }, - { "resource-group", Settings.ResourceGroupName }, { "server", ServerName }, { "database", TestDatabaseName }, { PostgresOptionDefinitions.AuthTypeText, AuthTypes.MicrosoftEntra }, @@ -278,8 +274,6 @@ public async Task Should_ExecuteQuery_WithAggregation_Successfully() "postgres_database_query", new() { - { "subscription", Settings.SubscriptionId }, - { "resource-group", Settings.ResourceGroupName }, { "server", ServerName }, { "database", TestDatabaseName }, { PostgresOptionDefinitions.AuthTypeText, AuthTypes.MicrosoftEntra }, @@ -305,8 +299,6 @@ public async Task Should_ExecuteQuery_WithJoin_Successfully() "postgres_database_query", new() { - { "subscription", Settings.SubscriptionId }, - { "resource-group", Settings.ResourceGroupName }, { "server", ServerName }, { "database", TestDatabaseName }, { PostgresOptionDefinitions.AuthTypeText, AuthTypes.MicrosoftEntra }, @@ -423,8 +415,6 @@ public async Task Should_RejectNonSelectQuery_WithValidationError() JsonElement error = await this.CallToolAsyncWithErrorExpected("postgres_database_query", new() { - { "subscription", Settings.SubscriptionId }, - { "resource-group", Settings.ResourceGroupName }, { "server", ServerName }, { "database", TestDatabaseName }, { PostgresOptionDefinitions.AuthTypeText, AuthTypes.MicrosoftEntra }, @@ -502,8 +492,6 @@ public async Task Should_HandleInvalidTableName_Gracefully() "postgres_table_schema_get", new() { - { "subscription", Settings.SubscriptionId }, - { "resource-group", Settings.ResourceGroupName }, { "server", ServerName }, { "database", TestDatabaseName }, { PostgresOptionDefinitions.AuthTypeText, AuthTypes.MicrosoftEntra }, diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs index 77de8d43a5..6031246a1d 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Table/TableSchemaGetCommandTests.cs @@ -22,8 +22,6 @@ public async Task ExecuteAsync_ReturnsSchema_WhenSchemaExists() Service.GetTableSchemaAsync(AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "table123", Arg.Any()).Returns(expectedSchema); var response = await ExecuteCommandAsync( - "--subscription", "sub123", - "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", @@ -40,8 +38,6 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenSchemaDoesNotExist() Service.GetTableSchemaAsync(AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "table123", Arg.Any()).Returns([]); var response = await ExecuteCommandAsync( - "--subscription", "sub123", - "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", @@ -53,8 +49,6 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenSchemaDoesNotExist() } [Theory] - [InlineData("--subscription")] - [InlineData("--resource-group")] [InlineData("--user")] [InlineData("--server")] [InlineData("--database")] @@ -62,8 +56,6 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenSchemaDoesNotExist() public async Task ExecuteAsync_ReturnsError_WhenParameterIsMissing(string missingParameter) { var response = await ExecuteCommandAsync(ArgBuilder.BuildArgs(missingParameter, - ("--subscription", "sub123"), - ("--resource-group", "rg1"), ($"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra), ("--user", "user1"), ("--server", "server123"), @@ -75,4 +67,16 @@ public async Task ExecuteAsync_ReturnsError_WhenParameterIsMissing(string missin Assert.Equal(HttpStatusCode.BadRequest, response.Status); Assert.Equal($"Missing Required options: {missingParameter}", response.Message); } + + [Fact] + public void Command_DoesNotExposeArmScopingOptions() + { + var optionNames = CommandDefinition.Options.Select(o => o.Name.TrimStart('-')).ToList(); + + Assert.DoesNotContain("subscription", optionNames); + Assert.DoesNotContain("resource-group", optionNames); + Assert.Contains(PostgresOptionDefinitions.UserName, optionNames); + Assert.Contains(PostgresOptionDefinitions.ServerName, optionNames); + Assert.Contains(PostgresOptionDefinitions.DatabaseName, optionNames); + } } From eea555b6bb7dc36c9069848340819100e3c49a64 Mon Sep 17 00:00:00 2001 From: Victor Colin Amador Date: Mon, 15 Jun 2026 19:40:06 -0600 Subject: [PATCH 3/4] Fixed formatting issues --- .../Services/PostgresServiceTests.cs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs index 0053620a99..6a31270133 100644 --- a/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs +++ b/tools/Azure.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.Tests/Services/PostgresServiceTests.cs @@ -202,6 +202,7 @@ public async Task ListServersAsync_ResourceGroupScope_ThrowsWhenRgNotFound() () => sut.ListServersAsync(subscriptionId, resourceGroup, TestContext.Current.CancellationToken)); Assert.Contains(resourceGroup, ex.Message); } + [Fact] public async Task ListTablesAsync_UsesParameterizedSchemaQuery_AndDefaultsAreNotInterpolated() { @@ -225,6 +226,9 @@ public async Task ListTablesAsync_UsesParameterizedSchemaQuery_AndDefaultsAreNot Assert.Equal("analytics", command.Parameters[0].Value); } } + + /// + /// Subclass that replaces the un-mockable ARM SDK extension-method calls with /// in-memory sequences, isolating logic. /// internal sealed class TestablePostgresService( From 7d20a19c5a4097eda64948c19770fd99a22b8b5e Mon Sep 17 00:00:00 2001 From: Victor Colin Amador Date: Mon, 15 Jun 2026 21:31:56 -0600 Subject: [PATCH 4/4] Updated changelog --- .../changelog-entries/vcolin7-postgres-schema-listing.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml b/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml index 1fa8dfbcc4..ed30b4f066 100644 --- a/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml +++ b/servers/Azure.Mcp.Server/changelog-entries/vcolin7-postgres-schema-listing.yaml @@ -1,3 +1,3 @@ changes: - section: "Features Added" - description: "Added an optional `--schema` parameter to `azmcp postgres list` for listing tables in non-public PostgreSQL schemas (defaults to `public`), and removed unused `--subscription`/`--resource-group` parameters from the PostgreSQL data-plane commands (`azmcp postgres database query`, `azmcp postgres table schema get`) and service methods." + description: "Added an optional `--schema` parameter to `azmcp postgres list` for listing tables in non-public PostgreSQL schemas (defaults to `public`), and removed unused `--subscription`/`--resource-group` parameters from the PostgreSQL data-plane commands: `database query`, `table schema get`."