Generalise DbContextResolver and query timing

Shared.Tests/DbContextResolverTests.cs

@@ -1,4 +1,7 @@
-using Microsoft.EntityFrameworkCore;
+using Microsoft.EntityFrameworkCore;
+using Microsoft.EntityFrameworkCore.Diagnostics;
+using Microsoft.EntityFrameworkCore.SqlServer;
+using Microsoft.EntityFrameworkCore.Infrastructure;
 using Microsoft.Extensions.DependencyInjection;
 using Imposter.Abstractions;
 using Shared.EntityFramework;
@@ -19,6 +22,20 @@ public TestDbContext(DbContextOptions<TestDbContext> options) : base(options) {
     }
 }
 
+public sealed class DisposableTestDbContext : DbContext
+{
+    public DisposableTestDbContext(DbContextOptions<DisposableTestDbContext> options) : base(options) { }
+
+    public bool WasDisposed { get; private set; }
+
+    public override void Dispose()
+    {
+        this.WasDisposed = true;
+        base.Dispose();
+    }
+}
+
+public sealed class NoOpCommandInterceptor : DbCommandInterceptor { }
 public class DbContextResolverTests {
     [Fact]
     public void ResolvedDbContext_RejectsNullScope()
@@ -111,4 +128,115 @@ public void Resolve_WithConnectionIdentifier_UpdatesInitialCatalog() {
         result.Context.Database.GetDbConnection().Database.ShouldBe("DefaultDb-Tenant1");
         result.Dispose();
     }
+
+    [Fact]
+    public void Resolve_WithSuffix_UsesConfiguredRetryOptionsAndInterceptors()
+    {
+        IConfigurationRoot configuration = new ConfigurationBuilder()
+            .AddInMemoryCollection(new Dictionary<string, string?>
+            {
+                ["ConnectionStrings:ReadModel"] = "Server=localhost;Initial Catalog=ReadModel;Integrated Security=True;TrustServerCertificate=True"
+            })
+            .Build();
+        NoOpCommandInterceptor interceptor = new();
+        DbContextResolverOptions resolverOptions = new()
+        {
+            RetryOptionsFactory = key =>
+            {
+                key.ShouldBe("ReadModel");
+                return new SqlServerRetryOptions
+                {
+                    MaxRetryCount = 2,
+                    MaxRetryDelay = TimeSpan.FromSeconds(5),
+                    AdditionalTransientErrorNumbers = [4060]
+                };
+            },
+            Interceptors = [interceptor]
+        };
+        using ServiceProvider rootProvider = new ServiceCollection().BuildServiceProvider();
+        DbContextResolver<TestDbContext> resolver = new(rootProvider, configuration, resolverOptions);
+
+        using ResolvedDbContext<TestDbContext> resolved = resolver.Resolve("ReadModel", "estate-123");
+
+        resolved.Context.Database.GetConnectionString().ShouldContain("Initial Catalog=ReadModel-estate-123");
+        SqlServerRetryingExecutionStrategy strategy = (SqlServerRetryingExecutionStrategy)resolved.Context.Database.CreateExecutionStrategy();
+        strategy.MaxRetryCount.ShouldBe(2);
+        strategy.MaxRetryDelay.ShouldBe(TimeSpan.FromSeconds(5));
+        strategy.AdditionalErrorNumbers.ShouldBe(new[] { 4060 });
+        resolved.Context.GetService<IInterceptors>().Aggregate<IDbCommandInterceptor>().ShouldBe(interceptor);
+    }
+
+    [Fact]
+    public void Resolve_WithSuffix_UsesDefaultRetryPolicyWhenNoRetryOptionsAreConfigured()
+    {
+        IConfigurationRoot configuration = new ConfigurationBuilder()
+            .AddInMemoryCollection(new Dictionary<string, string?>
+            {
+                ["ConnectionStrings:ReadModel"] = "Server=localhost;Initial Catalog=ReadModel;Integrated Security=True;TrustServerCertificate=True"
+            })
+            .Build();
+        using ServiceProvider rootProvider = new ServiceCollection().BuildServiceProvider();
+        DbContextResolver<TestDbContext> resolver = new(rootProvider, configuration);
+
+        using ResolvedDbContext<TestDbContext> resolved = resolver.Resolve("ReadModel", "estate-123");
+
+        SqlServerRetryingExecutionStrategy strategy = (SqlServerRetryingExecutionStrategy)
+            resolved.Context.Database.CreateExecutionStrategy();
+        strategy.MaxRetryCount.ShouldBe(6);
+        strategy.MaxRetryDelay.ShouldBe(TimeSpan.FromSeconds(30));
+    }
+
+    [Fact]
+    public void Resolve_WithSuffix_DisposesTheDynamicallyCreatedContext()
+    {
+        IConfigurationRoot configuration = new ConfigurationBuilder()
+            .AddInMemoryCollection(new Dictionary<string, string?>
+            {
+                ["ConnectionStrings:ReadModel"] = "Server=localhost;Initial Catalog=ReadModel;Integrated Security=True;TrustServerCertificate=True"
+            })
+            .Build();
+        using ServiceProvider rootProvider = new ServiceCollection().BuildServiceProvider();
+        DbContextResolver<DisposableTestDbContext> resolver = new(rootProvider, configuration);
+
+        ResolvedDbContext<DisposableTestDbContext> resolved = resolver.Resolve("ReadModel", "estate-123");
+        DisposableTestDbContext context = resolved.Context;
+        resolved.Dispose();
+
+        context.WasDisposed.ShouldBeTrue();
+    }
+
+    [Fact]
+    public void Resolve_WithMalformedConnectionString_ThrowsBeforeResolvingContext()
+    {
+        IConfigurationRoot configuration = new ConfigurationBuilder()
+            .AddInMemoryCollection(new Dictionary<string, string?>
+            {
+                ["ConnectionStrings:ReadModel"] = "NotAValidConnectionString"
+            })
+            .Build();
+        using ServiceProvider rootProvider = new ServiceCollection().BuildServiceProvider();
+        DbContextResolver<TestDbContext> resolver = new(rootProvider, configuration);
+
+        Should.Throw<ArgumentException>(() => resolver.Resolve("ReadModel", "estate-123"));
+    }
+
+    [Fact]
+    public void Resolve_WithNonSqlUnsuffixedConnectionString_ResolvesRegisteredContext()
+    {
+        ServiceCollection services = new();
+        services.AddDbContext<TestDbContext>(options => options.UseInMemoryDatabase("TestDbNonSqlConnectionString"));
+        IConfigurationRoot configuration = new ConfigurationBuilder()
+            .AddInMemoryCollection(new Dictionary<string, string?>
+            {
+                ["ConnectionStrings:ReadModel"] = "Mode=Memory;Cache=Shared;Data Source=test"
+            })
+            .Build();
+        using ServiceProvider rootProvider = services.BuildServiceProvider();
+        DbContextResolver<TestDbContext> resolver = new(rootProvider, configuration);
+
+        using ResolvedDbContext<TestDbContext> resolved = resolver.Resolve("ReadModel");
+
+        resolved.Context.ShouldNotBeNull();
+        resolved.Context.Database.ProviderName.ShouldContain("InMemory");
+    }
 }

