Review feedback for Azure.EFCore.Npgsql integration

playground/PostgresEndToEnd/PostgresEndToEnd.ApiService/PostgresEndToEnd.ApiService.csproj

@@ -9,6 +9,7 @@
   <ItemGroup>
     <AspireProjectOrPackageReference Include="Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
     <ProjectReference Include="..\..\Playground.ServiceDefaults\Playground.ServiceDefaults.csproj" />

playground/TestShop/CatalogDb/CatalogDb.csproj

@@ -15,6 +15,7 @@
     </PackageReference>
 
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
   </ItemGroup>

playground/TestShop/CatalogModel/CatalogModel.csproj

@@ -9,6 +9,7 @@
   <ItemGroup>
     <AspireProjectOrPackageReference Include="Aspire.Npgsql.EntityFrameworkCore.PostgreSQL" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
   </ItemGroup>

playground/TestShop/CatalogService/CatalogService.csproj

@@ -18,6 +18,7 @@
     <ProjectReference Include="..\TestShop.ServiceDefaults\TestShop.ServiceDefaults.csproj" />
 
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
   </ItemGroup>

playground/bicep/BicepSample.ApiService/BicepSample.ApiService.csproj

@@ -15,6 +15,7 @@
     <AspireProjectOrPackageReference Include="Aspire.Microsoft.EntityFrameworkCore.SqlServer" />
     <AspireProjectOrPackageReference Include="Aspire.Npgsql.EntityFrameworkCore.PostgreSQL" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
     <AspireProjectOrPackageReference Include="Aspire.StackExchange.Redis" />

playground/cdk/CdkSample.ApiService/CdkSample.ApiService.csproj

@@ -15,6 +15,7 @@
     <AspireProjectOrPackageReference Include="Aspire.Microsoft.EntityFrameworkCore.SqlServer" />
     <AspireProjectOrPackageReference Include="Aspire.Npgsql.EntityFrameworkCore.PostgreSQL" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
     <AspireProjectOrPackageReference Include="Aspire.StackExchange.Redis" />

playground/publishers/Publishers.ApiService/Publishers.ApiService.csproj

@@ -9,6 +9,7 @@
   <ItemGroup>
     <AspireProjectOrPackageReference Include="Aspire.Npgsql.EntityFrameworkCore.PostgreSQL" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
     <ProjectReference Include="..\..\Playground.ServiceDefaults\Playground.ServiceDefaults.csproj" />

playground/publishers/Publishers.DbSetup/Publishers.DbSetup.csproj

@@ -10,6 +10,7 @@
   <ItemGroup>
     <ProjectReference Include="..\..\..\src\Components\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL.csproj" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
     <ProjectReference Include="..\Publishers.Common\Publishers.Common.csproj" />

playground/waitfor/WaitForSandbox.ApiService/WaitForSandbox.ApiService.csproj

@@ -9,6 +9,7 @@
   <ItemGroup>
     <AspireProjectOrPackageReference Include="Aspire.Npgsql.EntityFrameworkCore.PostgreSQL" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
     <ProjectReference Include="..\..\Playground.ServiceDefaults\Playground.ServiceDefaults.csproj" />

playground/waitfor/WaitForSandbox.DbSetup/WaitForSandbox.DbSetup.csproj

@@ -10,6 +10,7 @@
   <ItemGroup>
     <ProjectReference Include="..\..\..\src\Components\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL.csproj" />
     <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+    <!-- This is to override CentralPackageTransitivePinningEnabled -->
     <PackageReference Include="Npgsql.DependencyInjection" VersionOverride="$(Npgsql8Version)" />
     <PackageReference Include="Npgsql.OpenTelemetry" VersionOverride="$(Npgsql8Version)" />
     <ProjectReference Include="..\WaitForSandbox.Common\WaitForSandbox.Common.csproj" />

src/Components/Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL/AspireAzureEFPostgreSqlExtensions.cs

@@ -6,13 +6,13 @@
 using Aspire;
 using Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL;
 using Aspire.Npgsql.EntityFrameworkCore.PostgreSQL;
-using Azure.Identity;
 using Microsoft.EntityFrameworkCore;
-using Npgsql;
 using Npgsql.EntityFrameworkCore.PostgreSQL.Infrastructure.Internal;
 
 #if NET9_0_OR_GREATER
 using Microsoft.Extensions.DependencyInjection;
+#else
+using Npgsql;
 #endif
 
 namespace Microsoft.Extensions.Hosting;
