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