Shared.Tests/QueryTimingInterceptorTests.cs

@@ -0,0 +1,107 @@
+#nullable enable
+using System;
+using System.Data;
+using System.Data.Common;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Data.SqlClient;
+using Microsoft.EntityFrameworkCore.Diagnostics;
+using Microsoft.Extensions.Logging;
+using Shared.EntityFramework;
+using Shouldly;
+using Xunit;
+
+namespace Shared.Tests;
+
+public sealed class QueryTimingInterceptorTests
+{
+    [Fact]
+    public async Task Interceptor_LogsSyncAndAsyncReaderScalarAndNonQueryCommandsWithoutSqlByDefault()
+    {
+        RecordingLoggerProvider provider = new();
+        using ILoggerFactory loggerFactory = LoggerFactory.Create(builder => builder.AddProvider(provider));
+        QueryTimingInterceptor interceptor = new(loggerFactory.CreateLogger<QueryTimingInterceptor>());
+        using DbCommand command = new SqlCommand("SELECT sensitive_value");
+
+        using DbDataReader reader = new DataTable().CreateDataReader();
+        interceptor.ReaderExecuted(command, CreateEventData(command, DbCommandMethod.ExecuteReader), reader);
+        await interceptor.ReaderExecutedAsync(command, CreateEventData(command, DbCommandMethod.ExecuteReader), reader, TestContext.Current.CancellationToken);
+        interceptor.ScalarExecuted(command, CreateEventData(command, DbCommandMethod.ExecuteScalar), 1);
+        await interceptor.ScalarExecutedAsync(command, CreateEventData(command, DbCommandMethod.ExecuteScalar), 1, TestContext.Current.CancellationToken);
+        interceptor.NonQueryExecuted(command, CreateEventData(command, DbCommandMethod.ExecuteNonQuery), 1);
+        await interceptor.NonQueryExecutedAsync(command, CreateEventData(command, DbCommandMethod.ExecuteNonQuery), 1, TestContext.Current.CancellationToken);
+
+        provider.Count.ShouldBe(6);
+        provider.Messages.ShouldAllBe(message => message.Contains("completed in 600 ms"));
+        provider.Messages.ShouldAllBe(message => !message.Contains("sensitive_value"));
+    }
+
+    [Fact]
+    public void Interceptor_UsesConfiguredThresholdAndCanIncludeCommandText()
+    {
+        RecordingLoggerProvider provider = new();
+        using ILoggerFactory loggerFactory = LoggerFactory.Create(builder => builder.AddProvider(provider));
+        QueryTimingInterceptor interceptor = new(loggerFactory.CreateLogger<QueryTimingInterceptor>(), new QueryTimingInterceptorOptions
+        {
+            ThresholdMilliseconds = 750,
+            IncludeCommandText = true
+        });
+        using DbCommand command = new SqlCommand("SELECT explicitly_enabled_value");
+
+        interceptor.ScalarExecuted(command, CreateEventData(command, DbCommandMethod.ExecuteScalar, TimeSpan.FromMilliseconds(749)), 1);
+        provider.Count.ShouldBe(0);
+
+        interceptor.ScalarExecuted(command, CreateEventData(command, DbCommandMethod.ExecuteScalar, TimeSpan.FromMilliseconds(750)), 1);
+        provider.Count.ShouldBe(1);
+        provider.Messages[0].ShouldContain("explicitly_enabled_value");
+    }
+
+    private static CommandExecutedEventData CreateEventData(DbCommand command,
+                                                              DbCommandMethod method,
+                                                              TimeSpan? duration = null)
+    {
+        return new CommandExecutedEventData(
+            null!,
+            (_, _) => String.Empty,
+            null!,
+            command,
+            command.CommandText,
+            null,
+            method,
+            Guid.NewGuid(),
+            Guid.NewGuid(),
+            null,
+            false,
+            false,
+            DateTimeOffset.UtcNow,
+            duration ?? TimeSpan.FromMilliseconds(600),
+            CommandSource.Unknown);
+    }
+
+    private sealed class RecordingLoggerProvider : ILoggerProvider
+    {
+        public int Count => this.Messages.Count;
+        public System.Collections.Generic.List<string> Messages { get; } = new();
+
+        public ILogger CreateLogger(string categoryName) => new RecordingLogger(this.Messages);
+        public void Dispose() { }
+    }
+
+    private sealed class RecordingLogger : ILogger
+    {
+        private readonly System.Collections.Generic.List<string> _messages;
+
+        public RecordingLogger(System.Collections.Generic.List<string> messages) => _messages = messages;
+        public IDisposable? BeginScope<TState>(TState state) where TState : notnull => null;
+        public bool IsEnabled(LogLevel logLevel) => true;
+        public void Log<TState>(LogLevel logLevel,
+                                EventId eventId,
+                                TState state,
+                                Exception? exception,
+                                Func<TState, Exception?, string> formatter)
+        {
+            if (logLevel == LogLevel.Warning)
+                _messages.Add(formatter(state, exception));
+        }
+    }
+}