@@ -121,51 +121,25 @@ private static void CopySettings(NpgsqlEntityFrameworkCorePostgreSQLSettings sou
 
     private static void ConfigureDbContextOptionsBuilder(AzureNpgsqlEntityFrameworkCorePostgreSQLSettings settings, DbContextOptionsBuilder dbContextOptionsBuilder)
     {
-        // The connection string requires the username to be provided. Since it will depend on the Managed Identity that is used
-        // we attempt to get the username from the access token.
-
-        var credential = settings.Credential ?? new DefaultAzureCredential();
-
 #pragma warning disable EF1001 // Internal EF Core API usage.
 
         // Get the connection string from the Npgsql options extension in case it was set using UseNpgsql(connStr) and Enrich()
-        var extensionsConnectionString = dbContextOptionsBuilder.Options.GetExtension<NpgsqlOptionsExtension>()?.ConnectionString;
+        var connectionString = settings.ConnectionString ?? dbContextOptionsBuilder.Options.GetExtension<NpgsqlOptionsExtension>()?.ConnectionString;
 
 #pragma warning restore EF1001 // Internal EF Core API usage.
 
-        var dataSourceBuilder = new NpgsqlDataSourceBuilder(settings.ConnectionString ?? extensionsConnectionString);
-
-        if (string.IsNullOrEmpty(dataSourceBuilder.ConnectionStringBuilder.Username))
+#if NET9_0_OR_GREATER
+        dbContextOptionsBuilder.UseNpgsql(options =>
         {
-            // Ensure to use the management scope, so the token contains user names for all managed identity types - e.g. user and service principal
-            var token = credential.GetToken(ManagedIdentityTokenCredentialHelpers.ManagementTokenRequestContext, default);
-
-            if (ManagedIdentityTokenCredentialHelpers.TryGetUsernameFromToken(token.Token, out var username))
-            {
-                dataSourceBuilder.ConnectionStringBuilder.Username = username;
-            }
-            else
-            {
-                // Otherwise check using the PostgresSql scope
-                token = credential.GetToken(ManagedIdentityTokenCredentialHelpers.DatabaseForPostgresSqlTokenRequestContext, default);
-
-                if (ManagedIdentityTokenCredentialHelpers.TryGetUsernameFromToken(token.Token, out username))
-                {
-                    dataSourceBuilder.ConnectionStringBuilder.Username = username;
-                }
-            }
-
-            // If we still don't have a username, we let Npgsql handle the error when trying to connect.
-        }
+            options.ConfigureDataSource(dataSourceBuilder => dataSourceBuilder.ConfigureEntraIdAuthentication(settings.Credential));
+        });
+#else
+        var dataSourceBuilder = new NpgsqlDataSourceBuilder(connectionString);
 
-        if (string.IsNullOrEmpty(dataSourceBuilder.ConnectionStringBuilder.Password))
+        if (dataSourceBuilder.ConfigureEntraIdAuthentication(settings.Credential))
         {
-            dataSourceBuilder.UsePasswordProvider(
-                passwordProvider: _ => credential.GetToken(ManagedIdentityTokenCredentialHelpers.DatabaseForPostgresSqlTokenRequestContext, default).Token,
-                passwordProviderAsync: async (_, ct) => (await credential.GetTokenAsync(ManagedIdentityTokenCredentialHelpers.DatabaseForPostgresSqlTokenRequestContext, default).ConfigureAwait(false)).Token
-            );
-
             dbContextOptionsBuilder.UseNpgsql(dataSourceBuilder.Build());
         }
+#endif
     }
 }

src/Components/Aspire.Azure.Npgsql/AspireAzureNpgsqlExtensions.cs

@@ -5,7 +5,6 @@
 using Aspire;
 using Aspire.Azure.Npgsql;
 using Aspire.Npgsql;
-using Azure.Identity;
 using Microsoft.Extensions.DependencyInjection;
 using Npgsql;
 
@@ -38,7 +37,7 @@ public static void AddAzureNpgsqlDataSource(this IHostApplicationBuilder builder
         {
             Debug.Assert(azureSettings != null);
 
-            ConfigureDataSourceBuilder(azureSettings, dataSourceBuilder);
+            dataSourceBuilder.ConfigureEntraIdAuthentication(azureSettings.Credential);
 
             configureDataSourceBuilder?.Invoke(dataSourceBuilder);
         });
@@ -67,7 +66,7 @@ public static void AddKeyedAzureNpgsqlDataSource(this IHostApplicationBuilder bu
         {
             Debug.Assert(azureSettings != null);
 
-            ConfigureDataSourceBuilder(azureSettings, dataSourceBuilder);
+            dataSourceBuilder.ConfigureEntraIdAuthentication(azureSettings.Credential);
 
             configureDataSourceBuilder?.Invoke(dataSourceBuilder);
         });
@@ -96,46 +95,4 @@ private static void CopySettings(NpgsqlSettings source, NpgsqlSettings destinati
         destination.DisableMetrics = source.DisableMetrics;
         destination.DisableTracing = source.DisableTracing;
     }
