Make region optional in Snowflake connection

[dwreeves]

Snowflake allows and if anything seems to encourage not adding in the region parameter to a Snowflake connection. The reason is because Snowflake recently moved to a system where the "region" is no longer part of the URI.

Under the hood, Airflow just passes to the Snowflake SQLAlchemy URL(), and under the hood Snowflake's URL() allows for region to be optional.

Fix

region should be allowed to be optional, basically. I wanted to do an approach that did not rely on magically trimming a trailing ., so something that explicitly checks for when region is not specified.

There were two approaches I could go with that meets the requirement: either fix it in Python, or fix it in profiles/snowflake.py. I went with the former approach because doing lots of things in Jinja2 is usually harder to maintain than doing the equivalent in Python, but it would have looked like this:

"account": "{{ env_var('SNOWFLAKE_ACCOUNT') ~ '.' ~ env_var('SNOWFLAKE_REGION')"
           " if env_var('SNOWFLAKE_REGION')"
           " else env_var('SNOWFLAKE_ACCOUNT') }}",

Pyproject.toml change

I also added pytest to the project as it wasn't specified as a dependency, but the tests/ folder uses it.

[chrishronek]

@dwreeves could you give me an example of what the formatted snowflake URI might look like (obviously without an actual password)? The way this is set up right now ensures that if there's a working snowflake connection in Airflow, then it can be used for cosmos. When I take away the region for a snowflake region, the connection is no longer valid in Airflow

[dwreeves]

@pohek321 Hmm, that's interesting.

To be clear, if you are on the legacy system for Snowflake account IDs, and you want to construct a URI without the region field, you move it over to the account part of the URI.

I'm on provider package 2.7.0 and my secret looks something like this and it works fine:

{"conn_type": "snowflake", "user": "SOME_USER", "password": "SOME_PASSWORD", "extra": "{\"account\": \"12345.us-east-1\", \"warehouse\": \"SOME_WAREHOUSE\", \"database\": \"SOME_DATABASE\", \"role\": \"SOME_ROLE\"}"}

Note that it's {"account": "12345.us-east-1"}. Removing the region from that part of the URI for legacy Snowflake connections only works for us-west-1 I believe.

For good measure, my Airflow dbt integration does the following. (1) Places the secrets directly in the .yml, i.e. I'm not using {{ env_var(...) }}, (2) dump the .yml file in /tmp (FWIW, I think I like your approach better than mine):

conn = self.snowflake_hook.get_connection(conn_id)
profile = self.dbt_profile_name: {
    "outputs": {
        self.dbt_target_name: {
            "account": conn.extra_dejson.get("account"),
            ...
        }
    }
}

Also FWIW, I also am the last person who committed to the docs for the Snowflake connection 😅 https://github.com/apache/airflow/blob/main/docs/apache-airflow-providers-snowflake/connections/snowflake.rst which I edited because I struggled a lot with this as it's very confusing. 🥴

[chrishronek]

@dwreeves, I appreciate your patience while I wrap my head around this.

I think this change would make sense, so in Airflow, your connection would look something like this, right:

[dwreeves]

The host does not do anything for the Snowflake connection in Airflow. I was very confused by this at first myself, which is the main thing that motivated me to rewrite the docs, although it's unfortunate it's still there in the UI (hm, wonder if that can be changed to be hidden? 👀).

So remove the host, and for the account you should set equal to ab12345.us-east-1. And that will make the connection work.

Sorry for this PR being so confusing, but I know I can't be the only one who has my Snowflake Airflow connection set up to not use the region, so want to make sure this works trivially not just for me but for others as well. 😅

[dwreeves]

Turns out yes, we can remove the host from the UI. 😄 apache/airflow#29208

[chrishronek]

@dwreeves, good to know! I did what you suggested, and it worked! I did not realize that the host parameter was optional.

However, I was able to break it (and by break, I mean the connection test in Airflow succeeds, but Cosmos cannot authenticate). It broke by putting <account>.<region> in the account field as well as <region> in the region field -- interestingly enough, Airflow was OK with this (Cosmos was not).

So I added a couple of lines to handle that scenario! If you are cool with the commit, then I'll approve and merge, and we'll tag a new release so you can access these features.

Additionally, thank you so much for your contributions! If you need anything else (or have ideas for improvements), don't hesitate to reach out or open issues. 🍻

[chrishronek]

Thanks, @dwreeves! Approved.