npgsql · roji · Jul 4, 2022 · Jul 4, 2022 · NinoFloris · Jul 4, 2022
@@ -88,6 +88,7 @@
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=Datemultirange/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=daterange/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=DDEX/@EntryIndexedValue">True</s:Boolean>
+	<s:Boolean x:Key="/Default/UserDictionary/Words/=failover/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=IANA/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=lquery/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=lseg/@EntryIndexedValue">True</s:Boolean>

@@ -444,7 +444,7 @@ class AuthenticationCompleteException : Exception { }
 
     async ValueTask<string?> GetPassword(string username, bool async, CancellationToken cancellationToken = default)
     {
-        var password = await _dataSource.GetPasswordAsync(async, cancellationToken);
+        var password = await _dataSource.GetPassword(async, cancellationToken);
 
         if (password is not null)
             return password;

@@ -237,7 +237,7 @@ internal void FlagAsWritableForMultiplexing()
     /// <summary>
     /// The connector source (e.g. pool) from where this connector came, and to which it will be returned.
     /// Note that in multi-host scenarios, this references the host-specific <see cref="PoolingDataSource"/> rather than the
-    /// <see cref="MultiHostDataSource"/>,
+    /// <see cref="NpgsqlMultiHostDataSource"/>,
     /// </summary>
     readonly NpgsqlDataSource _dataSource;
 

@@ -11,12 +11,21 @@ sealed class MultiHostDataSourceWrapper : NpgsqlDataSource
 {
     internal override bool OwnsConnectors => false;
 
-    readonly MultiHostDataSource _wrappedSource;
+    readonly NpgsqlMultiHostDataSource _wrappedSource;
 
-    public MultiHostDataSourceWrapper(NpgsqlConnectionStringBuilder settings, string connString, MultiHostDataSource source)
-        : base(settings, connString, source.Configuration)
+    public MultiHostDataSourceWrapper(NpgsqlMultiHostDataSource source, TargetSessionAttributes targetSessionAttributes)
+        : base(CloneSettingsForTargetSessionAttributes(source.Settings, targetSessionAttributes), source.Configuration)
         => _wrappedSource = source;
 
+    static NpgsqlConnectionStringBuilder CloneSettingsForTargetSessionAttributes(
+        NpgsqlConnectionStringBuilder settings,
+        TargetSessionAttributes targetSessionAttributes)
+    {
+        var clonedSettings = settings.Clone();
+        clonedSettings.TargetSessionAttributesParsed = targetSessionAttributes;
+        return clonedSettings;
+    }
+
     internal override (int Total, int Idle, int Busy) Statistics => _wrappedSource.Statistics;
 
     internal override void Clear() => _wrappedSource.Clear();

@@ -49,10 +49,9 @@ public bool IsBootstrapped
 
     internal MultiplexingDataSource(
         NpgsqlConnectionStringBuilder settings,
-        string connString,
         NpgsqlDataSourceConfiguration dataSourceConfig,
-        MultiHostDataSource? parentPool = null)
-        : base(settings, connString, dataSourceConfig, parentPool)
+        NpgsqlMultiHostDataSource? parentPool = null)
+        : base(settings, dataSourceConfig, parentPool)
     {
         Debug.Assert(Settings.Multiplexing);
 

@@ -199,15 +199,16 @@ void SetupDataSource()
         // The connection string may be equivalent to one that has already been seen though (e.g. different
         // ordering). Have NpgsqlConnectionStringBuilder produce a canonical string representation
         // and recheck.
-        // Note that we remove TargetSessionAttributes and LoadBalanceHosts to make all connection strings
-        // that are otherwise identical point to the same pool.
+        // Note that we remove TargetSessionAttributes to make all connection strings that are otherwise identical point to the same pool.
         var canonical = settings.ConnectionStringForMultipleHosts;
 
         if (PoolManager.Pools.TryGetValue(canonical, out _dataSource))
         {
-            // We're wrapping the original pool in the other, as the original doesn't have the TargetSessionAttributes
-            if (_dataSource is MultiHostDataSource mhdsCanonical)
-                _dataSource = new MultiHostDataSourceWrapper(Settings, _connectionString, mhdsCanonical);
+            // If this is a multi-host data source and the user specified a TargetSessionAttributes, create a wrapper in front of the
+            // MultiHostDataSource with that TargetSessionAttributes.
+            if (_dataSource is NpgsqlMultiHostDataSource multiHostDataSource && settings.TargetSessionAttributesParsed.HasValue)
+                _dataSource = multiHostDataSource.For(settings.TargetSessionAttributesParsed.Value);
+
             // The pool was found, but only under the canonical key - we're using a different version
             // for the first time. Map it via our own key for next time.
             _dataSource = PoolManager.Pools.GetOrAdd(_connectionString, _dataSource);
@@ -228,7 +229,7 @@ void SetupDataSource()
             Debug.Assert(_dataSource is not MultiHostDataSourceWrapper);
             // If the pool we created was the one that ended up being stored we need to increment the appropriate counter.
             // Avoids a race condition where multiple threads will create a pool but only one will be stored.
-            if (_dataSource is MultiHostDataSource multiHostConnectorPool)
+            if (_dataSource is NpgsqlMultiHostDataSource multiHostConnectorPool)
                 foreach (var hostPool in multiHostConnectorPool.Pools)
                     NpgsqlEventSource.Log.DataSourceCreated(hostPool);
             else
@@ -237,9 +238,10 @@ void SetupDataSource()
         else
             newDataSource.Dispose();
 
-        // We're wrapping the original pool in the other, as the original doesn't have the TargetSessionAttributes
-        if (_dataSource is MultiHostDataSource mhds)
-            _dataSource = new MultiHostDataSourceWrapper(Settings, _connectionString, mhds);
+        // If this is a multi-host data source and the user specified a TargetSessionAttributes, create a wrapper in front of the
+        // MultiHostDataSource with that TargetSessionAttributes.
+        if (_dataSource is NpgsqlMultiHostDataSource multiHostDataSource2 && settings.TargetSessionAttributesParsed.HasValue)
+            _dataSource = multiHostDataSource2.For(settings.TargetSessionAttributesParsed.Value);
 
         _dataSource = PoolManager.Pools.GetOrAdd(_connectionString, _dataSource);
     }

@@ -938,7 +938,7 @@ public string? TargetSessionAttributes
         }
     }
 