-
-    private static void ConfigureDataSourceBuilder(AzureNpgsqlSettings settings, NpgsqlDataSourceBuilder dataSourceBuilder)
-    {
-        // The connection string requires the username to be provided. Since it will depend on the Managed Identity that is used
-        // we attempt to get the username from the access token.
-
-        var credential = settings.Credential ?? new DefaultAzureCredential();
-
-        if (string.IsNullOrEmpty(dataSourceBuilder.ConnectionStringBuilder.Username))
-        {
-            // Ensure to use the management scope, so the token contains user names for all managed identity types - e.g. user and service principal
-            var token = credential.GetToken(ManagedIdentityTokenCredentialHelpers.ManagementTokenRequestContext, default);
-
-            if (ManagedIdentityTokenCredentialHelpers.TryGetUsernameFromToken(token.Token, out var username))
-            {
-                dataSourceBuilder.ConnectionStringBuilder.Username = username;
-            }
-            else
-            {
-                // Otherwise check using the PostgresSql scope
-                token = credential.GetToken(ManagedIdentityTokenCredentialHelpers.DatabaseForPostgresSqlTokenRequestContext, default);
-
-                if (ManagedIdentityTokenCredentialHelpers.TryGetUsernameFromToken(token.Token, out username))
-                {
-                    dataSourceBuilder.ConnectionStringBuilder.Username = username;
-                }
-            }
-
-            // If we still don't have a username, we let Npgsql handle the error when trying to connect.
-            // The user will be hinted to provide a username by using the configureDataSourceBuilder callback.
-        }
-
-        if (string.IsNullOrEmpty(dataSourceBuilder.ConnectionStringBuilder.Password))
-        {
-            // The token is not cached since it is refreshed for each new physical connection, or when it has expired.
-
-            dataSourceBuilder.UsePasswordProvider(
-                passwordProvider: _ => credential.GetToken(ManagedIdentityTokenCredentialHelpers.DatabaseForPostgresSqlTokenRequestContext, default).Token,
-                passwordProviderAsync: async (_, ct) => (await credential.GetTokenAsync(ManagedIdentityTokenCredentialHelpers.DatabaseForPostgresSqlTokenRequestContext, default).ConfigureAwait(false)).Token
-            );
-        }
-    }
 }

src/Components/Aspire.Npgsql.EntityFrameworkCore.PostgreSQL/Aspire.Npgsql.EntityFrameworkCore.PostgreSQL.csproj

@@ -27,6 +27,7 @@
   </ItemGroup>
 
   <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+  <!-- This is to override CentralPackageTransitivePinningEnabled -->
   <ItemGroup Condition="'$(TargetFramework)' == 'net8.0'">
     <PackageVersion Update="Npgsql.DependencyInjection" Version="$(Npgsql8Version)" />
     <PackageVersion Update="Npgsql.OpenTelemetry" Version="$(Npgsql8Version)" />

src/Components/Common/EntityFrameworkUtils.cs

@@ -36,6 +36,7 @@ public static TSettings GetDbContextSettings<TContext, TSettings>(this IHostAppl
 
         return settings;
     }
+
     /// <summary>
     /// Ensures a <see cref="DbContext"/> is registered in DI.
     /// </summary>

src/Components/Common/ManagedIdentityTokenCredentialHelpers.cs

@@ -3,18 +3,69 @@
 
 using System.Text.Json;
 using Azure.Core;
+using Azure.Identity;
+using Npgsql;
 
 namespace Aspire;
 
