Skip to content

Commit

Permalink
Applying PR suggestions
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-jmartinezramirez
sfc-gh-jmartinezramirez committed Sep 18, 2024
1 parent 188688f commit 448bb8e
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 15 deletions.
4 changes: 2 additions & 2 deletions Snowflake.Data.Tests/UnitTests/ConnectionPoolManagerMFATest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ public void TestPoolManagerShouldOnlyUsePasscodeAsArgumentForFirstSessionWhenNot
});
// Act
var session = _connectionPoolManager.GetSession(ConnectionStringMFABasicWithoutPasscode, null, SecureStringHelper.Encode(TestPasscode));
Thread.Sleep(3000);
Thread.Sleep(5000);

// Assert

Expand All @@ -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<Exception>(() =>_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]
Expand Down
31 changes: 18 additions & 13 deletions Snowflake.Data/Core/Session/SessionPool.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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) ||
Expand All @@ -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);
}
Expand All @@ -177,24 +177,29 @@ private void ValidatePoolingIfPasscodeProvided(SFSessionProperties sessionProper
internal async Task<SFSession> 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<SessionCreationToken> tokens)
Expand Down

0 comments on commit 448bb8e

Please sign in to comment.