Skip to content

Commit

Permalink
Add ability to specify host/port for Snowflake connection (#1063)
Browse files Browse the repository at this point in the history
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:
localstack-samples/localstack-snowflake-samples#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"
...
```
  • Loading branch information
whummer authored Jun 24, 2024
1 parent 5f5616c commit 06170e8
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 0 deletions.
2 changes: 2 additions & 0 deletions cosmos/profiles/snowflake/user_pass.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
24 changes: 24 additions & 0 deletions tests/profiles/snowflake/test_snowflake_user_pass.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 06170e8

Please sign in to comment.