enable-key-pair (#9)

Author: mcasara

CHANGELOG.md

@@ -1,3 +1,12 @@
+2.5.0 (2026-05-06)
+-------------------
+
+*Changes*
+- Update dependencies:
+  - snowflake-connector-python[pandas,secure-local-storage]
+- Add support for key pair authentication
+
+
 2.4.0 (2024-06-05)
 -------------------
 

setup.py

@@ -23,7 +23,7 @@
       python_requires='>=3.7',
       install_requires=[
           'pipelinewise-singer-python==1.*',
-          'snowflake-connector-python[pandas,secure-local-storage]==3.10.1',
+          'snowflake-connector-python[pandas,secure-local-storage]==4.4.0',
           'inflection==0.5.1',
           'joblib==1.2.0',
           'boto3==1.28.20',

target_snowflake/db_sync.py

@@ -44,17 +44,16 @@ def validate_config(config):
     # Use table stage if none s3_bucket and stage defined
     elif not config.get('s3_bucket', None) and not config.get('stage', None):
         required_config_keys = snowflake_required_config_keys
-    elif not(config.get("use_browser_authentication")) and not(config.get("password")):
-        errors.append("'password' configuration was not provided and it is mandatory when "
-                      "SSO browser authentication is not intended ("
-                      "'use_browser_authentication' is 'False'). Please provide a value "
-                      "for 'password' for basic authentication or "
-                      "set 'use_browser_authentication' to True for SSO browser authentication.")
+        if not(config.get("use_browser_authentication")) and not(config.get("password")) and not(config.get("private_key")):
+            errors.append("'password' or 'private_key' configuration was not provided and it is mandatory when "
+                          "SSO browser authentication is not intended ("
+                          "'use_browser_authentication' is 'False'). Please provide a value "
+                          "for 'password' or 'private_key' for basic authentication or "
+                          "set 'use_browser_authentication' to True for SSO browser authentication.")
     else:
         errors.append("Only one of 's3_bucket' or 'stage' keys defined in config. "
                       "Use both of them if you want to use an external stage when loading data into snowflake "
                       "or don't use any of them if you want ot use table stages.")
-
     # Check if mandatory keys exist
     for k in required_config_keys:
         if not config.get(k, None):
@@ -289,6 +288,39 @@ def __init__(self, connection_config, stream_schema_message=None, table_cache=No
         else:
             self.upload_client = SnowflakeUploadClient(connection_config, self)
 
+
+    def _configure_private_key_auth(self, private_key: str):
+        """Configure private key authentication.
+
+        Args:
+            private_key: RSA private key as raw PEM-encoded string
+
+        Raises:
+            ValueError: If private key is invalid or cannot be loaded
+        """
+        private_key_content = private_key.strip()
+
+        # Validate that the key is in PEM format
+        if not private_key_content.startswith("-----BEGIN"):
+            raise ValueError(
+                "Private key must be in raw PEM format (starting with '-----BEGIN').",
+            )
+
+        # Remove PEM headers/footers and whitespace to extract base64-encoded DER
+        # PEM format is just: -----BEGIN...-----\n<base64 DER>\n-----END...-----
+        key_base64 = re.sub(
+            r"-----BEGIN.*?-----\s*|\s*-----END.*?-----\s*|\s+",
+            "",
+            private_key_content,
+        )
+
+        if not key_base64:
+            raise ValueError(
+                "Private key appears to be empty after removing PEM headers.",
+            )
+
+        return key_base64
+
     def open_connection(self):
         """Open snowflake connection"""
         stream = None
@@ -298,6 +330,7 @@ def open_connection(self):
         return snowflake.connector.connect(
             user=self.connection_config['user'],
             password=self.connection_config.get('password') if not self.connection_config.get('use_browser_authentication') else None,
+            private_key=self._configure_private_key_auth(self.connection_config.get('private_key')) if not self.connection_config.get('use_browser_authentication') else None,
             account=self.connection_config['account'],
             database=self.connection_config['dbname'],
             warehouse=self.connection_config['warehouse'],

tests/unit/test_db_sync.py

@@ -42,6 +42,7 @@ def test_config_validation(self):
             'dbname': "dummy-value",
             'user': "dummy-value",
             'password': "dummy-value",
+            'private_key': "dummy-value",
             'warehouse': "dummy-value",
             'default_target_schema': "dummy-value",
             'file_format': "dummy-value"
@@ -92,15 +93,18 @@ def test_config_validation(self):
         config_with_archive_load_files['archive_load_files'] = True
         self.assertGreater(len(validator(config_with_external_stage)), 0)
 
-        # Configuration without password nor use_browser_authentication=True
-        config_without_password_nor_browser_auth = minimal_config.copy()
-        config_without_password_nor_browser_auth.pop("password")
-        self.assertGreater(len(validator(minimal_config)), 0)
+        # Configuration without password nor private_key nor use_browser_authentication=True
+        config_without_password_nor_private_key_nor_browser_auth = minimal_config.copy()
+        config_without_password_nor_private_key_nor_browser_auth.pop("password")
+        config_without_password_nor_private_key_nor_browser_auth.pop("private_key")
+        self.assertGreater(len(validator(config_without_password_nor_private_key_nor_browser_auth)), 0)
 
-        # Configuration without password nor use_browser_authentication=True
-        config_without_password_with_browser_auth = minimal_config.copy()
-        config_without_password_with_browser_auth["use_browser_authentication"] = True
-        self.assertEqual(len(validator(minimal_config)), 0)
+        # Configuration without password nor private_key and use_browser_authentication=True
+        config_without_password_nor_private_key_nor_browser_auth = minimal_config.copy()
+        config_without_password_nor_private_key_nor_browser_auth.pop("password")
+        config_without_password_nor_private_key_nor_browser_auth.pop("private_key")
+        config_without_password_nor_private_key_nor_browser_auth["use_browser_authentication"] = True
+        self.assertEqual(len(validator(config_without_password_nor_private_key_nor_browser_auth)), 0)
 
     def test_column_type_mapping(self):
         """Test JSON type to Snowflake column type mappings"""
@@ -218,6 +222,7 @@ def test_parallelism(self, query_patch):
             'dbname': "dummy-value",
             'user': "dummy-value",
             'password': "dummy-value",
+            'private_key': "dummy-value",
             'warehouse': "dummy-value",
             'default_target_schema': "dummy-value",
             'file_format': "dummy-value"
@@ -291,6 +296,7 @@ def test_record_primary_key_string(self, query_patch):
             'dbname': "dummy-value",
             'user': "dummy-value",
             'password': "dummy-value",
+            'private_key': "dummy-value",
             'warehouse': "dummy-value",
             'default_target_schema': "dummy-value",
             'file_format': "dummy-value"
@@ -350,6 +356,7 @@ def test_merge_failure_message(self, load_file_merge_patch, query_patch):
             'dbname': "dummy_dbname",
             'user': "dummy_user",
             'password': "dummy_password",
+            'private_key': "dummy_private_key",
             'warehouse': "dummy_warehouse",
             'default_target_schema': "dummy_default_target_schema",
             'file_format': "dummy_file_format",
@@ -388,6 +395,7 @@ def test_copy_failure_message(self, load_file_copy_patch, query_patch):
             'dbname': "dummy_dbname",
             'user': "dummy_user",
             'password': "dummy_password",
+            'private_key': "dummy_private_key",
             'warehouse': "dummy_warehouse",
             'default_target_schema': "dummy_default_target_schema",
             'file_format': "dummy_file_format",
@@ -423,6 +431,7 @@ def test_sync_table_with_no_changes_to_pk(self, query_patch):
             'dbname': "dummy-db",
             'user': "dummy-user",
             'password': "dummy-passwd",
+            'private_key': "dummy-private_key",
             'warehouse': "dummy-wh",
             'default_target_schema': "dummy-schema",
             'file_format': "dummy-file-format"
@@ -471,6 +480,7 @@ def test_sync_table_with_new_pk_in_stream(self, query_patch):
             'dbname': "dummy-db",
             'user': "dummy-user",
             'password': "dummy-passwd",
+            'private_key': "dummy-private_key",
             'warehouse': "dummy-wh",
             'default_target_schema': "dummy-schema",
             'file_format': "dummy-file-format"
@@ -542,6 +552,7 @@ def test_sync_table_with_stream_that_changes_to_have_no_pk(self, query_patch):
             'dbname': "dummy-db",
             'user': "dummy-user",
             'password': "dummy-passwd",
+            'private_key': "dummy-private_key",
             'warehouse': "dummy-wh",
             'default_target_schema': "dummy-schema",
             'file_format': "dummy-file-format"
@@ -591,6 +602,7 @@ def test_sync_table_with_stream_that_has_no_pk_but_get_a_new_one(self, query_pat
             'dbname': "dummy-db",
             'user': "dummy-user",
             'password': "dummy-passwd",
+            'private_key': "dummy-private_key",
             'warehouse': "dummy-wh",
             'default_target_schema': "dummy-schema",
             'file_format': "dummy-file-format"