This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5472d642898 Fix Snowflake connections ignoring region with
`snowflake-connector-python` 5 (#75334)
5472d642898 is described below
commit 5472d642898faa9da3eada2ce2e98a76d6cfdd10
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 21 21:16:42 2026 +0100
Fix Snowflake connections ignoring region with `snowflake-connector-python`
5 (#75334)
snowflake-connector-python 5.0 no longer accepts `region`: it warns
"Unknown parameter 'region'" and builds the host from `account` alone, so a
connection with the `region` extra dials the wrong host. On connector 5 the
hook
now passes the host earlier connectors built from the region instead.
Connector 5 also warns on every connection unless
`enable_server_session_keep_alive_auto_detection` is set, and plans to
change
its default. The hook now sets it to True, the behaviour of earlier
connectors, and exposes it as a connection extra.
* Tolerate a missing region in Snowflake connector-5 params
* Fix mypy return type in Snowflake connect-kwargs test helper
---
providers/snowflake/docs/connections/snowflake.rst | 6 +
providers/snowflake/provider.yaml | 8 +
.../providers/snowflake/get_provider_info.py | 5 +
.../airflow/providers/snowflake/hooks/snowflake.py | 67 +++++++-
.../tests/unit/snowflake/hooks/test_snowflake.py | 175 ++++++++++++++++++++-
5 files changed, 257 insertions(+), 4 deletions(-)
diff --git a/providers/snowflake/docs/connections/snowflake.rst
b/providers/snowflake/docs/connections/snowflake.rst
index 90002bfde79..51b2b749f3a 100644
--- a/providers/snowflake/docs/connections/snowflake.rst
+++ b/providers/snowflake/docs/connections/snowflake.rst
@@ -100,6 +100,12 @@ Extra (optional)
* ``host``: Target Snowflake hostname to connect to (e.g., for local
testing with LocalStack).
* ``port``: Target Snowflake port to connect to (e.g., for local testing
with LocalStack).
* ``ocsp_fail_open``: Specify `ocsp_fail_open
<https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect#label-python-ocsp-choosing-fail-open-or-fail-close-mode>`_.
+ * ``enable_server_session_keep_alive_auto_detection``: Whether closing the
connection checks for
+ async queries still running on it first. ``true`` (the default) keeps
the server session open
+ while any are running, so they finish; ``false`` always logs out on
close, and Snowflake may
+ cancel those queries. Takes effect with
``snowflake-connector-python>=5.0.0``; earlier connectors
+ always keep the session while async queries run. The hook sets ``true``
when the extra is
+ absent, so the connector's planned default change does not alter
behaviour.
* ``proxy_host``: Proxy hostname to use for connecting to Snowflake.
* ``proxy_port``: Proxy port to use for connecting to Snowflake.
* ``proxy_user``: Proxy username for authentication with the proxy server.
diff --git a/providers/snowflake/provider.yaml
b/providers/snowflake/provider.yaml
index 71cd310e438..4f6d31a0123 100644
--- a/providers/snowflake/provider.yaml
+++ b/providers/snowflake/provider.yaml
@@ -277,6 +277,14 @@ connection-types:
- boolean
- 'null'
description: Turns off OCSP certificate checks
+ enable_server_session_keep_alive_auto_detection:
+ label: Keep session while async queries run
+ schema:
+ type:
+ - boolean
+ - 'null'
+ default: true
+ description: Skip logout on close while async queries are running
(snowflake-connector-python 5.0+)
proxy_host:
label: Proxy Host
schema:
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
index e0ed30efd8e..bc8edf6ef53 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
@@ -155,6 +155,11 @@ def get_provider_info():
"schema": {"type": ["boolean", "null"]},
"description": "Turns off OCSP certificate checks",
},
+ "enable_server_session_keep_alive_auto_detection": {
+ "label": "Keep session while async queries run",
+ "schema": {"type": ["boolean", "null"], "default":
True},
+ "description": "Skip logout on close while async
queries are running (snowflake-connector-python 5.0+)",
+ },
"proxy_host": {"label": "Proxy Host", "schema": {"type":
["string", "null"]}},
"proxy_port": {"label": "Proxy Port", "schema": {"type":
["integer", "null"]}},
"proxy_user": {"label": "Proxy User", "schema": {"type":
["string", "null"]}},
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
index 9aac0e6dfe1..52d2f03d541 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
@@ -53,6 +53,10 @@ from airflow.utils.strings import to_boolean
OAUTH_REQUEST_TIMEOUT = 30 # seconds, avoid hanging tasks on token request
OAUTH_EXPIRY_BUFFER = 30
SUPPORTED_GRANT_TYPES = {"refresh_token", "client_credentials", "password"}
+LOGOUT_AUTO_DETECTION_PARAM = "enable_server_session_keep_alive_auto_detection"
+# Logout behaviour of connectors before 5.0: keep the server session while
async queries
+# started on the connection are still running. Connector 5 plans to change its
default.
+LOGOUT_AUTO_DETECTION_DEFAULT = True
T = TypeVar("T")
@@ -71,6 +75,20 @@ def _try_to_boolean(value: Any):
return value
+def _is_connector_5_or_later() -> bool:
+ """
+ Return True when the installed ``snowflake-connector-python`` is 5.0 or
later.
+
+ Connector 5 no longer accepts ``region`` and is the first to accept
+ ``enable_server_session_keep_alive_auto_detection``. Remove this and the
pre-5 branches
+ once the provider requires ``snowflake-connector-python>=5.0``.
+ """
+ # Lazy: importing snowflake.connector loads the whole connector.
+ from snowflake.connector.version import VERSION
+
+ return VERSION >= (5,)
+
+
def _is_retryable_oauth_error(exception: BaseException) -> bool:
"""Return True if exception is retryable for OAuth token request."""
if isinstance(exception, (ConnectionError, Timeout)):
@@ -295,6 +313,12 @@ class SnowflakeHook(DbApiHook):
"insecure_mode": BooleanField(
label=lazy_gettext("Insecure mode"), description="Turns off
OCSP certificate checks"
),
+ LOGOUT_AUTO_DETECTION_PARAM: BooleanField(
+ label=lazy_gettext("Keep session while async queries run"),
+ description="Skip logout on close while async queries are
running "
+ "(snowflake-connector-python 5.0+)",
+ default=True,
+ ),
"proxy_host": StringField(lazy_gettext("Proxy Host"),
widget=BS3TextFieldWidget()),
"proxy_port": IntegerField(
lazy_gettext("Proxy Port"),
@@ -552,6 +576,11 @@ class SnowflakeHook(DbApiHook):
if ocsp_fail_open is not None:
conn_config["ocsp_fail_open"] = _try_to_boolean(ocsp_fail_open)
+ logout_auto_detection = extra_dict.get(LOGOUT_AUTO_DETECTION_PARAM)
+
+ if logout_auto_detection is not None:
+ conn_config[LOGOUT_AUTO_DETECTION_PARAM] =
_try_to_boolean(logout_auto_detection)
+
def _add_workload_identity_config(
self,
conn_config: dict[str, Any],
@@ -815,6 +844,7 @@ class SnowflakeHook(DbApiHook):
"client_store_temporary_credential",
"json_result_force_utf8_decoding",
"ocsp_fail_open",
+ LOGOUT_AUTO_DETECTION_PARAM,
"proxy_host",
"proxy_port",
"proxy_user",
@@ -823,11 +853,38 @@ class SnowflakeHook(DbApiHook):
}
)
+ def _get_connector_params(self) -> dict[str, Any]:
+ """
+ Return ``_get_conn_params()`` in the form the installed
``snowflake-connector-python`` accepts.
+
+ Connector 5 ignores ``region`` and builds the host from ``account``
alone, so the hook
+ passes the host earlier connectors built from the region instead. It
also warns on every
+ connection unless ``enable_server_session_keep_alive_auto_detection``
is set, which earlier
+ connectors don't accept.
+ """
+ conn_params: dict[str, Any] = self._get_conn_params()
+
+ if not _is_connector_5_or_later():
+ conn_params.pop(LOGOUT_AUTO_DETECTION_PARAM, None)
+ return conn_params
+
+ region = conn_params.pop("region", None)
+ # Same host as construct_hostname() in earlier connectors: "us-west-2"
is the default
+ # deployment and adds nothing, and the region replaces any region in
the account.
+ if region and region != "us-west-2" and "host" not in conn_params:
+ account =
_validate_account_component(conn_params["account"].split(".")[0], "account")
+ region = _validate_account_component(region, "region")
+ top_level_domain = "cn" if region.lower().startswith("cn-") else
"com"
+ conn_params["account"] = account
+ conn_params["host"] =
f"{account}.{region}.snowflakecomputing.{top_level_domain}"
+ conn_params.setdefault(LOGOUT_AUTO_DETECTION_PARAM,
LOGOUT_AUTO_DETECTION_DEFAULT)
+ return conn_params
+
def get_conn(self) -> SnowflakeConnection:
"""Return a snowflake.connection object."""
from snowflake.connector import connect
- conn_config = self._get_conn_params()
+ conn_config = self._get_connector_params()
conn = connect(**conn_config)
return conn
@@ -849,6 +906,12 @@ class SnowflakeHook(DbApiHook):
if "ocsp_fail_open" in conn_params:
engine_kwargs.setdefault("connect_args", {})
engine_kwargs["connect_args"]["ocsp_fail_open"] =
conn_params["ocsp_fail_open"]
+ if _is_connector_5_or_later():
+ engine_kwargs.setdefault("connect_args", {})
+ engine_kwargs["connect_args"].setdefault(
+ LOGOUT_AUTO_DETECTION_PARAM,
+ conn_params.get(LOGOUT_AUTO_DETECTION_PARAM,
LOGOUT_AUTO_DETECTION_DEFAULT),
+ )
for key in ["session_parameters", "private_key"]:
if conn_params.get(key):
engine_kwargs.setdefault("connect_args", {})
@@ -866,7 +929,7 @@ class SnowflakeHook(DbApiHook):
from airflow import __version__ as airflow_version
from airflow.providers.snowflake import __version__ as provider_version
- conn_config = self._get_conn_params()
+ conn_config = self._get_connector_params()
session = Session.builder.configs(conn_config).create()
# add query tag for observability
session.update_query_tag(
diff --git a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
index 13a70cd2cb9..897a7f8a224 100644
--- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
@@ -35,7 +35,7 @@ from requests.exceptions import ConnectionError, HTTPError
from airflow.models import Connection
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, timezone
-from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook,
_is_connector_5_or_later
if TYPE_CHECKING:
from pathlib import Path
@@ -73,6 +73,40 @@ CONN_PARAMS_OAUTH_BASE = {
CONN_PARAMS_OAUTH = CONN_PARAMS_OAUTH_BASE | {"refresh_token": "secrettoken"}
[email protected]
+def connector_5():
+ with mock.patch(
+ "airflow.providers.snowflake.hooks.snowflake._is_connector_5_or_later",
+ autospec=True,
+ return_value=True,
+ ):
+ yield
+
+
[email protected]
+def connector_before_5():
+ with mock.patch(
+ "airflow.providers.snowflake.hooks.snowflake._is_connector_5_or_later",
+ autospec=True,
+ return_value=False,
+ ):
+ yield
+
+
+def _get_conn_connect_kwargs(extra: dict[str, Any]) -> dict[str, Any]:
+ """Return the kwargs ``SnowflakeHook.get_conn`` passes to
``snowflake.connector.connect``."""
+ connection_kwargs = deepcopy(BASE_CONNECTION_KWARGS)
+ connection_kwargs["extra"].update(extra)
+
+ with (
+ mock.patch.dict("os.environ",
AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()),
+ mock.patch("snowflake.connector.connect") as mock_connect,
+ ):
+ SnowflakeHook(snowflake_conn_id="test_conn").get_conn()
+
+ return dict(mock_connect.call_args.kwargs)
+
+
@pytest.fixture
def unencrypted_temporary_private_key(tmp_path: Path) -> Path:
key = rsa.generate_private_key(backend=default_backend(),
public_exponent=65537, key_size=2048)
@@ -107,6 +141,20 @@ def
base64_encoded_encrypted_private_key(encrypted_temporary_private_key: Path)
return
base64.b64encode(encrypted_temporary_private_key.read_bytes()).decode("utf-8")
[email protected](
+ ("version", "expected"),
+ [
+ pytest.param((3, 17, 0, None), False, id="3.17"),
+ pytest.param((4, 8, 0, None), False, id="4.8"),
+ pytest.param((5, 0, 0, None), True, id="5.0"),
+ pytest.param((5, 1, 0, None), True, id="5.1"),
+ ],
+)
+def test_is_connector_5_or_later(version, expected):
+ with mock.patch("snowflake.connector.version.VERSION", version):
+ assert _is_connector_5_or_later() is expected
+
+
class TestPytestSnowflakeHook:
@pytest.mark.parametrize(
("connection_kwargs", "expected_uri", "expected_conn_params"),
@@ -837,6 +885,7 @@ class TestPytestSnowflakeHook:
== "PARTNER_NAME"
)
+ @pytest.mark.usefixtures("connector_before_5")
def test_get_conn_should_call_connect(self):
with (
mock.patch.dict(
@@ -849,6 +898,108 @@ class TestPytestSnowflakeHook:
mock_connect.assert_called_once_with(**hook._get_conn_params())
assert mock_connect.return_value == conn
+ @pytest.mark.parametrize(
+ ("extra", "expected_account", "expected_host"),
+ [
+ pytest.param({"region": ""}, "airflow", None, id="no-region"),
+ pytest.param(
+ {"region": "us-east-2.aws"},
+ "airflow",
+ "airflow.us-east-2.aws.snowflakecomputing.com",
+ id="region",
+ ),
+ pytest.param(
+ {"region": "cn-north-1.aws"},
+ "airflow",
+ "airflow.cn-north-1.aws.snowflakecomputing.cn",
+ id="china-region",
+ ),
+ pytest.param(
+ {"account": "airflow.eu-west-1", "region": "us-east-2.aws"},
+ "airflow",
+ "airflow.us-east-2.aws.snowflakecomputing.com",
+ id="region-replaces-account-region",
+ ),
+ pytest.param({"region": "us-west-2"}, "airflow", None,
id="default-region"),
+ pytest.param(
+ {"region": "us-east-2.aws", "host": "localhost"},
+ "airflow",
+ "localhost",
+ id="explicit-host-ignores-region",
+ ),
+ ],
+ )
+ @pytest.mark.usefixtures("connector_5")
+ def test_get_conn_on_connector_5_builds_host_from_region(self, extra,
expected_account, expected_host):
+ connect_kwargs = _get_conn_connect_kwargs(extra)
+
+ assert "region" not in connect_kwargs
+ assert connect_kwargs["account"] == expected_account
+ assert connect_kwargs.get("host") == expected_host
+
+ @pytest.mark.usefixtures("connector_5")
+ def test_get_conn_on_connector_5_rejects_region_that_changes_host(self):
+ with pytest.raises(ValueError, match="Invalid Snowflake region"):
+ _get_conn_connect_kwargs({"region": "evil.example/x"})
+
+ @pytest.mark.parametrize(
+ ("extra", "expected"),
+ [
+ pytest.param({}, True, id="unset-keeps-earlier-behaviour"),
+ pytest.param({"enable_server_session_keep_alive_auto_detection":
False}, False, id="false"),
+ pytest.param({"enable_server_session_keep_alive_auto_detection":
"false"}, False, id="false-str"),
+ pytest.param({"enable_server_session_keep_alive_auto_detection":
True}, True, id="true"),
+ ],
+ )
+ @pytest.mark.usefixtures("connector_5")
+ def test_get_conn_on_connector_5_sets_logout_auto_detection(self, extra,
expected):
+ connect_kwargs = _get_conn_connect_kwargs(extra)
+
+ assert
connect_kwargs["enable_server_session_keep_alive_auto_detection"] is expected
+
+ @pytest.mark.usefixtures("connector_before_5")
+ def test_get_conn_before_connector_5_drops_logout_auto_detection(self):
+ connect_kwargs =
_get_conn_connect_kwargs({"enable_server_session_keep_alive_auto_detection":
False})
+
+ assert "enable_server_session_keep_alive_auto_detection" not in
connect_kwargs
+
+ @pytest.mark.parametrize(
+ ("extra", "engine_kwargs", "expected"),
+ [
+ pytest.param({}, None, True, id="unset"),
+ pytest.param(
+ {"enable_server_session_keep_alive_auto_detection": False},
None, False, id="connection-false"
+ ),
+ pytest.param(
+ {},
+ {"connect_args":
{"enable_server_session_keep_alive_auto_detection": False}},
+ False,
+ id="caller-false",
+ ),
+ ],
+ )
+ @pytest.mark.usefixtures("connector_5")
+ def test_get_sqlalchemy_engine_on_connector_5_passes_logout_auto_detection(
+ self, extra, engine_kwargs, expected
+ ):
+ connection_kwargs = deepcopy(BASE_CONNECTION_KWARGS)
+ connection_kwargs["extra"].update(extra)
+
+ with (
+ mock.patch.dict("os.environ",
AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()),
+ mock.patch(
+ "airflow.providers.snowflake.hooks.snowflake.create_engine",
autospec=True
+ ) as mock_create_engine,
+ ):
+
SnowflakeHook(snowflake_conn_id="test_conn").get_sqlalchemy_engine(engine_kwargs=engine_kwargs)
+
+ mock_create_engine.assert_called_once_with(
+ "snowflake://user:[email protected]_region/db/public"
+
"?application=AIRFLOW&authenticator=snowflake&role=af_role&warehouse=af_wh",
+ connect_args={"enable_server_session_keep_alive_auto_detection":
expected},
+ )
+
+ @pytest.mark.usefixtures("connector_before_5")
def test_get_sqlalchemy_engine_should_support_pass_auth(self):
with (
mock.patch.dict(
@@ -864,6 +1015,7 @@ class TestPytestSnowflakeHook:
)
assert mock_create_engine.return_value == conn
+ @pytest.mark.usefixtures("connector_before_5")
def test_get_sqlalchemy_engine_should_support_insecure_mode(self):
connection_kwargs = deepcopy(BASE_CONNECTION_KWARGS)
connection_kwargs["extra"]["extra__snowflake__insecure_mode"] = "True"
@@ -881,6 +1033,7 @@ class TestPytestSnowflakeHook:
)
assert mock_create_engine.return_value == conn
+ @pytest.mark.usefixtures("connector_before_5")
def
test_get_sqlalchemy_engine_should_support_json_result_force_utf8_decoding(self):
connection_kwargs = deepcopy(BASE_CONNECTION_KWARGS)
connection_kwargs["extra"]["extra__snowflake__json_result_force_utf8_decoding"]
= "True"
@@ -898,6 +1051,7 @@ class TestPytestSnowflakeHook:
)
assert mock_create_engine.return_value == conn
+ @pytest.mark.usefixtures("connector_before_5")
def test_get_sqlalchemy_engine_should_support_session_parameters(self):
connection_kwargs = deepcopy(BASE_CONNECTION_KWARGS)
connection_kwargs["extra"]["session_parameters"] = {"TEST_PARAM":
"AA", "TEST_PARAM_B": 123}
@@ -929,6 +1083,7 @@ class TestPytestSnowflakeHook:
assert "private_key" in
mock_create_engine.call_args.kwargs["connect_args"]
assert mock_create_engine.return_value == conn
+ @pytest.mark.usefixtures("connector_before_5")
def test_get_sqlalchemy_engine_should_support_ocsp_fail_open(self):
connection_kwargs = deepcopy(BASE_CONNECTION_KWARGS)
connection_kwargs["extra"]["ocsp_fail_open"] = "False"
@@ -1286,6 +1441,7 @@ class TestPytestSnowflakeHook:
hook.get_openlineage_database_specific_lineage(mock.MagicMock())
@pytest.mark.skipif(sys.version_info >= (3, 12), reason="Snowpark Python
doesn't support Python 3.12 yet")
+ @pytest.mark.usefixtures("connector_5")
@mock.patch("snowflake.snowpark.Session.builder")
def test_get_snowpark_session(self, mock_session_builder):
from airflow import __version__ as airflow_version
@@ -1301,7 +1457,22 @@ class TestPytestSnowflakeHook:
session = hook.get_snowpark_session()
assert session == mock_session
-
mock_session_builder.configs.assert_called_once_with(hook._get_conn_params())
+ mock_session_builder.configs.assert_called_once_with(
+ {
+ "account": "airflow",
+ "application": "AIRFLOW",
+ "authenticator": "snowflake",
+ "database": "db",
+ "enable_server_session_keep_alive_auto_detection": True,
+ "host": "airflow.af_region.snowflakecomputing.com",
+ "password": "pw",
+ "role": "af_role",
+ "schema": "public",
+ "session_parameters": None,
+ "user": "user",
+ "warehouse": "af_wh",
+ }
+ )
# Verify that update_query_tag was called with the expected tag
dictionary
mock_session.update_query_tag.assert_called_once_with(