-internal sealed class ManagedIdentityTokenCredentialHelpers
+internal static class ManagedIdentityTokenCredentialHelpers
 {
-    public const string AzureDatabaseForPostgresSqlScope = "https://ossrdbms-aad.database.windows.net/.default";
-    public const string AzureManagementScope = "https://management.azure.com/.default";
+    private const string AzureDatabaseForPostgresSqlScope = "https://ossrdbms-aad.database.windows.net/.default";
+    private const string AzureManagementScope = "https://management.azure.com/.default";
 
-    public static readonly TokenRequestContext DatabaseForPostgresSqlTokenRequestContext = new([AzureDatabaseForPostgresSqlScope]);
-    public static readonly TokenRequestContext ManagementTokenRequestContext = new([AzureManagementScope]);
+    private static readonly TokenRequestContext s_databaseForPostgresSqlTokenRequestContext = new([AzureDatabaseForPostgresSqlScope]);
+    private static readonly TokenRequestContext s_managementTokenRequestContext = new([AzureManagementScope]);
 
-    public static bool TryGetUsernameFromToken(string jwtToken, out string? username)
+    public static bool ConfigureEntraIdAuthentication(this NpgsqlDataSourceBuilder dataSourceBuilder, TokenCredential? credential)
+    {
+        credential ??= new DefaultAzureCredential();
+        var configuredAuth = false;
+
+        // The connection string requires the username to be provided. Since it will depend on the Managed Identity that is used
+        // we attempt to get the username from the access token if it's not defined.
+
+        if (string.IsNullOrEmpty(dataSourceBuilder.ConnectionStringBuilder.Username))
+        {
+            // Ensure to use the management scope, so the token contains user names for all managed identity types - e.g. user and service principal
+            var token = credential.GetToken(s_managementTokenRequestContext, default);
+
+            if (TryGetUsernameFromToken(token.Token, out var username))
+            {
+                dataSourceBuilder.ConnectionStringBuilder.Username = username;
+                configuredAuth = true;
+            }
+            else
+            {
+                // Otherwise check using the PostgresSql scope
+                token = credential.GetToken(s_databaseForPostgresSqlTokenRequestContext, default);
+
+                if (TryGetUsernameFromToken(token.Token, out username))
+                {
+                    dataSourceBuilder.ConnectionStringBuilder.Username = username;
+                    configuredAuth = true;
+                }
+            }
+
+            // If we still don't have a username, we let Npgsql handle the error when trying to connect.
+            // The user will be hinted to provide a username by using the configureDataSourceBuilder callback.
+        }
+
+        if (string.IsNullOrEmpty(dataSourceBuilder.ConnectionStringBuilder.Password))
+        {
+            // The token is not cached since it is refreshed for each new physical connection, or when it has expired.
+
+            dataSourceBuilder.UsePasswordProvider(
+                passwordProvider: _ => credential.GetToken(s_databaseForPostgresSqlTokenRequestContext, default).Token,
+                passwordProviderAsync: async (_, ct) => (await credential.GetTokenAsync(s_databaseForPostgresSqlTokenRequestContext, default).ConfigureAwait(false)).Token
+            );
+
+            configuredAuth = true;
+        }
+
+        return configuredAuth;
+    }
+
+    private static bool TryGetUsernameFromToken(string jwtToken, out string? username)
     {
         username = null;
 

tests/Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL.Tests/Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL.Tests.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
     <TargetFrameworks>$(AllTargetFrameworks)</TargetFrameworks>
@@ -7,6 +7,11 @@
   <ItemGroup>
     <Compile Include="$(RepoRoot)src\Aspire.Hosting.PostgreSQL\PostgresContainerImageTags.cs" />
     <Compile Include="..\Aspire.Npgsql.Tests\PostgreSQLContainerFixture.cs" />
+    <Compile Include="..\Aspire.Azure.Npgsql.Tests\FakeTokenCredential.cs" />
+    <Compile Include="..\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL.Tests\WorkaroundToReadProtectedField.cs" />
+    <Compile Include="..\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL.Tests\TestDbContext.cs" />
+    <Compile Include="..\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL.Tests\CustomExecutionStrategy.cs" />
+    <Compile Include="..\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL.Tests\CustomRetryExecutionStrategy.cs" />
 
     <None Include="$(RepoRoot)src\Components\Aspire.Npgsql.EntityFrameworkCore.PostgreSQL\ConfigurationSchema.json" CopyToOutputDirectory="PreserveNewest" />
 
@@ -18,6 +23,7 @@
   </ItemGroup>
 
   <!-- Npgsql EF needs to match the same major version as the underlying Npgsql assemblies. -->
+  <!-- This is to override CentralPackageTransitivePinningEnabled -->
   <ItemGroup Condition="'$(TargetFramework)' == 'net8.0'">
     <PackageVersion Update="Npgsql.DependencyInjection" Version="$(Npgsql8Version)" />
     <PackageVersion Update="Npgsql.OpenTelemetry" Version="$(Npgsql8Version)" />

tests/Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL.Tests/AspireAzureEFPostgreSqlExtensionsTests.cs

@@ -1,7 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using Aspire.Components.Common.Tests;
+using Aspire.Npgsql.Tests;
+using Aspire.TestUtilities;
 using Microsoft.EntityFrameworkCore;
 using Microsoft.EntityFrameworkCore.Infrastructure.Internal;
 using Microsoft.EntityFrameworkCore.Storage;

tests/Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL.Tests/ConformanceTests.cs

@@ -1,7 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using Aspire.Components.Common.Tests;
 using Aspire.Components.ConformanceTests;
 using Aspire.Npgsql.Tests;
 using Aspire.TestUtilities;

tests/Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL.Tests/ConformanceTests_TypeSpecificConfig.cs

@@ -1,8 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using Aspire.Components.Common.Tests;
 using Aspire.Npgsql.Tests;
+using Aspire.TestUtilities;
 using Microsoft.Extensions.Configuration;
 
 namespace Aspire.Azure.Npgsql.EntityFrameworkCore.PostgreSQL.Tests;