From 06170e83306c5277286c64a383b7ddf82e5cb625 Mon Sep 17 00:00:00 2001 From: Waldemar Hummer <waldemar.hummer@gmail.com> Date: Mon, 24 Jun 2024 15:36:08 +0200 Subject: [PATCH] Add ability to specify host/port for Snowflake connection (#1063) Add ability to specify `host`/`port` for Snowflake connection. At LocalStack, we have recently started building a Snowflake emulator that allows running SF queries entirely on the local machine: https://blog.localstack.cloud/2024-05-22-introducing-localstack-for-snowflake/ . As part of a sample application we're building, we have an Apache Airflow DAG that uses Cosmos (and DBT) to connect to the local Snowflake emulator running on `localhost`. Here is a link to the sample app: https://github.com/localstack-samples/localstack-snowflake-samples/pull/12 Currently, we're hardcoding this integration in the user DAG file itself, [see here](https://github.com/localstack-samples/localstack-snowflake-samples/pull/12/files#diff-559d4f883ad589522b8a9d33f87fe95b0da72ac29b775e98b273a8eb3ede9924R10-R19): ``` ... from cosmos.profiles.snowflake.user_pass import SnowflakeUserPasswordProfileMapping ... SnowflakeUserPasswordProfileMapping.airflow_param_mapping["host"] = "extra.host" SnowflakeUserPasswordProfileMapping.airflow_param_mapping["port"] = "extra.port" ... ``` --- cosmos/profiles/snowflake/user_pass.py | 2 ++ .../snowflake/test_snowflake_user_pass.py | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/cosmos/profiles/snowflake/user_pass.py b/cosmos/profiles/snowflake/user_pass.py index a6042495e..3fc6595c9 100644 --- a/cosmos/profiles/snowflake/user_pass.py +++ b/cosmos/profiles/snowflake/user_pass.py @@ -40,6 +40,8 @@ class SnowflakeUserPasswordProfileMapping(BaseProfileMapping): "warehouse": "extra.warehouse", "schema": "schema", "role": "extra.role", + "host": "extra.host", + "port": "extra.port", } def can_claim_connection(self) -> bool: diff --git a/tests/profiles/snowflake/test_snowflake_user_pass.py b/tests/profiles/snowflake/test_snowflake_user_pass.py index 8113d8528..6514bdf8d 100644 --- a/tests/profiles/snowflake/test_snowflake_user_pass.py +++ b/tests/profiles/snowflake/test_snowflake_user_pass.py @@ -231,3 +231,27 @@ def test_appends_region() -> None: "database": conn.extra_dejson.get("database"), "warehouse": conn.extra_dejson.get("warehouse"), } + + +def test_appends_host_and_port() -> None: + """ + Tests that host/port extras are appended to the connection settings. + """ + conn = Connection( + conn_id="my_snowflake_connection", + conn_type="snowflake", + login="my_user", + password="my_password", + schema="my_schema", + extra=json.dumps( + { + "host": "snowflake.localhost.localstack.cloud", + "port": 4566, + } + ), + ) + + with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): + profile_mapping = SnowflakeUserPasswordProfileMapping(conn) + assert profile_mapping.profile["host"] == "snowflake.localhost.localstack.cloud" + assert profile_mapping.profile["port"] == 4566