feat(snowflake): adding oauth token bypass to snowflake (#12048)

datahub-project · Dec 6, 2024 · 1ed55f4 · 1ed55f4
1 parent b755c68
commit 1ed55f4
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 0 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_connection.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_connection.py
@@ -43,6 +43,7 @@
     "EXTERNAL_BROWSER_AUTHENTICATOR": EXTERNAL_BROWSER_AUTHENTICATOR,
     "KEY_PAIR_AUTHENTICATOR": KEY_PAIR_AUTHENTICATOR,
     "OAUTH_AUTHENTICATOR": OAUTH_AUTHENTICATOR,
+    "OAUTH_AUTHENTICATOR_TOKEN": OAUTH_AUTHENTICATOR,
 }
 
 _SNOWFLAKE_HOST_SUFFIX = ".snowflakecomputing.com"
@@ -104,6 +105,10 @@ class SnowflakeConnectionConfig(ConfigModel):
         description="Connect args to pass to Snowflake SqlAlchemy driver",
         exclude=True,
     )
+    token: Optional[str] = pydantic.Field(
+        default=None,
+        description="OAuth token from external identity provider. Not recommended for most use cases because it will not be able to refresh once expired.",
+    )
 
     def get_account(self) -> str:
         assert self.account_id
@@ -148,6 +153,18 @@ def authenticator_type_is_valid(cls, v, values):
         logger.info(f"using authenticator type '{v}'")
         return v
 
+    @pydantic.validator("token", always=True)
+    def validate_token_oauth_config(cls, v, values):
+        auth_type = values.get("authentication_type")
+        if auth_type == "OAUTH_AUTHENTICATOR_TOKEN":
+            if not v:
+                raise ValueError("Token required for OAUTH_AUTHENTICATOR_TOKEN.")
+        elif v is not None:
+            raise ValueError(
+                "Token can only be provided when using OAUTH_AUTHENTICATOR_TOKEN"
+            )
+        return v
+
     @staticmethod
     def _check_oauth_config(oauth_config: Optional[OAuthConfiguration]) -> None:
         if oauth_config is None:
@@ -333,6 +350,17 @@ def get_native_connection(self) -> NativeSnowflakeConnection:
                 application=_APPLICATION_NAME,
                 **connect_args,
             )
+        elif self.authentication_type == "OAUTH_AUTHENTICATOR_TOKEN":
+            return snowflake.connector.connect(
+                user=self.username,
+                account=self.account_id,
+                authenticator="oauth",
+                token=self.token,  # Token generated externally and provided directly to the recipe
+                warehouse=self.warehouse,
+                role=self.role,
+                application=_APPLICATION_NAME,
+                **connect_args,
+            )
         elif self.authentication_type == "OAUTH_AUTHENTICATOR":
             return self.get_oauth_connection()
         elif self.authentication_type == "KEY_PAIR_AUTHENTICATOR":

diff --git a/metadata-ingestion/tests/unit/snowflake/test_snowflake_source.py b/metadata-ingestion/tests/unit/snowflake/test_snowflake_source.py
@@ -130,6 +130,60 @@ def test_snowflake_oauth_happy_paths():
     )
 
 
+def test_snowflake_oauth_token_happy_path():
+    assert SnowflakeV2Config.parse_obj(
+        {
+            "account_id": "test",
+            "authentication_type": "OAUTH_AUTHENTICATOR_TOKEN",
+            "token": "valid-token",
+            "username": "test-user",
+            "oauth_config": None,
+        }
+    )
+
+
+def test_snowflake_oauth_token_without_token():
+    with pytest.raises(
+        ValidationError, match="Token required for OAUTH_AUTHENTICATOR_TOKEN."
+    ):
+        SnowflakeV2Config.parse_obj(
+            {
+                "account_id": "test",
+                "authentication_type": "OAUTH_AUTHENTICATOR_TOKEN",
+                "username": "test-user",
+            }
+        )
+
+
+def test_snowflake_oauth_token_with_wrong_auth_type():
+    with pytest.raises(
+        ValueError,
+        match="Token can only be provided when using OAUTH_AUTHENTICATOR_TOKEN.",
+    ):
+        SnowflakeV2Config.parse_obj(
+            {
+                "account_id": "test",
+                "authentication_type": "OAUTH_AUTHENTICATOR",
+                "token": "some-token",
+                "username": "test-user",
+            }
+        )
+
+
+def test_snowflake_oauth_token_with_empty_token():
+    with pytest.raises(
+        ValidationError, match="Token required for OAUTH_AUTHENTICATOR_TOKEN."
+    ):
+        SnowflakeV2Config.parse_obj(
+            {
+                "account_id": "test",
+                "authentication_type": "OAUTH_AUTHENTICATOR_TOKEN",
+                "token": "",
+                "username": "test-user",
+            }
+        )
+
+
 default_config_dict: Dict[str, Any] = {
     "username": "user",
     "password": "password",