Skip to content

DI Refactor: Azure.Mcp.Tools.Postgres #2079

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Copilot wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

DI Refactor: Azure.Mcp.Tools.Postgres #2079

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f39253c
Initial plan
Copilot Mar 17, 2026
85c67a6
Refactor Azure.Mcp.Tools.Postgres to use constructor DI for IPostgres…
Copilot Mar 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions tools/Azure.Mcp.Tools.Postgres/src/Commands/Database/DatabaseQueryCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,9 @@

namespace Azure.Mcp.Tools.Postgres.Commands.Database;

public sealed class DatabaseQueryCommand(ILogger<DatabaseQueryCommand> logger) : BaseDatabaseCommand<DatabaseQueryOptions>(logger)
public sealed class DatabaseQueryCommand(IPostgresService postgresService, ILogger<DatabaseQueryCommand> logger) : BaseDatabaseCommand<DatabaseQueryOptions>(logger)
{
private readonly IPostgresService _postgresService = postgresService;
private const string CommandTitle = "Query PostgreSQL Database";

public override string Id => "81a28bca-014c-4738-9e1a-654d77cb2dd8";
Expand Down Expand Up @@ -58,10 +59,9 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,

try
{
IPostgresService pgService = context.GetService<IPostgresService>() ?? throw new InvalidOperationException("PostgreSQL service is not available.");
// Validate the query early to avoid sending unsafe SQL to the server.
SqlQueryValidator.EnsureReadOnlySelect(options.Query);
List<string> queryResult = await pgService.ExecuteQueryAsync(
List<string> queryResult = await _postgresService.ExecuteQueryAsync(
options.Subscription!,
options.ResourceGroup!,
options.AuthType!,
Expand Down
11 changes: 5 additions & 6 deletions tools/Azure.Mcp.Tools.Postgres/src/Commands/PostgresListCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,9 @@

namespace Azure.Mcp.Tools.Postgres.Commands;

public sealed class PostgresListCommand(ILogger<PostgresListCommand> logger) : BasePostgresCommand<BasePostgresOptions>(logger)
public sealed class PostgresListCommand(IPostgresService postgresService, ILogger<PostgresListCommand> logger) : BasePostgresCommand<BasePostgresOptions>(logger)
{
private readonly IPostgresService _postgresService = postgresService;
public override string Id => "8a12c3f4-2e5d-4b3a-9f2c-5e6d7f8a9b0c";

public override string Name => "list";
Expand Down Expand Up @@ -73,13 +74,11 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,

var options = BindOptions(parseResult);

IPostgresService postgresService = context.GetService<IPostgresService>() ?? throw new InvalidOperationException("PostgreSQL service is not available.");

// Route based on provided parameters
if (!string.IsNullOrEmpty(options.Database))
{
// List tables in specified database
List<string> tables = await postgresService.ListTablesAsync(
List<string> tables = await _postgresService.ListTablesAsync(
options.Subscription!,
options.ResourceGroup!,
options.AuthType!,
Expand All @@ -96,7 +95,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
else if (!string.IsNullOrEmpty(options.Server))
{
// List databases on specified server
List<string> databases = await postgresService.ListDatabasesAsync(
List<string> databases = await _postgresService.ListDatabasesAsync(
options.Subscription!,
options.ResourceGroup!,
options.AuthType!,
Expand All @@ -112,7 +111,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
else
{
// List all servers in the subscription (optionally scoped to a resource group)
List<string> servers = await postgresService.ListServersAsync(
List<string> servers = await _postgresService.ListServersAsync(
options.Subscription!,
options.ResourceGroup,
cancellationToken);
Expand Down
6 changes: 3 additions & 3 deletions tools/Azure.Mcp.Tools.Postgres/src/Commands/Server/ServerConfigGetCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,9 @@

namespace Azure.Mcp.Tools.Postgres.Commands.Server;

public sealed class ServerConfigGetCommand(ILogger<ServerConfigGetCommand> logger) : BaseServerCommand<ServerConfigGetOptions>(logger)
public sealed class ServerConfigGetCommand(IPostgresService postgresService, ILogger<ServerConfigGetCommand> logger) : BaseServerCommand<ServerConfigGetOptions>(logger)
{
private readonly IPostgresService _postgresService = postgresService;
private const string CommandTitle = "Get PostgreSQL Server Configuration";

public override string Id => "049a0d10-0a6e-4278-a0a3-15ce6b2e5ee1";
Expand Down Expand Up @@ -44,8 +45,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
try
{

IPostgresService pgService = context.GetService<IPostgresService>() ?? throw new InvalidOperationException("PostgreSQL service is not available.");
var config = await pgService.GetServerConfigAsync(options.Subscription!, options.ResourceGroup!, options.User!, options.Server!, cancellationToken);
var config = await _postgresService.GetServerConfigAsync(options.Subscription!, options.ResourceGroup!, options.User!, options.Server!, cancellationToken);
context.Response.Results = config?.Length > 0 ?
ResponseResult.Create(new(config), PostgresJsonContext.Default.ServerConfigGetCommandResult) :
null;
Expand Down
6 changes: 3 additions & 3 deletions tools/Azure.Mcp.Tools.Postgres/src/Commands/Server/ServerParamGetCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@

namespace Azure.Mcp.Tools.Postgres.Commands.Server;

public sealed class ServerParamGetCommand(ILogger<ServerParamGetCommand> logger) : BaseServerCommand<ServerParamGetOptions>(logger)
public sealed class ServerParamGetCommand(IPostgresService postgresService, ILogger<ServerParamGetCommand> logger) : BaseServerCommand<ServerParamGetOptions>(logger)
{
private readonly IPostgresService _postgresService = postgresService;
private const string CommandTitle = "Get PostgreSQL Server Parameter";

public override string Id => "af3a581d-ab64-4939-9765-974815d9c7be";
Expand Down Expand Up @@ -58,8 +59,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,

try
{
IPostgresService pgService = context.GetService<IPostgresService>() ?? throw new InvalidOperationException("PostgreSQL service is not available.");
var parameterValue = await pgService.GetServerParameterAsync(options.Subscription!, options.ResourceGroup!, options.User!, options.Server!, options.Param!, cancellationToken);
var parameterValue = await _postgresService.GetServerParameterAsync(options.Subscription!, options.ResourceGroup!, options.User!, options.Server!, options.Param!, cancellationToken);
context.Response.Results = parameterValue?.Length > 0 ?
ResponseResult.Create(new(parameterValue), PostgresJsonContext.Default.ServerParamGetCommandResult) :
null;
Expand Down
6 changes: 3 additions & 3 deletions tools/Azure.Mcp.Tools.Postgres/src/Commands/Server/ServerParamSetCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@

namespace Azure.Mcp.Tools.Postgres.Commands.Server;

public sealed class ServerParamSetCommand(ILogger<ServerParamSetCommand> logger) : BaseServerCommand<ServerParamSetOptions>(logger)
public sealed class ServerParamSetCommand(IPostgresService postgresService, ILogger<ServerParamSetCommand> logger) : BaseServerCommand<ServerParamSetOptions>(logger)
{
private readonly IPostgresService _postgresService = postgresService;
private const string CommandTitle = "Set PostgreSQL Server Parameter";

public override string Id => "2134621b-518f-48ac-a66a-82c40fcb58bb";
Expand Down Expand Up @@ -60,8 +61,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,

try
{
IPostgresService pgService = context.GetService<IPostgresService>() ?? throw new InvalidOperationException("PostgreSQL service is not available.");
var result = await pgService.SetServerParameterAsync(options.Subscription!, options.ResourceGroup!, options.User!, options.Server!, options.Param!, options.Value!, cancellationToken);
var result = await _postgresService.SetServerParameterAsync(options.Subscription!, options.ResourceGroup!, options.User!, options.Server!, options.Param!, options.Value!, cancellationToken);
context.Response.Results = !string.IsNullOrEmpty(result) ?
ResponseResult.Create(new(result, options.Param!, options.Value!), PostgresJsonContext.Default.ServerParamSetCommandResult) :
null;
Expand Down
6 changes: 3 additions & 3 deletions tools/Azure.Mcp.Tools.Postgres/src/Commands/Table/TableSchemaGetCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@

namespace Azure.Mcp.Tools.Postgres.Commands.Table;

public sealed class TableSchemaGetCommand(ILogger<TableSchemaGetCommand> logger) : BaseDatabaseCommand<TableSchemaGetOptions>(logger)
public sealed class TableSchemaGetCommand(IPostgresService postgresService, ILogger<TableSchemaGetCommand> logger) : BaseDatabaseCommand<TableSchemaGetOptions>(logger)
{
private readonly IPostgresService _postgresService = postgresService;
private const string CommandTitle = "Get PostgreSQL Table Schema";

public override string Id => "643a3497-44e1-4727-b3d6-c2e5dba6cab2";
Expand Down Expand Up @@ -56,8 +57,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
{


IPostgresService pgService = context.GetService<IPostgresService>() ?? throw new InvalidOperationException("PostgreSQL service is not available.");
List<string> schema = await pgService.GetTableSchemaAsync(
List<string> schema = await _postgresService.GetTableSchemaAsync(
options.Subscription!,
options.ResourceGroup!,
options.AuthType!,
Expand Down
10 changes: 5 additions & 5 deletions ...s.Postgres/tests/Azure.Mcp.Tools.Postgres.UnitTests/Database/DatabaseQueryCommandTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ public async Task ExecuteAsync_ReturnsQueryResults_WhenQueryIsValid()
_postgresService.ExecuteQueryAsync("sub123", "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "SELECT * FROM test;", Arg.Any<CancellationToken>())
.Returns(expectedResults);

var command = new DatabaseQueryCommand(_logger);
var command = new DatabaseQueryCommand(_postgresService, _logger);
var args = command.GetCommand().Parse(["--subscription", "sub123", "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", "--database", "db123", "--query", "SELECT * FROM test;"]);
var context = new CommandContext(_serviceProvider);
var response = await command.ExecuteAsync(context, args, TestContext.Current.CancellationToken);
Expand All @@ -66,7 +66,7 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenQueryFails()
_postgresService.ExecuteQueryAsync("sub123", "rg1", AuthTypes.MicrosoftEntra, "user1", null, "server1", "db123", "SELECT * FROM test;", Arg.Any<CancellationToken>())
.Returns([]);

var command = new DatabaseQueryCommand(_logger);
var command = new DatabaseQueryCommand(_postgresService, _logger);

var args = command.GetCommand().Parse(["--subscription", "sub123", "--resource-group", "rg1", $"--{PostgresOptionDefinitions.AuthTypeText}", AuthTypes.MicrosoftEntra, "--user", "user1", "--server", "server1", "--database", "db123", "--query", "SELECT * FROM test;"]);
var context = new CommandContext(_serviceProvider);
Expand All @@ -91,7 +91,7 @@ public async Task ExecuteAsync_ReturnsEmpty_WhenQueryFails()
[InlineData("--query")]
public async Task ExecuteAsync_ReturnsError_WhenParameterIsMissing(string missingParameter)
{
var command = new DatabaseQueryCommand(_logger);
var command = new DatabaseQueryCommand(_postgresService, _logger);
var args = command.GetCommand().Parse(ArgBuilder.BuildArgs(missingParameter,
("--subscription", "sub123"),
("--resource-group", "rg1"),
Expand Down Expand Up @@ -154,7 +154,7 @@ public async Task ExecuteAsync_ReturnsError_WhenParameterIsMissing(string missin
[InlineData("SELECT * FROM pg_user_mappings")] // FDW credential exposure
public async Task ExecuteAsync_InvalidQuery_ValidationError(string badQuery)
{
var command = new DatabaseQueryCommand(_logger);
var command = new DatabaseQueryCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand All @@ -178,7 +178,7 @@ public async Task ExecuteAsync_InvalidQuery_ValidationError(string badQuery)
public async Task ExecuteAsync_LongQuery_ValidationError()
{
var longSelect = "SELECT " + new string('a', 6000) + " FROM test"; // exceeds max length
var command = new DatabaseQueryCommand(_logger);
var command = new DatabaseQueryCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand Down
30 changes: 15 additions & 15 deletions ...e.Mcp.Tools.Postgres/tests/Azure.Mcp.Tools.Postgres.UnitTests/PostgresListCommandTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ public async Task ExecuteAsync_ListsServers_WhenNoServerOrDatabaseProvided()
var expectedServers = new List<string> { "postgres-server-1", "postgres-server-2", "postgres-server-3" };
_postgresService.ListServersAsync("sub123", "rg1", Arg.Any<CancellationToken>()).Returns(expectedServers);

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1"
Expand Down Expand Up @@ -67,7 +67,7 @@ public async Task ExecuteAsync_ListsAllServersInSubscription_WhenNoResourceGroup
var expectedServers = new List<string> { "postgres-server-1", "postgres-server-2" };
_postgresService.ListServersAsync("sub123", null, Arg.Any<CancellationToken>()).Returns(expectedServers);

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123"
]);
Expand All @@ -91,7 +91,7 @@ public async Task ExecuteAsync_ListsAllServersInSubscription_WhenNoResourceGroup
[Fact]
public async Task ExecuteAsync_ReturnsError_WhenServerProvidedWithoutUser()
{
var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--server", "server1"
Expand All @@ -118,7 +118,7 @@ public async Task ExecuteAsync_ListsDatabases_WhenServerProvided()
"server1",
Arg.Any<CancellationToken>()).Returns(expectedDatabases);

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand Down Expand Up @@ -157,7 +157,7 @@ public async Task ExecuteAsync_ListsTables_WhenServerAndDatabaseProvided()
"db1",
Arg.Any<CancellationToken>()).Returns(expectedTables);

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand Down Expand Up @@ -188,7 +188,7 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoServersExist()
{
_postgresService.ListServersAsync("sub123", "rg1", Arg.Any<CancellationToken>()).Returns([]);

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1"
Expand Down Expand Up @@ -222,7 +222,7 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoDatabasesExist()
"server1",
Arg.Any<CancellationToken>()).Returns([]);

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand Down Expand Up @@ -260,7 +260,7 @@ public async Task ExecuteAsync_ReturnsNull_WhenNoTablesExist()
"db1",
Arg.Any<CancellationToken>()).Returns([]);

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand Down Expand Up @@ -293,7 +293,7 @@ public async Task ExecuteAsync_ReturnsError_WhenListServersThrows()
_postgresService.ListServersAsync("sub123", "rg1", Arg.Any<CancellationToken>())
.ThrowsAsync(new Exception("Test error"));

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1"
Expand Down Expand Up @@ -321,7 +321,7 @@ public async Task ExecuteAsync_ReturnsError_WhenListDatabasesThrows()
Arg.Any<CancellationToken>())
.ThrowsAsync(new Exception("Test error"));

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand Down Expand Up @@ -353,7 +353,7 @@ public async Task ExecuteAsync_ReturnsError_WhenListTablesThrows()
Arg.Any<CancellationToken>())
.ThrowsAsync(new Exception("Test error"));

var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse([
"--subscription", "sub123",
"--resource-group", "rg1",
Expand All @@ -375,7 +375,7 @@ public async Task ExecuteAsync_ReturnsError_WhenListTablesThrows()
[InlineData("--subscription")]
public async Task ExecuteAsync_ReturnsError_WhenRequiredParameterIsMissing(string missingParameter)
{
var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
var args = command.GetCommand().Parse(ArgBuilder.BuildArgs(missingParameter,
("--subscription", "sub123")
));
Expand All @@ -391,7 +391,7 @@ public async Task ExecuteAsync_ReturnsError_WhenRequiredParameterIsMissing(strin
[Fact]
public void Metadata_IsConfiguredCorrectly()
{
var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);

Assert.False(command.Metadata.Destructive);
Assert.True(command.Metadata.ReadOnly);
Expand All @@ -400,14 +400,14 @@ public void Metadata_IsConfiguredCorrectly()
[Fact]
public void Name_IsCorrect()
{
var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
Assert.Equal("list", command.Name);
}

[Fact]
public void Description_IsCorrect()
{
var command = new PostgresListCommand(_logger);
var command = new PostgresListCommand(_postgresService, _logger);
Assert.Contains("List PostgreSQL servers", command.Description);
Assert.Contains("databases, or tables", command.Description);
}
Expand Down
Loading
Loading