From 448bb8e0237eb2eacdb9c2e8aa0be242af7cbed6 Mon Sep 17 00:00:00 2001 From: Juan Martinez Ramirez Date: Tue, 3 Sep 2024 15:33:55 -0600 Subject: [PATCH] Applying PR suggestions --- .../UnitTests/ConnectionPoolManagerMFATest.cs | 4 +-- Snowflake.Data/Core/Session/SessionPool.cs | 31 +++++++++++-------- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/Snowflake.Data.Tests/UnitTests/ConnectionPoolManagerMFATest.cs b/Snowflake.Data.Tests/UnitTests/ConnectionPoolManagerMFATest.cs index 139c2a4f1..0fc73b25b 100644 --- a/Snowflake.Data.Tests/UnitTests/ConnectionPoolManagerMFATest.cs +++ b/Snowflake.Data.Tests/UnitTests/ConnectionPoolManagerMFATest.cs @@ -102,7 +102,7 @@ public void TestPoolManagerShouldOnlyUsePasscodeAsArgumentForFirstSessionWhenNot }); // Act var session = _connectionPoolManager.GetSession(ConnectionStringMFABasicWithoutPasscode, null, SecureStringHelper.Encode(TestPasscode)); - Thread.Sleep(3000); + Thread.Sleep(5000); // Assert @@ -119,7 +119,7 @@ public void TestPoolManagerShouldThrowExceptionIfForcePoolingWithPasscodeNotUsin var connectionString = "db=D1;warehouse=W1;account=A1;user=U1;password=P1;role=R1;minPoolSize=2;passcode=12345;POOLINGENABLED=true"; // Act and assert var thrown = Assert.Throws(() =>_connectionPoolManager.GetSession(connectionString, null)); - Assert.That(thrown.Message, Does.Contain("Could not use connection pool because passcode was provided using a different authenticator than username_password_mfa")); + Assert.That(thrown.Message, Does.Contain("Passcode with MinPoolSize feature of connection pool allowed only for username_password_mfa authentication")); } [Test] diff --git a/Snowflake.Data/Core/Session/SessionPool.cs b/Snowflake.Data/Core/Session/SessionPool.cs index 8d314d235..f4d514ea8 100644 --- a/Snowflake.Data/Core/Session/SessionPool.cs +++ b/Snowflake.Data/Core/Session/SessionPool.cs @@ -139,7 +139,7 @@ internal SFSession GetSession(string connStr, SecureString password, SecureStrin { s_logger.Debug("SessionPool::GetSession" + PoolIdentification()); var sessionProperties = SFSessionProperties.ParseConnectionString(connStr, password); - ValidatePoolingIfPasscodeProvided(sessionProperties); + ValidateMinPoolSizeWithPasscode(sessionProperties); if (!GetPooling()) return NewNonPoolingSession(connStr, password, passcode); var isMfaAuthentication = sessionProperties.TryGetValue(SFSessionProperty.AUTHENTICATOR, out var authenticator) && authenticator == MFACacheAuthenticator.AUTH_NAME; @@ -158,7 +158,7 @@ internal SFSession GetSession(string connStr, SecureString password, SecureStrin return session; } - private void ValidatePoolingIfPasscodeProvided(SFSessionProperties sessionProperties) + private void ValidateMinPoolSizeWithPasscode(SFSessionProperties sessionProperties) { if (!GetPooling() || !IsMultiplePoolsVersion() || _poolConfig.MinPoolSize == 0) return; var isUsingPasscode = (sessionProperties.IsNonEmptyValueProvided(SFSessionProperty.PASSCODE) || @@ -168,7 +168,7 @@ private void ValidatePoolingIfPasscodeProvided(SFSessionProperties sessionProper authenticator == MFACacheAuthenticator.AUTH_NAME; if(isUsingPasscode && !isMfaAuthenticator) { - const string ErrorMessage = "Could not use connection pool because passcode was provided using a different authenticator than username_password_mfa"; + const string ErrorMessage = "Passcode with MinPoolSize feature of connection pool allowed only for username_password_mfa authentication"; s_logger.Error(ErrorMessage + PoolIdentification()); throw new Exception(ErrorMessage); } @@ -177,24 +177,29 @@ private void ValidatePoolingIfPasscodeProvided(SFSessionProperties sessionProper internal async Task GetSessionAsync(string connStr, SecureString password, SecureString passcode, CancellationToken cancellationToken) { s_logger.Debug("SessionPool::GetSessionAsync" + PoolIdentification()); - SFSession session = null; var sessionProperties = SFSessionProperties.ParseConnectionString(connStr, password); - ValidatePoolingIfPasscodeProvided(sessionProperties); + ValidateMinPoolSizeWithPasscode(sessionProperties); if (!GetPooling()) return await NewNonPoolingSessionAsync(connStr, password, passcode, cancellationToken).ConfigureAwait(false); - var sessionOrCreateTokens = GetIdleSession(connStr); + var isMfaAuthentication = sessionProperties.TryGetValue(SFSessionProperty.AUTHENTICATOR, out var authenticator) && authenticator == MFACacheAuthenticator.AUTH_NAME; + var sessionOrCreateTokens = GetIdleSession(connStr, isMfaAuthentication ? 1 : Int32.MaxValue); WarnAboutOverridenConfig(); - if (sessionProperties.TryGetValue(SFSessionProperty.AUTHENTICATOR, out var authenticator) && - authenticator == MFACacheAuthenticator.AUTH_NAME) - session = sessionOrCreateTokens.Session ?? - await NewSessionAsync(connStr, password, passcode, sessionOrCreateTokens.SessionCreationToken(), cancellationToken) - .ConfigureAwait(false); + if (sessionOrCreateTokens.Session != null) { _sessionPoolEventHandler.OnSessionProvided(this); } - ScheduleNewIdleSessions(connStr, password, RegisterSessionCreationsToEnsureMinPoolSize()); - return session ?? sessionOrCreateTokens.Session ?? await NewSessionAsync(connStr, password, passcode, sessionOrCreateTokens.SessionCreationToken(), cancellationToken).ConfigureAwait(false); + ScheduleNewIdleSessions(connStr, password, sessionOrCreateTokens.BackgroundSessionCreationTokens()); + WarnAboutOverridenConfig(); + var session = sessionOrCreateTokens.Session ?? + await NewSessionAsync(connStr, password, passcode, sessionOrCreateTokens.SessionCreationToken(), cancellationToken) + .ConfigureAwait(false); + if (isMfaAuthentication) + { + ScheduleNewIdleSessions(connStr, password, RegisterSessionCreationsToEnsureMinPoolSize()); + } + return session; + } private void ScheduleNewIdleSessions(string connStr, SecureString password, List tokens)