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

@@ -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.

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:

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"]}},

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(

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"}
 
 
+@pytest.fixture
+def connector_5():
+    with mock.patch(
+        "airflow.providers.snowflake.hooks.snowflake._is_connector_5_or_later",
+        autospec=True,
+        return_value=True,
+    ):
+        yield
+
+
+@pytest.fixture
+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")
 
 
+@pytest.mark.parametrize(
+    ("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:pw@airflow.af_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(