-    internal TargetSessionAttributes? TargetSessionAttributesParsed { get; private set; }
+    internal TargetSessionAttributes? TargetSessionAttributesParsed { get; set; }
 
     internal static TargetSessionAttributes ParseTargetSessionAttributes(string s)
         => s switch
@@ -1865,46 +1865,4 @@ enum ReplicationMode
     Logical
 }
 
-/// <summary>
-/// Specifies server type preference.
-/// </summary>
-enum TargetSessionAttributes : byte
-{
-    /// <summary>
-    /// Any successful connection is acceptable.
-    /// </summary>
-    Any = 0,
-
-    /// <summary>
-    /// Session must accept read-write transactions by default (that is, the server must not be in hot standby mode and the
-    /// <c>default_transaction_read_only</c> parameter must be off).
-    /// </summary>
-    ReadWrite = 1,
-
-    /// <summary>
-    /// Session must not accept read-write transactions by default (the converse).
-    /// </summary>
-    ReadOnly = 2,
-
-    /// <summary>
-    /// Server must not be in hot standby mode.
-    /// </summary>
-    Primary = 3,
-
-    /// <summary>
-    /// Server must be in hot standby mode.
-    /// </summary>
-    Standby = 4,
-
-    /// <summary>
-    /// First try to find a primary server, but if none of the listed hosts is a primary server, try again in <see cref="Any"/> mode.
-    /// </summary>
-    PreferPrimary = 5,
-
-    /// <summary>
-    /// First try to find a standby server, but if none of the listed hosts is a standby server, try again in <see cref="Any"/> mode.
-    /// </summary>
-    PreferStandby = 6,
-}
-
 #endregion
@@ -50,12 +50,11 @@ private protected readonly Dictionary<Transaction, List<NpgsqlConnector>> _pendi
 
     internal NpgsqlDataSource(
         NpgsqlConnectionStringBuilder settings,
-        string connectionString,
         NpgsqlDataSourceConfiguration dataSourceConfig)
     {
         Settings = settings;
         ConnectionString = settings.PersistSecurityInfo
-            ? connectionString
+            ? settings.ToString()
             : settings.ToStringWithoutPassword();
 
         Configuration = dataSourceConfig;
@@ -179,11 +178,11 @@ public string Password
         }
     }
 
