From 17e86ccd9fab7ab48518405843efed33b01e5dbe Mon Sep 17 00:00:00 2001 From: mohamedmeqlad99 Date: Sat, 9 Aug 2025 07:27:37 +0300 Subject: [PATCH 1/3] feat(snowflake): Add Snowpark Container Services authentication support --- .../impl/snowflake/configuration.py | 54 ++++++++++++++----- 1 file changed, 42 insertions(+), 12 deletions(-) diff --git a/dlt/destinations/impl/snowflake/configuration.py b/dlt/destinations/impl/snowflake/configuration.py index 4f18c8b08b..2ef4d3ffa2 100644 --- a/dlt/destinations/impl/snowflake/configuration.py +++ b/dlt/destinations/impl/snowflake/configuration.py @@ -1,4 +1,5 @@ import dataclasses +import os from pathlib import Path from typing import Final, Optional, Any, Dict, ClassVar, List @@ -48,6 +49,23 @@ def parse_native_representation(self, native_value: Any) -> None: setattr(self, param, self.query.get(param)) def on_resolved(self) -> None: + # Auto-detect Snowpark Container Services token + token_path = "/snowflake/session/token" + if ( + not self.token + and os.path.exists(token_path) + ): + try: + with open(token_path, "r") as f: + self.token = f.read().strip() + # Use env vars if host not set + self.host = self.host or os.getenv("SNOWFLAKE_HOST") or os.getenv("SNOWFLAKE_ACCOUNT") + self.authenticator = self.authenticator or "oauth" + except Exception as ex: + raise ConfigurationValueError( + f"Failed to read Snowpark Container Services token: {ex}" + ) + if self.private_key_path: try: self.private_key = Path(self.private_key_path).read_text("ascii") @@ -56,10 +74,12 @@ def on_resolved(self) -> None: "Make sure that `private_key` in dlt recognized format is at" f" `{self.private_key_path}`. Note that binary formats are not supported" ) - if not self.password and not self.private_key and not self.authenticator: + + # Update validation to include token authentication + if not (self.password or self.private_key or self.authenticator or self.token): raise ConfigurationValueError( "`SnowflakeCredentials` requires one of the following to be specified: `password`," - " `private_key`, `authenticator` (OAuth2)." + " `private_key`, `authenticator` (OAuth2), or a Snowpark Container Services token." ) def get_query(self) -> Dict[str, Any]: @@ -70,21 +90,31 @@ def get_query(self) -> Dict[str, Any]: return query def to_connector_params(self) -> Dict[str, Any]: - # gather all params in query query = self.get_query() if self.private_key: query["private_key"] = decode_private_key(self.private_key, self.private_key_passphrase) - - # we do not want passphrase to be passed query.pop("private_key_passphrase", None) - conn_params: Dict[str, Any] = dict( - query, - user=self.username, - password=self.password, - account=self.host, - database=self.database, - ) + # Support token-based authentication + if self.token: + conn_params = dict( + query, + token=self.token, + account=self.host, + authenticator=self.authenticator or "oauth", + database=self.database, + ) + # Remove user/password when using token auth + conn_params.pop("user", None) + conn_params.pop("password", None) + else: + conn_params = dict( + query, + user=self.username, + password=self.password, + account=self.host, + database=self.database, + ) if self.application != "" and "application" not in conn_params: conn_params["application"] = self.application From 3e1f5587fa82d4338fba6b0e8dbfe6555a79d64b Mon Sep 17 00:00:00 2001 From: mohamedmeqlad99 Date: Sat, 9 Aug 2025 07:30:28 +0300 Subject: [PATCH 2/3] adding test file --- .../test_snowflake_snowpark_token.py | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 tests/destinations/test_snowflake_snowpark_token.py diff --git a/tests/destinations/test_snowflake_snowpark_token.py b/tests/destinations/test_snowflake_snowpark_token.py new file mode 100644 index 0000000000..417574c568 --- /dev/null +++ b/tests/destinations/test_snowflake_snowpark_token.py @@ -0,0 +1,60 @@ +import os +import tempfile +from pathlib import Path +from dlt.destinations.impl.snowflake.configuration import SnowflakeCredentials + +def test_snowpark_token(monkeypatch): + """Test Snowpark Container Services token authentication""" + with tempfile.NamedTemporaryFile("w+", delete=False) as tf: + tf.write("test-token") + tf.flush() + token_path = tf.name + + # Set up test environment + monkeypatch.setenv("SNOWFLAKE_HOST", "test-account") + monkeypatch.setenv("SNOWFLAKE_ACCOUNT", "test-account") + + # Mock token path existence and file reading + monkeypatch.setattr("os.path.exists", lambda p: p == token_path) + monkeypatch.setattr("builtins.open", lambda p, mode="r": open(token_path, mode)) + + # Test credentials resolution + creds = SnowflakeCredentials() + creds.on_resolved.__globals__["token_path"] = token_path + creds.on_resolved() + + assert creds.token == "test-token" + assert creds.host == "test-account" + assert creds.authenticator == "oauth" + + # Test connector parameters + params = creds.to_connector_params() + assert params["token"] == "test-token" + assert params["account"] == "test-account" + assert params["authenticator"] == "oauth" + assert "user" not in params + assert "password" not in params + + # Cleanup + Path(token_path).unlink() + +def test_explicit_credentials_preferred(monkeypatch): + """Test that explicit credentials are preferred over Snowpark token""" + with tempfile.NamedTemporaryFile("w+", delete=False) as tf: + tf.write("test-token") + tf.flush() + token_path = tf.name + + # Test with explicit credentials + creds = SnowflakeCredentials( + token="explicit-token", + host="explicit-account" + ) + creds.on_resolved.__globals__["token_path"] = token_path + creds.on_resolved() + + assert creds.token == "explicit-token" + assert creds.host == "explicit-account" + + # Cleanup + Path(token_path).unlink() \ No newline at end of file From 54818574b722cde4b0eac05fe17b6c0b4da45d19 Mon Sep 17 00:00:00 2001 From: mohamedmeqlad99 Date: Sat, 9 Aug 2025 07:31:51 +0300 Subject: [PATCH 3/3] update docs --- .../docs/destinations/snowflake_config.md | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 docs/website/docs/destinations/snowflake_config.md diff --git a/docs/website/docs/destinations/snowflake_config.md b/docs/website/docs/destinations/snowflake_config.md new file mode 100644 index 0000000000..2d99dba98b --- /dev/null +++ b/docs/website/docs/destinations/snowflake_config.md @@ -0,0 +1,63 @@ +# Snowflake Configuration + +## Authentication Methods + +Snowflake destination supports several authentication methods: + +### Username and Password +```python +pipeline = dlt.pipeline( + pipeline_name="my_pipeline", + destination="snowflake", + credentials={ + "username": "my_user", + "password": "my_password", + "account": "my_account" + } +) +``` + +### Private Key +```python +pipeline = dlt.pipeline( + pipeline_name="my_pipeline", + destination="snowflake", + credentials={ + "username": "my_user", + "private_key_path": "path/to/rsa_key.p8", + "account": "my_account" + } +) +``` + +### Snowpark Container Services + +When running inside Snowpark Container Services, dlt will automatically detect and use the mounted OAuth token and environment variables: + +```python +pipeline = dlt.pipeline( + pipeline_name="my_pipeline", + destination="snowflake", + dataset_name="my_dataset" +) + +# No credentials needed - dlt will automatically use: +# - Token from /snowflake/session/token +# - Account from SNOWFLAKE_HOST or SNOWFLAKE_ACCOUNT env vars +``` + +You can still provide explicit credentials to override the auto-detection: + +```python +pipeline = dlt.pipeline( + pipeline_name="my_pipeline", + destination="snowflake", + credentials={ + "token": "my_token", # Override container token + "account": "my_account" # Override SNOWFLAKE_HOST + } +) +``` + +## Additional Configuration +# ...existing documentation... \ No newline at end of file