Add support for using a managed identity to connect to Azure Database for PostgreSQL

src/GenerateBatchVmSkus/GenerateBatchVmSkus.csproj

@@ -8,7 +8,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Azure.Identity" Version="1.10.1" />
+    <PackageReference Include="Azure.Identity" Version="1.10.2" />
     <PackageReference Include="Azure.ResourceManager.Batch" Version="1.1.1" />
     <PackageReference Include="Azure.ResourceManager.Compute" Version="1.1.0" />
     <PackageReference Include="Azure.Security.KeyVault.Secrets" Version="4.4.0" />

src/Tes.ApiClients/Tes.ApiClients.csproj

@@ -8,7 +8,7 @@
 
   <ItemGroup>
     <PackageReference Include="Azure.Core" Version="1.35.0" />
-    <PackageReference Include="Azure.Identity" Version="1.10.1" />
+    <PackageReference Include="Azure.Identity" Version="1.10.2" />
     <PackageReference Include="Microsoft.Extensions.Caching.Abstractions" Version="7.0.0" />
     <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="7.0.0" />
     <PackageReference Include="Microsoft.Extensions.Logging" Version="7.0.0" />

src/Tes.Runner.Test/Tes.Runner.Test.csproj

@@ -11,7 +11,7 @@
 
   <ItemGroup>
     <PackageReference Include="Azure.Storage.Blobs" Version="12.16.0" />
-    <PackageReference Include="Microsoft.EntityFrameworkCore" Version="7.0.5" />
+    <PackageReference Include="Microsoft.EntityFrameworkCore" Version="7.0.12" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.5.0" />
     <PackageReference Include="Moq" Version="4.18.4" />
     <PackageReference Include="MSTest.TestAdapter" Version="3.0.2" />

src/Tes.Runner/Tes.Runner.csproj

@@ -7,7 +7,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Azure.Identity" Version="1.10.1" />
+    <PackageReference Include="Azure.Identity" Version="1.10.2" />
     <PackageReference Include="Azure.Storage.Blobs" Version="12.16.0" />
     <PackageReference Include="Docker.DotNet" Version="3.125.14" />
     <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="7.0.0" />

src/Tes/Models/PostgreSqlOptions.cs

@@ -27,5 +27,6 @@ public static string GetConfigurationSectionName(string serviceName = "Tes")
         public string DatabaseName { get; set; } = "tes_db";
         public string DatabaseUserLogin { get; set; }
         public string DatabaseUserPassword { get; set; }
+        public bool UseManagedIdentity { get; set; }
     }
 }

src/Tes/Models/TesTaskPostgres.cs

@@ -9,7 +9,7 @@ namespace Tes.Models
     /// <summary>
     /// Database schema for encapsulating a TesTask as Json for Postgresql.
     /// </summary>
