Support multiple connection strings by managing multiple data sources internally

Author: roji

src/EFCore.PG/Properties/NpgsqlStrings.Designer.cs

@@ -247,7 +247,4 @@
   <data name="StoredProcedureReturnValueNotSupported" xml:space="preserve">
     <value>The entity type '{entityType}' is mapped to the stored procedure '{sproc}', which is configured with result columns. PostgreSQL stored procedures do not support return values; use output parameters instead.</value>
   </data>
-  <data name="DataSourceWithMultipleConnectionStrings" xml:space="preserve">
-    <value>Different connection strings are being used, but the provider uses has been configured with a feature that requires a singleton data source internally: {dataSourceFeature}</value>
-  </data>
 </root>

src/EFCore.PG/Properties/NpgsqlStrings.resx

@@ -1,7 +1,7 @@
+using System.Collections.Concurrent;
 using System.Data.Common;
 using Npgsql.EntityFrameworkCore.PostgreSQL.Infrastructure;
 using Npgsql.EntityFrameworkCore.PostgreSQL.Infrastructure.Internal;
-using Npgsql.EntityFrameworkCore.PostgreSQL.Internal;
 
 namespace Npgsql.EntityFrameworkCore.PostgreSQL.Storage.Internal;
 
@@ -27,11 +27,9 @@ namespace Npgsql.EntityFrameworkCore.PostgreSQL.Storage.Internal;
 /// </remarks>
 public class NpgsqlDataSourceManager : IDisposable, IAsyncDisposable
 {
-    private bool _isInitialized;
-    private string? _connectionString;
     private readonly IEnumerable<INpgsqlDataSourceConfigurationPlugin> _plugins;
-    private NpgsqlDataSource? _dataSource;
-    private readonly object _lock = new();
+    private readonly ConcurrentDictionary<string, NpgsqlDataSource> _dataSources = new();
+    private volatile int _isDisposed;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -70,39 +68,39 @@ public NpgsqlDataSourceManager(IEnumerable<INpgsqlDataSourceConfigurationPlugin>
             { ConnectionString: null } or null => null,
 
             // The following are features which require an NpgsqlDataSource, since they require configuration on NpgsqlDataSourceBuilder.
-            { EnumDefinitions.Count: > 0 } => GetSingletonDataSource(npgsqlOptionsExtension, "MapEnum"),
-            _ when _plugins.Any() => GetSingletonDataSource(npgsqlOptionsExtension, _plugins.First().GetType().Name),
+            { EnumDefinitions.Count: > 0 } => GetSingletonDataSource(npgsqlOptionsExtension),
+            _ when _plugins.Any() => GetSingletonDataSource(npgsqlOptionsExtension),
 
             // If there's no configured feature which requires us to use a data source internally, don't use one; this causes
             // NpgsqlRelationalConnection to use the connection string as before (no data source), allowing switching connection strings
             // with the same service provider etc.
             _ => null
         };
 
-    private DbDataSource GetSingletonDataSource(NpgsqlOptionsExtension npgsqlOptionsExtension, string dataSourceFeature)
+    private DbDataSource GetSingletonDataSource(NpgsqlOptionsExtension npgsqlOptionsExtension)
     {
-        if (!_isInitialized)
+        var connectionString = npgsqlOptionsExtension.ConnectionString;
+        Check.DebugAssert(connectionString is not null, "Connection string can't be null");
+
+        if (_dataSources.TryGetValue(connectionString, out var dataSource))
         {
-            lock (_lock)
-            {
-                if (!_isInitialized)
-                {
-                    _dataSource = CreateSingletonDataSource(npgsqlOptionsExtension);
-                    _connectionString = npgsqlOptionsExtension.ConnectionString;
-                    _isInitialized = true;
-                    return _dataSource;
-                }
-            }
+            return dataSource;
         }
 
-        Check.DebugAssert(_dataSource is not null, "_dataSource cannot be null at this point");
+        var newDataSource = CreateDataSource(npgsqlOptionsExtension);
 
-        if (_connectionString != npgsqlOptionsExtension.ConnectionString)
+        var addedDataSource = _dataSources.GetOrAdd(connectionString, newDataSource);
+        if (!ReferenceEquals(addedDataSource, newDataSource))
         {
-            throw new InvalidOperationException(NpgsqlStrings.DataSourceWithMultipleConnectionStrings(dataSourceFeature));
+            newDataSource.Dispose();
+        }
+        else if (_isDisposed == 1)
+        {
+            newDataSource.Dispose();
+            throw new ObjectDisposedException(nameof(NpgsqlDataSourceManager));
         }
 
-        return _dataSource;
+        return addedDataSource;
     }
 
     /// <summary>
@@ -111,7 +109,7 @@ private DbDataSource GetSingletonDataSource(NpgsqlOptionsExtension npgsqlOptions
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    protected virtual NpgsqlDataSource CreateSingletonDataSource(NpgsqlOptionsExtension npgsqlOptionsExtension)
+    protected virtual NpgsqlDataSource CreateDataSource(NpgsqlOptionsExtension npgsqlOptionsExtension)
     {
         var dataSourceBuilder = new NpgsqlDataSourceBuilder(npgsqlOptionsExtension.ConnectionString);
 
@@ -151,7 +149,15 @@ enumDefinition.StoreTypeSchema is null
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public void Dispose()
-        => _dataSource?.Dispose();
+    {
+        if (Interlocked.CompareExchange(ref _isDisposed, 1, 0) == 0)
+        {
+            foreach (var dataSource in _dataSources.Values)
+            {
+                dataSource.Dispose();
+            }
+        }
+    }
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -161,9 +167,12 @@ public void Dispose()
     /// </summary>
     public async ValueTask DisposeAsync()
     {
-        if (_dataSource != null)
+        if (Interlocked.CompareExchange(ref _isDisposed, 1, 0) == 0)
         {
-            await _dataSource.DisposeAsync().ConfigureAwait(false);
+            foreach (var dataSource in _dataSources.Values)
+            {
+                await dataSource.DisposeAsync().ConfigureAwait(false);
+            }
         }
     }
 }

src/EFCore.PG/Storage/Internal/NpgsqlDataSourceManager.cs

@@ -103,34 +103,60 @@ public void DbDataSource_from_application_service_provider_does_not_used_if_conn
     }
 
     [Fact]
-    public void Multiple_connection_strings_with_plugin_is_not_supported()
+    public void Multiple_connection_strings_with_plugin()
     {
         var context1 = new ConnectionStringSwitchingContext("Host=FakeHost1", withNetTopologySuite: true);
-        _ = context1.GetService<IRelationalConnection>();
-        var context2 = new ConnectionStringSwitchingContext("Host=FakeHost2", withNetTopologySuite: true);
-
-        var exception = Assert.Throws<InvalidOperationException>(() => context2.GetService<IRelationalConnection>());
-        Assert.Equal(NpgsqlStrings.DataSourceWithMultipleConnectionStrings("NetTopologySuiteDataSourceConfigurationPlugin"), exception.Message);
+        var connection1 = (NpgsqlRelationalConnection)context1.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost1", connection1.ConnectionString);
+        Assert.NotNull(connection1.DbDataSource);
+
+        var context2 = new ConnectionStringSwitchingContext("Host=FakeHost1", withNetTopologySuite: true);
+        var connection2 = (NpgsqlRelationalConnection)context2.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost1", connection2.ConnectionString);
+        Assert.Same(connection1.DbDataSource, connection2.DbDataSource);
+
+        var context3 = new ConnectionStringSwitchingContext("Host=FakeHost2", withNetTopologySuite: true);
+        var connection3 = (NpgsqlRelationalConnection)context3.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost2", connection3.ConnectionString);
+        Assert.NotSame(connection1.DbDataSource, connection3.DbDataSource);
     }
 
     [Fact]
-    public void Multiple_connection_strings_with_enum_is_not_supported()
+    public void Multiple_connection_strings_with_enum()
     {
         var context1 = new ConnectionStringSwitchingContext("Host=FakeHost1", withEnum: true);
-        _ = context1.GetService<IRelationalConnection>();
-        var context2 = new ConnectionStringSwitchingContext("Host=FakeHost2", withEnum: true);
-
-        var exception = Assert.Throws<InvalidOperationException>(() => context2.GetService<IRelationalConnection>());
-        Assert.Equal(NpgsqlStrings.DataSourceWithMultipleConnectionStrings("MapEnum"), exception.Message);
+        var connection1 = (NpgsqlRelationalConnection)context1.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost1", connection1.ConnectionString);
+        Assert.NotNull(connection1.DbDataSource);
+
+        var context2 = new ConnectionStringSwitchingContext("Host=FakeHost1", withEnum: true);
+        var connection2 = (NpgsqlRelationalConnection)context2.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost1", connection2.ConnectionString);
+        Assert.Same(connection1.DbDataSource, connection2.DbDataSource);
+
+        var context3 = new ConnectionStringSwitchingContext("Host=FakeHost2", withEnum: true);
+        var connection3 = (NpgsqlRelationalConnection)context3.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost2", connection3.ConnectionString);
+        Assert.NotSame(connection1.DbDataSource, connection3.DbDataSource);
     }
 
     [Fact]
-    public void Multiple_connection_strings_without_data_source_features_is_supported()
+    public void Multiple_connection_strings_without_data_source_features()
     {
         var context1 = new ConnectionStringSwitchingContext("Host=FakeHost1");
-        _ = context1.GetService<IRelationalConnection>();
-        var context2 = new ConnectionStringSwitchingContext("Host=FakeHost2");
-        _ = context2.GetService<IRelationalConnection>();
+        var connection1 = (NpgsqlRelationalConnection)context1.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost1", connection1.ConnectionString);
+        Assert.Null(connection1.DbDataSource);
+
+        var context2 = new ConnectionStringSwitchingContext("Host=FakeHost1");
+        var connection2 = (NpgsqlRelationalConnection)context2.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost1", connection2.ConnectionString);
+        Assert.Null(connection2.DbDataSource);
+
+        var context3 = new ConnectionStringSwitchingContext("Host=FakeHost2");
+        var connection3 = (NpgsqlRelationalConnection)context3.GetService<IRelationalConnection>();
+        Assert.Equal("Host=FakeHost2", connection3.ConnectionString);
+        Assert.Null(connection3.DbDataSource);
     }
 
     private class ConnectionStringSwitchingContext(string connectionString, bool withNetTopologySuite = false, bool withEnum = false)