Shared/EntityFramework/DbContextResolver.cs

@@ -1,13 +1,10 @@
-using Microsoft.EntityFrameworkCore;
+#nullable enable
+using Microsoft.EntityFrameworkCore;
+using Microsoft.EntityFrameworkCore.Diagnostics;
+using Microsoft.Data.SqlClient;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using System;
-using System.Collections.Generic;
-using System.Diagnostics;
-using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
-using Microsoft.Data.SqlClient;
 
 namespace Shared.EntityFramework;
 
@@ -17,43 +14,83 @@ public interface IDbContextResolver<TContext> where TContext : DbContext
     ResolvedDbContext<TContext> Resolve(String connectionStringKey, String databaseNameSuffix);
 }
 
-public class DbContextResolver<TContext> : IDbContextResolver<TContext> where TContext : DbContext {
+public class DbContextResolver<TContext> : IDbContextResolver<TContext> where TContext : DbContext
+{
     private readonly IServiceProvider _rootProvider;
     private readonly IConfiguration _config;
+    private readonly DbContextResolverOptions _options;
+
+    public DbContextResolver(IServiceProvider rootProvider, IConfiguration config)
+        : this(rootProvider, config, new DbContextResolverOptions())
+    {
+    }
 
     public DbContextResolver(IServiceProvider rootProvider,
-                             IConfiguration config) {
-        _rootProvider = rootProvider;
-        _config = config;
+                             IConfiguration config,
+                             DbContextResolverOptions options)
+    {
+        _rootProvider = rootProvider ?? throw new ArgumentNullException(nameof(rootProvider));
+        _config = config ?? throw new ArgumentNullException(nameof(config));
+        _options = options ?? throw new ArgumentNullException(nameof(options));
     }
 
-    public ResolvedDbContext<TContext> Resolve(String connectionStringKey) {
+    public ResolvedDbContext<TContext> Resolve(String connectionStringKey)
+    {
         return this.Resolve(connectionStringKey, String.Empty);
     }
 
-    public ResolvedDbContext<TContext> Resolve(String connectionStringKey,
-                                               String databaseNameSuffix) {
-        IServiceScope scope = _rootProvider.CreateScope();
-        String connectionString = _config.GetConnectionString(connectionStringKey);
+    public ResolvedDbContext<TContext> Resolve(String connectionStringKey, String databaseNameSuffix)
+    {
+        String? connectionString = _config.GetConnectionString(connectionStringKey);
         if (String.IsNullOrWhiteSpace(connectionString))
             throw new InvalidOperationException($"Connection string for '{connectionStringKey}' not found.");
 
-        // Update the connection string with the identifier if needed
-        if (!String.IsNullOrWhiteSpace(databaseNameSuffix)) {
-            SqlConnectionStringBuilder builder = new(connectionString);
-            builder.InitialCatalog = $"{builder.InitialCatalog}-{databaseNameSuffix}";
-            connectionString = builder.ConnectionString;
+        if (String.IsNullOrWhiteSpace(databaseNameSuffix))
+        {
+            IServiceScope regularScope = _rootProvider.CreateScope();
+            try
+            {
+                return new ResolvedDbContext<TContext>(regularScope);
+            }
+            catch
+            {
+                regularScope.Dispose();
+                throw;
+            }
+        }
 
+        SqlConnectionStringBuilder connectionStringBuilder = new(connectionString);
+        connectionStringBuilder.InitialCatalog = $"{connectionStringBuilder.InitialCatalog}-{databaseNameSuffix}";
+        connectionString = connectionStringBuilder.ConnectionString;
 
-            // Create an isolated service collection and provider
-            ServiceCollection services = new();
-            services.AddDbContext<TContext>(options => { options.UseSharedSqlServer<TContext>(connectionString); });
+        IServiceScope scope = _rootProvider.CreateScope();
+        try
+        {
+            DbContextOptionsBuilder<TContext> optionsBuilder = new();
+            SqlServerRetryOptions retryOptions = _options.RetryOptionsFactory?.Invoke(connectionStringKey)
+                ?? new SqlServerRetryOptions
+                {
+                    MaxRetryCount = SqlServerDbContextOptionsBuilderExtensions.DefaultMaxRetryCount,
+                    MaxRetryDelay = SqlServerDbContextOptionsBuilderExtensions.DefaultMaxRetryDelay
+                };
 
-            ServiceProvider provider = services.BuildServiceProvider();
-            scope = provider.CreateScope();
-            // Standard resolution using DI container
-        }
+            optionsBuilder.UseSharedSqlServer<TContext>(connectionString, retry =>
+            {
+                retry.MaxRetryCount = retryOptions.MaxRetryCount;
+                retry.MaxRetryDelay = retryOptions.MaxRetryDelay;
+                retry.AdditionalTransientErrorNumbers = retryOptions.AdditionalTransientErrorNumbers;
+            });
 
-        return new ResolvedDbContext<TContext>(scope);
+            if (_options.Interceptors.Count > 0)
+                optionsBuilder.AddInterceptors(_options.Interceptors);
+
+            TContext context = ActivatorUtilities.CreateInstance<TContext>(scope.ServiceProvider, optionsBuilder.Options);
+            return new ResolvedDbContext<TContext>(scope, context);
+        }
+        catch
+        {
+            scope.Dispose();
+            throw;
+        }
     }
 }

Shared/EntityFramework/DbContextResolverOptions.cs

@@ -0,0 +1,16 @@
+#nullable enable
+using System;
+using System.Collections.Generic;
+using Microsoft.EntityFrameworkCore.Diagnostics;
+
+namespace Shared.EntityFramework;
+
+/// <summary>Options for dynamically resolved DbContexts.</summary>
+public sealed class DbContextResolverOptions
+{
+    /// <summary>Gets a retry policy for a connection string key. Returning an empty options object disables retries.</summary>
+    public Func<string, SqlServerRetryOptions?>? RetryOptionsFactory { get; init; }
+
+    /// <summary>Gets interceptors to apply to contexts created with a database suffix.</summary>
+    public IReadOnlyCollection<DbCommandInterceptor> Interceptors { get; init; } = Array.Empty<DbCommandInterceptor>();
+}