-    [Table(Repository.TesDbContext.TesTasksPostgresTableName)]
+    [Table("testasks")]
     public class TesTaskDatabaseItem
     {
         [Column("id")]

src/Tes/Repository/PostgreSqlCachingRepository.cs

@@ -9,13 +9,15 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.EntityFrameworkCore;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Polly;
 
 namespace Tes.Repository
 {
     public abstract class PostgreSqlCachingRepository<T> : IDisposable where T : class
     {
+        private readonly IServiceScopeFactory _scopeFactory = null!;
         private readonly TimeSpan _writerWaitTime = TimeSpan.FromMilliseconds(50);
         private readonly int _batchSize = 1000;
         private static readonly TimeSpan defaultCompletedTaskCacheExpiration = TimeSpan.FromDays(1);
@@ -30,17 +32,16 @@ public abstract class PostgreSqlCachingRepository<T> : IDisposable where T : cla
         private readonly Task _writerWorkerTask;
 
         protected enum WriteAction { Add, Update, Delete }
-
-        protected Func<TesDbContext> CreateDbContext { get; init; }
         protected readonly ICache<T> _cache;
         protected readonly ILogger _logger;
 
         private bool _disposedValue;
 
-        protected PostgreSqlCachingRepository(ILogger logger = default, ICache<T> cache = default)
+        protected PostgreSqlCachingRepository(ILogger logger = default, ICache<T> cache = default, IServiceScopeFactory scopeFactory = default)
         {
             _logger = logger;
             _cache = cache;
+            _scopeFactory = scopeFactory;
 
             // The only "normal" exit for _writerWorkerTask is "cancelled". Anything else should force the process to exit because it means that this repository will no longer write to the database!
             _writerWorkerTask = Task.Run(() => WriterWorkerAsync(_writerWorkerCancellationTokenSource.Token))
@@ -187,7 +188,8 @@ private async ValueTask WriteItemsAsync(IList<(T DbItem, WriteAction Action, Tas
             if (dbItems.Count == 0) { return; }
 
             cancellationToken.ThrowIfCancellationRequested();
-            using var dbContext = CreateDbContext();
+            using var scope = _scopeFactory.CreateScope();
+            using var dbContext = scope.ServiceProvider.GetRequiredService<TesDbContext>();
 
             // Manually set entity state to avoid potential NPG PostgreSql bug
             dbContext.ChangeTracker.AutoDetectChangesEnabled = false;

src/Tes/Repository/TesDbContext.cs

@@ -1,40 +1,40 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System;
+using Azure.Core;
+using Azure.Identity;
 using Microsoft.EntityFrameworkCore;
 using Tes.Models;
+using Tes.Utilities;
 
 namespace Tes.Repository
 {
     public class TesDbContext : DbContext
     {
-        public const string TesTasksPostgresTableName = "testasks";
+        private const int maxBatchSize = 1000;
+        private readonly PostgresConnectionStringUtility connectionStringUtility = null!;
 
         public TesDbContext()
         {
             // Default constructor, which is required to run the EF migrations tool,
             // "dotnet ef migrations add InitialCreate"
+            // DI will NOT use this constructor
         }
 
-        public TesDbContext(string connectionString)
+        public TesDbContext(PostgresConnectionStringUtility connectionStringUtility)
         {
-            ArgumentException.ThrowIfNullOrEmpty(connectionString, nameof(connectionString));
-            ConnectionString = connectionString;
+            this.connectionStringUtility = connectionStringUtility;
         }
 
-        public string ConnectionString { get; set; }
         public DbSet<TesTaskDatabaseItem> TesTasks { get; set; }
 
         protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
         {
-            if (!optionsBuilder.IsConfigured)
-            {
-                // use PostgreSQL
-                optionsBuilder
-                    .UseNpgsql(ConnectionString, options => options.MaxBatchSize(1000))
-                    .UseLowerCaseNamingConvention();
-            }
+            string connectionString = this.connectionStringUtility.GetConnectionString().Result;
+
+            optionsBuilder
+                .UseNpgsql(connectionString, options => options.MaxBatchSize(maxBatchSize))
+                .UseLowerCaseNamingConvention();
         }
     }
 }

src/Tes/Repository/TesTaskPostgreSqlRepository.cs

@@ -11,8 +11,8 @@ namespace Tes.Repository
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.EntityFrameworkCore;
+    using Microsoft.Extensions.DependencyInjection;
     using Microsoft.Extensions.Logging;
-    using Microsoft.Extensions.Options;
     using Polly;
     using Tes.Models;
     using Tes.Utilities;
@@ -23,34 +23,24 @@ namespace Tes.Repository
     /// <typeparam name="TesTask"></typeparam>
     public sealed class TesTaskPostgreSqlRepository : PostgreSqlCachingRepository<TesTaskDatabaseItem>, IRepository<TesTask>
     {
+        private readonly IServiceScopeFactory _scopeFactory = null!;
+
         /// <summary>
         /// Default constructor that also will create the schema if it does not exist
         /// </summary>
         /// <param name="options"></param>
         /// <param name="logger"></param>
         /// <param name="cache"></param>
-        public TesTaskPostgreSqlRepository(IOptions<PostgreSqlOptions> options, ILogger<TesTaskPostgreSqlRepository> logger, ICache<TesTaskDatabaseItem> cache = null)
+        public TesTaskPostgreSqlRepository(ILogger<TesTaskPostgreSqlRepository> logger = default, IServiceScopeFactory scopeFactory = default, ICache<TesTaskDatabaseItem> cache = null)
             : base(logger, cache)
         {
-            var connectionString = new ConnectionStringUtility().GetPostgresConnectionString(options);
-            CreateDbContext = () => { return new TesDbContext(connectionString); };
-            using var dbContext = CreateDbContext();
+            _scopeFactory = scopeFactory;
+            using var scope = _scopeFactory.CreateScope();
+            using var dbContext = scope.ServiceProvider.GetRequiredService<TesDbContext>();
             dbContext.Database.MigrateAsync().Wait();
             WarmCacheAsync(CancellationToken.None).Wait();
         }
 
-        /// <summary>
-        /// Constructor for testing to enable mocking DbContext 
-        /// </summary>
-        /// <param name="createDbContext">A delegate that creates a TesTaskPostgreSqlRepository context</param>
-        public TesTaskPostgreSqlRepository(Func<TesDbContext> createDbContext)
-            : base()
-        {
-            CreateDbContext = createDbContext;
-            using var dbContext = createDbContext();
-            dbContext.Database.MigrateAsync().Wait();
-        }
-
         private async Task WarmCacheAsync(CancellationToken cancellationToken)
         {
             if (_cache is null)
@@ -224,7 +214,8 @@ private async Task<TesTaskDatabaseItem> GetItemFromCacheOrDatabase(string id, bo
 
             if (!_cache?.TryGetValue(id, out item) ?? true)
             {
-                using var dbContext = CreateDbContext();
+                using var scope = _scopeFactory.CreateScope();
+                using var dbContext = scope.ServiceProvider.GetRequiredService<TesDbContext>();
 
                 // Search for Id within the JSON
                 item = await _asyncPolicy.ExecuteAsync(ct => dbContext.TesTasks.FirstOrDefaultAsync(t => t.Json.Id == id, ct), cancellationToken);
@@ -252,7 +243,8 @@ private async Task<IEnumerable<TesTask>> InternalGetItemsAsync(Expression<Func<T
             //orderBy = pagination is null ? orderBy : q => q.OrderBy(t => t.Json.CreationTime).ThenBy(t => t.Json.Id);
             orderBy = pagination is null ? orderBy : q => q.OrderBy(t => t.Json.Id);
 
-            using var dbContext = CreateDbContext();
+            using var scope = _scopeFactory.CreateScope();
+            using var dbContext = scope.ServiceProvider.GetRequiredService<TesDbContext>();
             return (await GetItemsAsync(dbContext.TesTasks, WhereTesTask(predicate), cancellationToken, orderBy, pagination)).Select(item => EnsureActiveItemInCache(item, t => t.Json.Id, t => t.Json.IsActiveState()).Json);
         }
 

src/Tes/Tes.csproj

@@ -5,18 +5,20 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="Azure.Identity" Version="1.10.2" />
     <PackageReference Include="EFCore.NamingConventions" Version="7.0.2" />
-    <PackageReference Include="Microsoft.EntityFrameworkCore" Version="7.0.3" />
-    <PackageReference Include="Microsoft.EntityFrameworkCore.Design" Version="7.0.3">
+    <PackageReference Include="Microsoft.EntityFrameworkCore" Version="7.0.12" />
+    <PackageReference Include="Microsoft.EntityFrameworkCore.Design" Version="7.0.12">
       <PrivateAssets>all</PrivateAssets>
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
     </PackageReference>
-    <PackageReference Include="Microsoft.EntityFrameworkCore.InMemory" Version="7.0.3" />
+    <PackageReference Include="Microsoft.EntityFrameworkCore.InMemory" Version="7.0.12" />
+    <PackageReference Include="Microsoft.EntityFrameworkCore.Relational" Version="7.0.12" />
     <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="7.0.0" />
     <PackageReference Include="Microsoft.Rest.ClientRuntime" Version="2.3.24" />
     <!--Mitigate reported security issues-->
     <PackageReference Include="Newtonsoft.Json" Version="13.0.2" />
-    <PackageReference Include="Npgsql.EntityFrameworkCore.PostgreSQL" Version="7.0.3" />
+    <PackageReference Include="Npgsql.EntityFrameworkCore.PostgreSQL" Version="7.0.11" />
     <PackageReference Include="Polly" Version="7.2.3" />
     <PackageReference Include="Polly.Extensions.Http" Version="3.0.0" />
   </ItemGroup>

src/Tes/Utilities/PostgresConnectionStringUtility.cs

@@ -3,35 +3,81 @@
 
 using System;
 using System.Text;
-using Microsoft.Extensions.Options;
+using System.Threading.Tasks;
+using Azure.Core;
 using Tes.Models;
 
 namespace Tes.Utilities
 {
-    public class ConnectionStringUtility
+    public class PostgresConnectionStringUtility
     {
-        public string GetPostgresConnectionString(IOptions<PostgreSqlOptions> options)
+        private const string azureDatabaseForPostgresqlScope = "https://ossrdbms-aad.database.windows.net/.default";
+        private readonly string connectionString = null!;
+        private readonly TokenCredential tokenCredential = null!;
+        public bool UseManagedIdentity { get; set; }     
+
+        public PostgresConnectionStringUtility(PostgreSqlOptions options, TokenCredential tokenCredential)
+        {
+            this.tokenCredential = tokenCredential;
+            connectionString = InternalGetConnectionString(options);
+            UseManagedIdentity = options.UseManagedIdentity;
+        }
+
+        public async Task<string> GetConnectionString()
         {
-            ArgumentException.ThrowIfNullOrEmpty(options.Value.ServerName, nameof(options.Value.ServerName));
-            ArgumentException.ThrowIfNullOrEmpty(options.Value.ServerNameSuffix, nameof(options.Value.ServerNameSuffix));
-            ArgumentException.ThrowIfNullOrEmpty(options.Value.ServerPort, nameof(options.Value.ServerPort));
-            ArgumentException.ThrowIfNullOrEmpty(options.Value.ServerSslMode, nameof(options.Value.ServerSslMode));
-            ArgumentException.ThrowIfNullOrEmpty(options.Value.DatabaseName, nameof(options.Value.DatabaseName));
-            ArgumentException.ThrowIfNullOrEmpty(options.Value.DatabaseUserLogin, nameof(options.Value.DatabaseUserLogin));
-            ArgumentException.ThrowIfNullOrEmpty(options.Value.DatabaseUserPassword, nameof(options.Value.DatabaseUserPassword));
-
-            if (options.Value.ServerName.Contains(options.Value.ServerNameSuffix, StringComparison.OrdinalIgnoreCase))
+            if (UseManagedIdentity)
             {
-                throw new ArgumentException($"'{nameof(options.Value.ServerName)}' should only contain the name of the server like 'myserver' and NOT the full host name like 'myserver{options.Value.ServerNameSuffix}'", nameof(options.Value.ServerName));
+                // Use AAD managed identity
+                // https://learn.microsoft.com/en-us/azure/postgresql/single-server/how-to-connect-with-managed-identity
+                // https://learn.microsoft.com/en-us/azure/postgresql/single-server/concepts-azure-ad-authentication
+
+                var accessToken = await tokenCredential.GetTokenAsync(
+                    new TokenRequestContext(scopes: new string[] { azureDatabaseForPostgresqlScope }), System.Threading.CancellationToken.None);
+
+                return $"{connectionString}Password={accessToken.Token};";
+            }
+
+            return connectionString;
+        }
+
+        private string InternalGetConnectionString(PostgreSqlOptions options)
+        {
+            ArgumentException.ThrowIfNullOrEmpty(options.ServerName, nameof(options.ServerName));
+            ArgumentException.ThrowIfNullOrEmpty(options.ServerNameSuffix, nameof(options.ServerNameSuffix));
+            ArgumentException.ThrowIfNullOrEmpty(options.ServerPort, nameof(options.ServerPort));
+            ArgumentException.ThrowIfNullOrEmpty(options.ServerSslMode, nameof(options.ServerSslMode));
+            ArgumentException.ThrowIfNullOrEmpty(options.DatabaseName, nameof(options.DatabaseName));
+            ArgumentException.ThrowIfNullOrEmpty(options.DatabaseUserLogin, nameof(options.DatabaseUserLogin));
+
+            if (!options.UseManagedIdentity)
+            {
+                // Ensure password is set if NOT using Managed Identity
+                ArgumentException.ThrowIfNullOrEmpty(options.DatabaseUserPassword, nameof(options.DatabaseUserPassword));
+            }
+
+            if (options.UseManagedIdentity && !string.IsNullOrWhiteSpace(options.DatabaseUserPassword))
+            {
+                // throw if password IS set when using Managed Identity
+                throw new ArgumentException("DatabaseUserPassword shall not be set if UseManagedIdentity is true");
+            }
+
+            if (options.ServerName.Contains(options.ServerNameSuffix, StringComparison.OrdinalIgnoreCase))
+            {
+                throw new ArgumentException($"'{nameof(options.ServerName)}' should only contain the name of the server like 'myserver' and NOT the full host name like 'myserver{options.ServerNameSuffix}'", nameof(options.ServerName));
             }
 
             var connectionStringBuilder = new StringBuilder();
-            connectionStringBuilder.Append($"Server={options.Value.ServerName}{options.Value.ServerNameSuffix};");
-            connectionStringBuilder.Append($"Database={options.Value.DatabaseName};");
-            connectionStringBuilder.Append($"Port={options.Value.ServerPort};");
-            connectionStringBuilder.Append($"User Id={options.Value.DatabaseUserLogin};");
-            connectionStringBuilder.Append($"Password={options.Value.DatabaseUserPassword};");
-            connectionStringBuilder.Append($"SSL Mode={options.Value.ServerSslMode};");
+            connectionStringBuilder.Append($"Server={options.ServerName}{options.ServerNameSuffix};");
+            connectionStringBuilder.Append($"Database={options.DatabaseName};");
+            connectionStringBuilder.Append($"Port={options.ServerPort};");
+            connectionStringBuilder.Append($"SSL Mode={options.ServerSslMode};");
+            connectionStringBuilder.Append($"User Id={options.DatabaseUserLogin};");
+
+            if (!options.UseManagedIdentity)
+            {
+                connectionStringBuilder.Append($"Password={options.DatabaseUserPassword};");
+            }
+
             return connectionStringBuilder.ToString();
         }
     }

src/TesApi.Tests/Repository/TesTaskPostgreSqlRepositoryIntegrationTests.cs

@@ -61,10 +61,9 @@ await PostgreSqlTestUtility.CreateTestDbAsync(
                 DatabaseUserPassword = adminPw
             };
 
-            var optionsMock = new Mock<IOptions<PostgreSqlOptions>>();
-            optionsMock.Setup(x => x.Value).Returns(options);
-            var connectionString = new ConnectionStringUtility().GetPostgresConnectionString(optionsMock.Object);
-            repository = new TesTaskPostgreSqlRepository(() => new TesDbContext(connectionString));
+            var optionsMock = new Mock<PostgreSqlOptions>();
+            var connectionString = new PostgresConnectionStringUtility(optionsMock.Object, null);
+            repository = new TesTaskPostgreSqlRepository();
             Console.WriteLine("Creation complete.");
         }
 

src/TesApi.Web/Startup.cs

@@ -20,6 +20,7 @@
 using Tes.ApiClients.Options;
 using Tes.Models;
 using Tes.Repository;
+using Tes.Utilities;
 using TesApi.Filters;
 using TesApi.Web.Management;
 using TesApi.Web.Management.Batch;
@@ -77,6 +78,10 @@ public void ConfigureServices(IServiceCollection services)
                     .Configure<MarthaOptions>(configuration.GetSection(MarthaOptions.SectionName))
 
                     .AddMemoryCache(o => o.ExpirationScanFrequency = TimeSpan.FromHours(12))
+
+                    .AddSingleton<TokenCredential, DefaultAzureCredential>()
+                    .AddSingleton<PostgresConnectionStringUtility>()
+                    .AddDbContext<TesDbContext>(ServiceLifetime.Scoped)
                     .AddSingleton<ICache<TesTaskDatabaseItem>, TesRepositoryCache<TesTaskDatabaseItem>>()
                     .AddSingleton<TesTaskPostgreSqlRepository>()
                     .AddSingleton<AzureProxy>()
@@ -108,7 +113,6 @@ public void ConfigureServices(IServiceCollection services)
                     .AddSingleton<AzureManagementClientsFactory>()
                     .AddSingleton<ConfigurationUtils>()
                     .AddSingleton<IAllowedVmSizesService, AllowedVmSizesService>()
-                    .AddSingleton<TokenCredential>(s => new DefaultAzureCredential())
                     .AddSingleton<TaskToNodeTaskConverter>()
                     .AddSingleton<TaskExecutionScriptingManager>()
                     .AddTransient<BatchNodeScriptBuilder>()