-    internal async ValueTask<string?> GetPasswordAsync(bool async, CancellationToken cancellationToken = default)
+    internal async ValueTask<string?> GetPassword(bool async, CancellationToken cancellationToken = default)
     {
         // A periodic password provider is configured, but the first refresh hasn't completed yet (race condition).
         // Wait until it completes.
-        if (_password is null && _passwordRefreshTask is not null)
+        if (_password is null && _periodicPasswordProvider is not null)
         {
             if (async)
                 await _passwordRefreshTask;

@@ -101,8 +101,6 @@ public NpgsqlDataSourceBuilder UsePeriodicPasswordProvider(
     /// </summary>
     public NpgsqlDataSource Build()
     {
-        var connectionString = ConnectionStringBuilder.ToString();
-
         ConnectionStringBuilder.PostProcessAndValidate();
 
         if (_periodicPasswordProvider is not null &&
@@ -123,17 +121,28 @@ public NpgsqlDataSource Build()
 
         if (ConnectionStringBuilder.Host!.Contains(","))
         {
+            if (ConnectionStringBuilder.TargetSessionAttributes is not null)
+                throw new InvalidOperationException(NpgsqlStrings.CannotSpecifyTargetSessionAttributes);
             if (ConnectionStringBuilder.Multiplexing)
                 throw new NotSupportedException("Multiplexing is not supported with multiple hosts");
             if (ConnectionStringBuilder.ReplicationMode != ReplicationMode.Off)
                 throw new NotSupportedException("Replication is not supported with multiple hosts");
-            return new MultiHostDataSource(ConnectionStringBuilder, connectionString, config);
+
+            return new NpgsqlMultiHostDataSource(ConnectionStringBuilder, config);
         }
 
         return ConnectionStringBuilder.Multiplexing
-            ? new MultiplexingDataSource(ConnectionStringBuilder, connectionString, config)
+            ? new MultiplexingDataSource(ConnectionStringBuilder, config)
             : ConnectionStringBuilder.Pooling
-                ? new PoolingDataSource(ConnectionStringBuilder, connectionString, config)
-                : new UnpooledDataSource(ConnectionStringBuilder, connectionString, config);
+                ? new PoolingDataSource(ConnectionStringBuilder, config)
+                : new UnpooledDataSource(ConnectionStringBuilder, config);
     }
+
+#pragma warning disable RS0016
+    /// <summary>
+    /// Builds and returns a <see cref="NpgsqlMultiHostDataSource" /> which is ready for use for load-balancing and failover scenarios.
+    /// </summary>
+    public NpgsqlMultiHostDataSource BuildMultiHost()
+        => Build() as NpgsqlMultiHostDataSource ?? throw new InvalidOperationException(NpgsqlStrings.MultipleHostsMustBeSpecified);
+#pragma warning restore RS0016
 }
@@ -11,18 +11,23 @@
 
 namespace Npgsql;
 
-sealed class MultiHostDataSource : NpgsqlDataSource
+#pragma warning disable CS1591
+#pragma warning disable RS0016
+
+public sealed class NpgsqlMultiHostDataSource : NpgsqlDataSource
 {
     internal override bool OwnsConnectors => false;
 
     readonly NpgsqlDataSource[] _pools;
 
     internal NpgsqlDataSource[] Pools => _pools;
 
+    readonly MultiHostDataSourceWrapper[] _wrappers;
+
     volatile int _roundRobinIndex = -1;
 
-    public MultiHostDataSource(NpgsqlConnectionStringBuilder settings, string connString, NpgsqlDataSourceConfiguration dataSourceConfig)
-        : base(settings, connString, dataSourceConfig)
+    internal NpgsqlMultiHostDataSource(NpgsqlConnectionStringBuilder settings, NpgsqlDataSourceConfiguration dataSourceConfig)
+        : base(settings, dataSourceConfig)
     {
         var hosts = settings.Host!.Split(',');
         _pools = new NpgsqlDataSource[hosts.Length];
@@ -40,11 +45,77 @@ public MultiHostDataSource(NpgsqlConnectionStringBuilder settings, string connSt
                 poolSettings.Host = host.ToString();
 
             _pools[i] = settings.Pooling
-                ? new PoolingDataSource(poolSettings, poolSettings.ConnectionString, dataSourceConfig, this)
-                : new UnpooledDataSource(poolSettings, poolSettings.ConnectionString, dataSourceConfig);
+                ? new PoolingDataSource(poolSettings, dataSourceConfig, this)
+                : new UnpooledDataSource(poolSettings, dataSourceConfig);
+        }
+
+        var targetSessionAttributeValues = Enum.GetValues(typeof(TargetSessionAttributes)).Cast<TargetSessionAttributes>().ToArray();
+        _wrappers = new MultiHostDataSourceWrapper[targetSessionAttributeValues.Max(t => (int)t) + 1];
+        foreach (var targetSessionAttribute in targetSessionAttributeValues)
+        {
+            _wrappers[(int)targetSessionAttribute] = new(this, targetSessionAttribute);
+        }
+    }
+
+    /// <summary>
+    /// Returns a new, unopened connection from this data source.
+    /// </summary>
+    /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
+    public NpgsqlConnection CreateConnection(TargetSessionAttributes targetSessionAttributes)
+        => NpgsqlConnection.FromDataSource(_wrappers[(int)targetSessionAttributes]);
+
+    /// <summary>
+    /// Returns a new, opened connection from this data source.
+    /// </summary>
+    /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
+    public NpgsqlConnection OpenConnection(TargetSessionAttributes targetSessionAttributes)
+    {
+        var connection = CreateConnection(targetSessionAttributes);
+
+        try
+        {
+            connection.Open();
+            return connection;
+        }
+        catch
+        {
+            connection.Dispose();
+            throw;
         }
     }
 
+    /// <summary>
+    /// Returns a new, opened connection from this data source.
+    /// </summary>
+    /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
+    /// <param name="cancellationToken">
+    /// An optional token to cancel the asynchronous operation. The default value is <see cref="CancellationToken.None"/>.
+    /// </param>
+    public async ValueTask<NpgsqlConnection> OpenConnectionAsync(
+        TargetSessionAttributes targetSessionAttributes,
+        CancellationToken cancellationToken = default)
+    {
+        var connection = CreateConnection(targetSessionAttributes);
+
+        try
+        {
+            await connection.OpenAsync(cancellationToken).ConfigureAwait(false);
+            return connection;
+        }
+        catch
+        {
+            await connection.DisposeAsync().ConfigureAwait(false);
+            throw;
+        }
+    }
+
+    /// <summary>
+    /// Returns an <see cref="NpgsqlDataSource" /> that wraps this multi-host one with the given server type.
+    /// </summary>
+    /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
+    public NpgsqlDataSource For(TargetSessionAttributes targetSessionAttributes)
+        => _wrappers[(int)targetSessionAttributes];
+
     static bool IsPreferred(ClusterState state, TargetSessionAttributes preferredType)
         => state switch
         {
@@ -252,13 +323,13 @@ int GetRoundRobinIndex()
     }
 
     internal override void Return(NpgsqlConnector connector)
-        => throw new NpgsqlException("Npgsql bug: a connector was returned to " + nameof(MultiHostDataSource));
+        => throw new NpgsqlException("Npgsql bug: a connector was returned to " + nameof(NpgsqlMultiHostDataSource));
 
     internal override bool TryGetIdleConnector([NotNullWhen(true)] out NpgsqlConnector? connector)
-        => throw new NpgsqlException("Npgsql bug: trying to get an idle connector from " + nameof(MultiHostDataSource));
+        => throw new NpgsqlException("Npgsql bug: trying to get an idle connector from " + nameof(NpgsqlMultiHostDataSource));
 
     internal override ValueTask<NpgsqlConnector?> OpenNewConnector(NpgsqlConnection conn, NpgsqlTimeout timeout, bool async, CancellationToken cancellationToken)
-        => throw new NpgsqlException("Npgsql bug: trying to open a new connector from " + nameof(MultiHostDataSource));
+        => throw new NpgsqlException("Npgsql bug: trying to open a new connector from " + nameof(NpgsqlMultiHostDataSource));
 
     internal override void Clear()
     {
@@ -319,5 +390,5 @@ static TargetSessionAttributes GetTargetSessionAttributes(NpgsqlConnection conne
         => connection.Settings.TargetSessionAttributesParsed ??
            (PostgresEnvironment.TargetSessionAttributes is string s
                ? NpgsqlConnectionStringBuilder.ParseTargetSessionAttributes(s)
-               : default);
+               : TargetSessionAttributes.Any);
 }
@@ -30,7 +30,7 @@ class PoolingDataSource : NpgsqlDataSource
     /// </summary>
     private protected readonly NpgsqlConnector?[] Connectors;
 
-    readonly MultiHostDataSource? _parentPool;
+    readonly NpgsqlMultiHostDataSource? _parentPool;
 
     /// <summary>
     /// Reader side for the idle connector channel. Contains nulls in order to release waiting attempts after
@@ -77,10 +77,9 @@ internal sealed override (int Total, int Idle, int Busy) Statistics
 
     internal PoolingDataSource(
         NpgsqlConnectionStringBuilder settings,
-        string connString,
         NpgsqlDataSourceConfiguration dataSourceConfig,
-        MultiHostDataSource? parentPool = null)
-        : base(settings, connString, dataSourceConfig)
+        NpgsqlMultiHostDataSource? parentPool = null)
+        : base(settings, dataSourceConfig)
     {
         if (settings.MaxPoolSize < settings.MinPoolSize)
             throw new ArgumentException($"Connection can't have 'Max Pool Size' {settings.MaxPoolSize} under 'Min Pool Size' {settings.MinPoolSize}");