1111
1212namespace Npgsql ;
1313
14- sealed class MultiHostDataSource : NpgsqlDataSource
14+ #pragma warning disable CS1591
15+ #pragma warning disable RS0016
16+
17+ public sealed class NpgsqlMultiHostDataSource : NpgsqlDataSource
1518{
1619 internal override bool OwnsConnectors => false ;
1720
1821 readonly NpgsqlDataSource [ ] _pools ;
1922
2023 internal NpgsqlDataSource [ ] Pools => _pools ;
2124
25+ readonly MultiHostDataSourceWrapper [ ] _wrappers ;
26+
2227 volatile int _roundRobinIndex = - 1 ;
2328
24- public MultiHostDataSource ( NpgsqlConnectionStringBuilder settings , string connString , NpgsqlDataSourceConfiguration dataSourceConfig )
25- : base ( settings , connString , dataSourceConfig )
29+ internal NpgsqlMultiHostDataSource ( NpgsqlConnectionStringBuilder settings , NpgsqlDataSourceConfiguration dataSourceConfig )
30+ : base ( settings , dataSourceConfig )
2631 {
2732 var hosts = settings . Host ! . Split ( ',' ) ;
2833 _pools = new NpgsqlDataSource [ hosts . Length ] ;
@@ -40,11 +45,77 @@ public MultiHostDataSource(NpgsqlConnectionStringBuilder settings, string connSt
4045 poolSettings . Host = host . ToString ( ) ;
4146
4247 _pools [ i ] = settings . Pooling
43- ? new PoolingDataSource ( poolSettings , poolSettings . ConnectionString , dataSourceConfig , this )
44- : new UnpooledDataSource ( poolSettings , poolSettings . ConnectionString , dataSourceConfig ) ;
48+ ? new PoolingDataSource ( poolSettings , dataSourceConfig , this )
49+ : new UnpooledDataSource ( poolSettings , dataSourceConfig ) ;
50+ }
51+
52+ var targetSessionAttributeValues = Enum . GetValues ( typeof ( TargetSessionAttributes ) ) . Cast < TargetSessionAttributes > ( ) . ToArray ( ) ;
53+ _wrappers = new MultiHostDataSourceWrapper [ targetSessionAttributeValues . Max ( t => ( int ) t ) + 1 ] ;
54+ foreach ( var targetSessionAttribute in targetSessionAttributeValues )
55+ {
56+ _wrappers [ ( int ) targetSessionAttribute ] = new ( this , targetSessionAttribute ) ;
57+ }
58+ }
59+
60+ /// <summary>
61+ /// Returns a new, unopened connection from this data source.
62+ /// </summary>
63+ /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
64+ public NpgsqlConnection CreateConnection ( TargetSessionAttributes targetSessionAttributes )
65+ => NpgsqlConnection . FromDataSource ( _wrappers [ ( int ) targetSessionAttributes ] ) ;
66+
67+ /// <summary>
68+ /// Returns a new, opened connection from this data source.
69+ /// </summary>
70+ /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
71+ public NpgsqlConnection OpenConnection ( TargetSessionAttributes targetSessionAttributes )
72+ {
73+ var connection = CreateConnection ( targetSessionAttributes ) ;
74+
75+ try
76+ {
77+ connection . Open ( ) ;
78+ return connection ;
79+ }
80+ catch
81+ {
82+ connection . Dispose ( ) ;
83+ throw ;
4584 }
4685 }
4786
87+ /// <summary>
88+ /// Returns a new, opened connection from this data source.
89+ /// </summary>
90+ /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
91+ /// <param name="cancellationToken">
92+ /// An optional token to cancel the asynchronous operation. The default value is <see cref="CancellationToken.None"/>.
93+ /// </param>
94+ public async ValueTask < NpgsqlConnection > OpenConnectionAsync (
95+ TargetSessionAttributes targetSessionAttributes ,
96+ CancellationToken cancellationToken = default )
97+ {
98+ var connection = CreateConnection ( targetSessionAttributes ) ;
99+
100+ try
101+ {
102+ await connection . OpenAsync ( cancellationToken ) . ConfigureAwait ( false ) ;
103+ return connection ;
104+ }
105+ catch
106+ {
107+ await connection . DisposeAsync ( ) . ConfigureAwait ( false ) ;
108+ throw ;
109+ }
110+ }
111+
112+ /// <summary>
113+ /// Returns an <see cref="NpgsqlDataSource" /> that wraps this multi-host one with the given server type.
114+ /// </summary>
115+ /// <param name="targetSessionAttributes">Specifies the server type (e.g. primary, standby).</param>
116+ public NpgsqlDataSource For ( TargetSessionAttributes targetSessionAttributes )
117+ => _wrappers [ ( int ) targetSessionAttributes ] ;
118+
48119 static bool IsPreferred ( ClusterState state , TargetSessionAttributes preferredType )
49120 => state switch
50121 {
@@ -252,13 +323,13 @@ int GetRoundRobinIndex()
252323 }
253324
254325 internal override void Return ( NpgsqlConnector connector )
255- => throw new NpgsqlException ( "Npgsql bug: a connector was returned to " + nameof ( MultiHostDataSource ) ) ;
326+ => throw new NpgsqlException ( "Npgsql bug: a connector was returned to " + nameof ( NpgsqlMultiHostDataSource ) ) ;
256327
257328 internal override bool TryGetIdleConnector ( [ NotNullWhen ( true ) ] out NpgsqlConnector ? connector )
258- => throw new NpgsqlException ( "Npgsql bug: trying to get an idle connector from " + nameof ( MultiHostDataSource ) ) ;
329+ => throw new NpgsqlException ( "Npgsql bug: trying to get an idle connector from " + nameof ( NpgsqlMultiHostDataSource ) ) ;
259330
260331 internal override ValueTask < NpgsqlConnector ? > OpenNewConnector ( NpgsqlConnection conn , NpgsqlTimeout timeout , bool async , CancellationToken cancellationToken )
261- => throw new NpgsqlException ( "Npgsql bug: trying to open a new connector from " + nameof ( MultiHostDataSource ) ) ;
332+ => throw new NpgsqlException ( "Npgsql bug: trying to open a new connector from " + nameof ( NpgsqlMultiHostDataSource ) ) ;
262333
263334 internal override void Clear ( )
264335 {
@@ -319,5 +390,5 @@ static TargetSessionAttributes GetTargetSessionAttributes(NpgsqlConnection conne
319390 => connection . Settings . TargetSessionAttributesParsed ??
320391 ( PostgresEnvironment . TargetSessionAttributes is string s
321392 ? NpgsqlConnectionStringBuilder . ParseTargetSessionAttributes ( s )
322- : default ) ;
393+ : TargetSessionAttributes . Any ) ;
323394}
0 commit comments