Accept pooling for key files with password

snowflakedb · May 15, 2024 · edb50c2 · edb50c2
1 parent db479e6
commit edb50c2
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 5 deletions.
diff --git a/Snowflake.Data.Tests/UnitTests/ConnectionPoolManagerTest.cs b/Snowflake.Data.Tests/UnitTests/ConnectionPoolManagerTest.cs
@@ -211,6 +211,18 @@ public void TestGetPoolingOnManagerLevelAlwaysTrue()
             Assert.IsFalse(sessionPool2.GetPooling());
         }
 
+        [Test]
+        [TestCase("authenticator=externalbrowser;account=test;user=test;")]
+        [TestCase("authenticator=snowflake_jwt;account=test;user=test;private_key_file=/some/file.key")]
+        public void TestDisabledPoolingWhenSecretesProvidedExternally(string connectionString)
+        {
+            // act
+            var pool = _connectionPoolManager.GetPool(connectionString, null);
+
+            // assert
+            Assert.IsFalse(pool.GetPooling());
+        }
+
         [Test]
         public void TestGetTimeoutOnManagerLevelWhenNotAllPoolsEqual()
         {

diff --git a/Snowflake.Data.Tests/UnitTests/Session/ConnectionPoolConfigExtractorTest.cs b/Snowflake.Data.Tests/UnitTests/Session/ConnectionPoolConfigExtractorTest.cs
@@ -230,6 +230,8 @@ public void TestExtractPoolingEnabled(string propertyValue, bool poolingEnabled)
         [TestCase("authenticator=snowflake_jwt;account=test;user=test;private_key=secretKey", true)]
         [TestCase("authenticator=snowflake_jwt;account=test;user=test;private_key=secretKey;poolingEnabled=true;", true)]
         [TestCase("authenticator=snowflake_jwt;account=test;user=test;private_key=secretKey;poolingEnabled=false;", false)]
+        [TestCase("authenticator=snowflake_jwt;account=test;user=test;private_key_file=/some/file.key;private_key_pwd=secretPwd", true)]
+        [TestCase("authenticator=snowflake_jwt;account=test;user=test;private_key_file=/some/file.key;private_key_pwd=", false)]
         public void TestDisablePoolingDefaultWhenSecretsProvidedExternally(string connectionString, bool poolingEnabled)
         {
             // act

diff --git a/Snowflake.Data/Core/Session/SFSessionHttpClientProperties.cs b/Snowflake.Data/Core/Session/SFSessionHttpClientProperties.cs
@@ -55,8 +55,8 @@ public void DisablePoolingDefaultIfSecretsProvidedExternally(SFSessionProperties
                 DisablePoolingIfNotExplicitlyEnabled(properties, "external browser");
 
             } else if (KeyPairAuthenticator.AUTH_NAME.Equals(authenticator)
-                       && properties.TryGetValue(SFSessionProperty.PRIVATE_KEY_FILE, out var privateKeyFile)
-                       && !string.IsNullOrEmpty(privateKeyFile))
+                       && properties.IsNonEmptyValueProvided(SFSessionProperty.PRIVATE_KEY_FILE)
+                       && !properties.IsNonEmptyValueProvided(SFSessionProperty.PRIVATE_KEY_PWD))
             {
                 DisablePoolingIfNotExplicitlyEnabled(properties, "key pair with private key in a file");
             }

diff --git a/Snowflake.Data/Core/Session/SFSessionProperty.cs b/Snowflake.Data/Core/Session/SFSessionProperty.cs
@@ -256,7 +256,7 @@ internal static SFSessionProperties ParseConnectionString(string connectionStrin
             }
 
             ValidateAuthenticator(properties);
-            properties.IsPoolingEnabledValueProvided = CheckPoolingEnabledValueProvided(properties);
+            properties.IsPoolingEnabledValueProvided = properties.IsNonEmptyValueProvided(SFSessionProperty.POOLINGENABLED);
             CheckSessionProperties(properties);
             ValidateFileTransferMaxBytesInMemoryProperty(properties);
             ValidateAccountDomain(properties);
@@ -310,8 +310,8 @@ private static void ValidateAuthenticator(SFSessionProperties properties)
             }
         }
 
-        private static bool CheckPoolingEnabledValueProvided(SFSessionProperties properties) =>
-            properties.TryGetValue(SFSessionProperty.POOLINGENABLED, out var poolingEnabledStr) && !string.IsNullOrEmpty(poolingEnabledStr);
+        internal bool IsNonEmptyValueProvided(SFSessionProperty property) =>
+            TryGetValue(property, out var propertyValueStr) && !string.IsNullOrEmpty(propertyValueStr);
 
         private static string BuildConnectionStringWithoutSecrets(ref string[] keys, ref string[] values)
         {