apache-airflow-providers-snowflake 6.3.1__tar.gz → 6.4.0__tar.gz

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-snowflake
-Version: 6.3.1
+Version: 6.4.0
 Summary: Provider package apache-airflow-providers-snowflake for Apache Airflow
 Keywords: airflow-provider,snowflake,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -30,8 +30,8 @@ Requires-Dist: snowflake-sqlalchemy>=1.4.0
 Requires-Dist: snowflake-snowpark-python>=1.17.0;python_version<'3.12'
 Requires-Dist: apache-airflow-providers-openlineage>=2.3.0 ; extra == "openlineage"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -63,7 +63,7 @@ Provides-Extra: openlineage
 
 Package ``apache-airflow-providers-snowflake``
 
-Release: ``6.3.1``
+Release: ``6.4.0``
 
 
 `Snowflake <https://www.snowflake.com/>`__
@@ -76,7 +76,7 @@ This is a provider package for ``snowflake`` provider. All classes for this prov
 are in ``airflow.providers.snowflake`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0/>`_.
 
 Installation
 ------------
@@ -125,5 +125,5 @@ Dependent package
 ==================================================================================================================  =================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0/changelog.html>`_.
 

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-snowflake``
 
-Release: ``6.3.1``
+Release: ``6.4.0``
 
 
 `Snowflake <https://www.snowflake.com/>`__
@@ -36,7 +36,7 @@ This is a provider package for ``snowflake`` provider. All classes for this prov
 are in ``airflow.providers.snowflake`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0/>`_.
 
 Installation
 ------------
@@ -85,4 +85,4 @@ Dependent package
 ==================================================================================================================  =================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0/changelog.html>`_.

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/docs/changelog.rst

@@ -27,6 +27,29 @@
 Changelog
 ---------
 
+6.4.0
+.....
+
+Features
+~~~~~~~~
+
+* ``Extend SnowflakeHook OAuth implementation to support external IDPs and client_credentials grant (#51620)``
+
+Bug Fixes
+~~~~~~~~~
+
+* ``fix: make query_ids in SnowflakeSqlApiOperator in deferrable mode consistent (#51542)``
+* ``fix: Duplicate region in Snowflake URI no longer breaks OpenLineage (#50831)``
+* ``Do not allow semicolons in CopyFromExternalStageToSnowflakeOperator fieldS (#51734)``
+
+Misc
+~~~~
+
+* ``Port ''ti.run'' to Task SDK execution path (#50141)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+
 6.3.1
 .....
 

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/docs/connections/snowflake.rst

@@ -58,6 +58,8 @@ Extra (optional)
     * ``warehouse``: Snowflake warehouse name.
     * ``role``: Snowflake role.
     * ``authenticator``: To connect using OAuth set this parameter ``oauth``.
+    * ``token_endpoint``: Specify token endpoint for external OAuth provider.
+    * ``grant_type``: Specify grant type for OAuth authentication. Currently supported: ``refresh_token`` (default), ``client_credentials``.
     * ``refresh_token``: Specify refresh_token for OAuth connection.
     * ``private_key_file``: Specify the path to the private key file.
     * ``private_key_content``: Specify the content of the private key file in base64 encoded format. You can use the following Python code to encode the private key:

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/docs/index.rst

@@ -78,7 +78,7 @@ apache-airflow-providers-snowflake package
 `Snowflake <https://www.snowflake.com/>`__
 
 
-Release: 6.3.1
+Release: 6.4.0
 
 Provider package
 ----------------
@@ -138,5 +138,5 @@ Downloading official packages
 You can download officially released packages and verify their checksums and signatures from the
 `Official Apache Download site <https://downloads.apache.org/airflow/providers/>`_
 
-* `The apache-airflow-providers-snowflake 6.3.1 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.3.1.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.3.1.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.3.1.tar.gz.sha512>`__)
-* `The apache-airflow-providers-snowflake 6.3.1 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.3.1-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.3.1-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.3.1-py3-none-any.whl.sha512>`__)
+* `The apache-airflow-providers-snowflake 6.4.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.4.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.4.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.4.0.tar.gz.sha512>`__)
+* `The apache-airflow-providers-snowflake 6.4.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.4.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.4.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_snowflake-6.4.0-py3-none-any.whl.sha512>`__)

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/provider.yaml

@@ -22,12 +22,13 @@ description: |
     `Snowflake <https://www.snowflake.com/>`__
 
 state: ready
-source-date-epoch: 1747133792
+source-date-epoch: 1749896974
 # Note that those versions are maintained by release manager - do not update them manually
 # with the exception of case where other provider in sources has >= new provider version.
 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
 # to be done in the same PR
 versions:
+  - 6.4.0
   - 6.3.1
   - 6.3.0
   - 6.2.2

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-snowflake"
-version = "6.3.1"
+version = "6.4.0"
 description = "Provider package apache-airflow-providers-snowflake for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -116,8 +116,8 @@ apache-airflow-providers-common-sql = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.3.1/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-snowflake/6.4.0/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/src/airflow/providers/snowflake/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "6.3.1"
+__version__ = "6.4.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/src/airflow/providers/snowflake/hooks/snowflake.py

@@ -136,6 +136,9 @@ class SnowflakeHook(DbApiHook):
                         "session_parameters": "session parameters",
                         "client_request_mfa_token": "client request mfa token",
                         "client_store_temporary_credential": "client store temporary credential (externalbrowser mode)",
+                        "grant_type": "refresh_token client_credentials",
+                        "token_endpoint": "token endpoint",
+                        "refresh_token": "refresh token",
                     },
                     indent=1,
                 ),
@@ -200,18 +203,32 @@ class SnowflakeHook(DbApiHook):
 
         return account_identifier
 
-    def get_oauth_token(self, conn_config: dict | None = None) -> str:
+    def get_oauth_token(
+        self,
+        conn_config: dict | None = None,
+        token_endpoint: str | None = None,
+        grant_type: str = "refresh_token",
+    ) -> str:
         """Generate temporary OAuth access token using refresh token in connection details."""
         if conn_config is None:
             conn_config = self._get_conn_params
 
-        url = f"https://{conn_config['account']}.snowflakecomputing.com/oauth/token-request"
+        url = token_endpoint or f"https://{conn_config['account']}.snowflakecomputing.com/oauth/token-request"
 
         data = {
-            "grant_type": "refresh_token",
-            "refresh_token": conn_config["refresh_token"],
+            "grant_type": grant_type,
             "redirect_uri": conn_config.get("redirect_uri", "https://localhost.com"),
         }
+
+        if grant_type == "refresh_token":
+            data |= {
+                "refresh_token": conn_config["refresh_token"],
+            }
+        elif grant_type == "client_credentials":
+            pass  # no setup necessary for client credentials grant.
+        else:
+            raise ValueError(f"Unknown grant_type: {grant_type}")
+
         response = requests.post(
             url,
             data=data,
@@ -226,7 +243,8 @@ class SnowflakeHook(DbApiHook):
         except requests.exceptions.HTTPError as e:  # pragma: no cover
             msg = f"Response: {e.response.content.decode()} Status Code: {e.response.status_code}"
             raise AirflowException(msg)
-        return response.json()["access_token"]
+        token = response.json()["access_token"]
+        return token
 
     @cached_property
     def _get_conn_params(self) -> dict[str, str | None]:
@@ -329,14 +347,21 @@ class SnowflakeHook(DbApiHook):
         if refresh_token:
             conn_config["refresh_token"] = refresh_token
             conn_config["authenticator"] = "oauth"
+
+        if conn_config.get("authenticator") == "oauth":
+            token_endpoint = self._get_field(extra_dict, "token_endpoint") or ""
             conn_config["client_id"] = conn.login
             conn_config["client_secret"] = conn.password
+            conn_config["token"] = self.get_oauth_token(
+                conn_config=conn_config,
+                token_endpoint=token_endpoint,
+                grant_type=extra_dict.get("grant_type", "refresh_token"),
+            )
+
             conn_config.pop("login", None)
             conn_config.pop("user", None)
             conn_config.pop("password", None)
 
-            conn_config["token"] = self.get_oauth_token(conn_config=conn_config)
-
         # configure custom target hostname and port, if specified
         snowflake_host = extra_dict.get("host")
         snowflake_port = extra_dict.get("port")

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py

@@ -137,6 +137,7 @@ class SnowflakeSqlApiHook(SnowflakeHook):
             When executing the statement, Snowflake replaces placeholders (? and :name) in
             the statement with these specified values.
         """
+        self.query_ids = []
         conn_config = self._get_conn_params
 
         req_id = uuid.uuid4()
@@ -222,14 +223,21 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         }
         return headers
 
-    def get_oauth_token(self, conn_config: dict[str, Any] | None = None) -> str:
+    def get_oauth_token(
+        self,
+        conn_config: dict[str, Any] | None = None,
+        token_endpoint: str | None = None,
+        grant_type: str = "refresh_token",
+    ) -> str:
         """Generate temporary OAuth access token using refresh token in connection details."""
         warnings.warn(
             "This method is deprecated. Please use `get_oauth_token` method from `SnowflakeHook` instead. ",
             AirflowProviderDeprecationWarning,
             stacklevel=2,
         )
-        return super().get_oauth_token(conn_config=conn_config)
+        return super().get_oauth_token(
+            conn_config=conn_config, token_endpoint=token_endpoint, grant_type=grant_type
+        )
 
     def get_request_url_header_params(self, query_id: str) -> tuple[dict[str, Any], dict[str, Any], str]:
         """

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/src/airflow/providers/snowflake/operators/snowflake.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 import time
 from collections.abc import Iterable, Mapping, Sequence
 from datetime import timedelta
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, SupportsAbs, cast
 
 from airflow.configuration import conf
@@ -390,6 +391,7 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
         self.bindings = bindings
         self.execute_async = False
         self.deferrable = deferrable
+        self.query_ids: list[str] = []
         if any([warehouse, database, role, schema, authenticator, session_parameters]):  # pragma: no cover
             hook_params = kwargs.pop("hook_params", {})  # pragma: no cover
             kwargs["hook_params"] = {
@@ -403,6 +405,16 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
             }
         super().__init__(conn_id=snowflake_conn_id, **kwargs)  # pragma: no cover
 
+    @cached_property
+    def _hook(self):
+        return SnowflakeSqlApiHook(
+            snowflake_conn_id=self.snowflake_conn_id,
+            token_life_time=self.token_life_time,
+            token_renewal_delta=self.token_renewal_delta,
+            deferrable=self.deferrable,
+            **self.hook_params,
+        )
+
     def execute(self, context: Context) -> None:
         """
         Make a POST API request to snowflake by using SnowflakeSQL and execute the query to get the ids.
@@ -410,13 +422,6 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
         By deferring the SnowflakeSqlApiTrigger class passed along with query ids.
         """
         self.log.info("Executing: %s", self.sql)
-        self._hook = SnowflakeSqlApiHook(
-            snowflake_conn_id=self.snowflake_conn_id,
-            token_life_time=self.token_life_time,
-            token_renewal_delta=self.token_renewal_delta,
-            deferrable=self.deferrable,
-            **self.hook_params,
-        )
         self.query_ids = self._hook.execute_query(
             self.sql,  # type: ignore[arg-type]
             statement_count=self.statement_count,
@@ -504,9 +509,11 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
                 msg = f"{event['status']}: {event['message']}"
                 raise AirflowException(msg)
             if "status" in event and event["status"] == "success":
-                hook = SnowflakeSqlApiHook(snowflake_conn_id=self.snowflake_conn_id)
-                query_ids = cast("list[str]", event["statement_query_ids"])
-                hook.check_query_output(query_ids)
+                self.query_ids = cast("list[str]", event["statement_query_ids"])
+                self._hook.check_query_output(self.query_ids)
                 self.log.info("%s completed successfully.", self.task_id)
+                # Re-assign query_ids to hook after coming back from deferral to be consistent for listeners.
+                if not self._hook.query_ids:
+                    self._hook.query_ids = self.query_ids
         else:
             self.log.info("%s completed successfully.", self.task_id)

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/src/airflow/providers/snowflake/transfers/copy_into_snowflake.py

@@ -27,6 +27,15 @@ from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 from airflow.providers.snowflake.utils.common import enclose_param
 
 
+def _validate_parameter(param_name: str, value: str | None) -> str | None:
+    """Validate that the parameter doesn't contain any invalid pattern."""
+    if value is None:
+        return None
+    if ";" in value:
+        raise ValueError(f"Invalid {param_name}: semicolons (;) not allowed.")
+    return value
+
+
 class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
     """
     Executes a COPY INTO command to load files from an external stage from clouds to Snowflake.
@@ -91,8 +100,8 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
     ):
         super().__init__(**kwargs)
         self.files = files
-        self.table = table
-        self.stage = stage
+        self.table = _validate_parameter("table", table)
+        self.stage = _validate_parameter("stage", stage)
         self.prefix = prefix
         self.file_format = file_format
         self.schema = schema
@@ -126,7 +135,7 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
         if self.schema:
             into = f"{self.schema}.{self.table}"
         else:
-            into = self.table
+            into = self.table  # type: ignore[assignment]
 
         if self.columns_array:
             into = f"{into}({', '.join(self.columns_array)})"

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/src/airflow/providers/snowflake/triggers/snowflake_trigger.py

@@ -74,7 +74,6 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
             self.token_renewal_delta,
         )
         try:
-            statement_query_ids: list[str] = []
             for query_id in self.query_ids:
                 while True:
                     statement_status = await self.get_query_status(query_id)
@@ -84,12 +83,10 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
                 if statement_status["status"] == "error":
                     yield TriggerEvent(statement_status)
                     return
-                if statement_status["status"] == "success":
-                    statement_query_ids.extend(statement_status["statement_handles"])
             yield TriggerEvent(
                 {
                     "status": "success",
-                    "statement_query_ids": statement_query_ids,
+                    "statement_query_ids": self.query_ids,
                 }
             )
         except Exception as e:

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/src/airflow/providers/snowflake/utils/openlineage.py

@@ -52,7 +52,15 @@ def fix_account_name(name: str) -> str:
             account, region = spl
             cloud = "aws"
         else:
-            account, region, cloud = spl
+            # region can easily get duplicated without crashing snowflake, so we need to handle that as well
+            # eg. account_locator.europe-west3.gcp.europe-west3.gcp will be ok for snowflake
+            account, region, cloud, *rest = spl
+            rest = [x for x in rest if x not in (region, cloud)]
+            if rest:  # Not sure what could be left here, but leaving this just in case
+                log.warning(
+                    "Unexpected parts found in Snowflake uri hostname and will be ignored by OpenLineage: %s",
+                    rest,
+                )
         return f"{account}.{region}.{cloud}"
 
     # Check for existing accounts with cloud names
@@ -72,13 +80,16 @@ def fix_snowflake_sqlalchemy_uri(uri: str) -> str:
     """
     Fix snowflake sqlalchemy connection URI to OpenLineage structure.
 
-    Snowflake sqlalchemy connection URI has following structure:
+    Snowflake sqlalchemy connection URI has the following structure:
     'snowflake://<user_login_name>:<password>@<account_identifier>/<database_name>/<schema_name>?warehouse=<warehouse_name>&role=<role_name>'
     We want account identifier normalized. It can have two forms:
-    - newer, in form of <organization>-<id>. In this case we want to do nothing.
-    - older, composed of <id>-<region>-<cloud> where region and cloud can be
+    - newer, in form of <organization_id>-<account_id>. In this case we want to do nothing.
+    - older, composed of <account_locator>.<region>.<cloud> where region and cloud can be
     optional in some cases. If <cloud> is omitted, it's AWS.
     If region and cloud are omitted, it's AWS us-west-1
+
+    Current doc on Snowflake account identifiers:
+    https://docs.snowflake.com/en/user-guide/admin-account-identifier
     """
     try:
         parts = urlparse(uri)

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/tests/unit/snowflake/decorators/test_snowpark.py

@@ -24,6 +24,7 @@ from unittest import mock
 import pytest
 
 from airflow.decorators import task
+from airflow.providers.snowflake.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils import timezone
 
 if TYPE_CHECKING:
@@ -156,7 +157,7 @@ class TestSnowparkDecorator:
         mock_snowflake_hook.return_value.get_snowpark_session.assert_called_once()
 
     @mock.patch("airflow.providers.snowflake.operators.snowpark.SnowflakeHook")
-    def test_snowpark_decorator_multiple_output(self, mock_snowflake_hook, dag_maker):
+    def test_snowpark_decorator_multiple_output(self, mock_snowflake_hook, dag_maker, request):
         @task.snowpark(
             task_id=TASK_ID,
             snowflake_conn_id=CONN_ID,
@@ -171,15 +172,23 @@ class TestSnowparkDecorator:
             assert session == mock_snowflake_hook.return_value.get_snowpark_session.return_value
             return {"a": 1, "b": "2"}
 
-        with dag_maker(dag_id=TEST_DAG_ID):
-            ret = func()
-
-        dr = dag_maker.create_dagrun()
-        ret.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
-        ti = dr.get_task_instances()[0]
-        assert ti.xcom_pull(key="a") == 1
-        assert ti.xcom_pull(key="b") == "2"
-        assert ti.xcom_pull() == {"a": 1, "b": "2"}
+        if AIRFLOW_V_3_0_PLUS:
+            run_task = request.getfixturevalue("run_task")
+            op = func().operator
+            run_task(task=op)
+            assert run_task.xcom.get(key="a") == 1
+            assert run_task.xcom.get(key="b") == "2"
+            assert run_task.xcom.get(key="return_value") == {"a": 1, "b": "2"}
+        else:
+            with dag_maker(dag_id=TEST_DAG_ID):
+                ret = func()
+
+            dr = dag_maker.create_dagrun()
+            ret.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+            ti = dr.get_task_instances()[0]
+            assert ti.xcom_pull(key="a") == 1
+            assert ti.xcom_pull(key="b") == "2"
+            assert ti.xcom_pull() == {"a": 1, "b": "2"}
         mock_snowflake_hook.assert_called_once()
         mock_snowflake_hook.return_value.get_snowpark_session.assert_called_once()
 

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/tests/unit/snowflake/hooks/test_snowflake.py

@@ -53,14 +53,13 @@ BASE_CONNECTION_KWARGS: dict = {
     },
 }
 
-CONN_PARAMS_OAUTH = {
+CONN_PARAMS_OAUTH_BASE = {
     "account": "airflow",
     "application": "AIRFLOW",
     "authenticator": "oauth",
     "database": "db",
     "client_id": "test_client_id",
     "client_secret": "test_client_pw",
-    "refresh_token": "secrettoken",
     "region": "af_region",
     "role": "af_role",
     "schema": "public",
@@ -68,6 +67,8 @@ CONN_PARAMS_OAUTH = {
     "warehouse": "af_wh",
 }
 
+CONN_PARAMS_OAUTH = CONN_PARAMS_OAUTH_BASE | {"refresh_token": "secrettoken"}
+
 
 @pytest.fixture
 def unencrypted_temporary_private_key(tmp_path: Path) -> Path:
@@ -559,6 +560,112 @@ class TestPytestSnowflakeHook:
         assert "region" in conn_params_extra_keys
         assert "account" in conn_params_extra_keys
 
+    @mock.patch("requests.post")
+    @mock.patch(
+        "airflow.providers.snowflake.hooks.snowflake.SnowflakeHook._get_conn_params",
+        new_callable=PropertyMock,
+    )
+    def test_get_conn_params_should_support_oauth_with_token_endpoint(
+        self, mock_get_conn_params, requests_post
+    ):
+        requests_post.return_value = Mock(
+            status_code=200,
+            json=lambda: {
+                "access_token": "supersecretaccesstoken",
+                "expires_in": 600,
+                "refresh_token": "secrettoken",
+                "token_type": "Bearer",
+                "username": "test_user",
+            },
+        )
+        connection_kwargs = {
+            **BASE_CONNECTION_KWARGS,
+            "login": "test_client_id",
+            "password": "test_client_secret",
+            "extra": {
+                "database": "db",
+                "account": "airflow",
+                "warehouse": "af_wh",
+                "region": "af_region",
+                "role": "af_role",
+                "refresh_token": "secrettoken",
+                "authenticator": "oauth",
+                "token_endpoint": "https://www.example.com/oauth/token",
+            },
+        }
+        mock_get_conn_params.return_value = connection_kwargs
+        with mock.patch.dict("os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()):
+            hook = SnowflakeHook(snowflake_conn_id="test_conn")
+            conn_params = hook._get_conn_params
+
+        conn_params_keys = conn_params.keys()
+        conn_params_extra = conn_params.get("extra", {})
+        conn_params_extra_keys = conn_params_extra.keys()
+
+        assert "authenticator" in conn_params_extra_keys
+        assert conn_params_extra["authenticator"] == "oauth"
+        assert conn_params_extra["token_endpoint"] == "https://www.example.com/oauth/token"
+
+        assert "user" not in conn_params_keys
+        assert "password" in conn_params_keys
+        assert "refresh_token" in conn_params_extra_keys
+        # Mandatory fields to generate account_identifier `https://<account>.<region>`
+        assert "region" in conn_params_extra_keys
+        assert "account" in conn_params_extra_keys
+
+    @mock.patch("requests.post")
+    @mock.patch(
+        "airflow.providers.snowflake.hooks.snowflake.SnowflakeHook._get_conn_params",
+        new_callable=PropertyMock,
+    )
+    def test_get_conn_params_should_support_oauth_with_client_credentials(
+        self, mock_get_conn_params, requests_post
+    ):
+        requests_post.return_value = Mock(
+            status_code=200,
+            json=lambda: {
+                "access_token": "supersecretaccesstoken",
+                "expires_in": 600,
+                "refresh_token": "secrettoken",
+                "token_type": "Bearer",
+                "username": "test_user",
+            },
+        )
+        connection_kwargs = {
+            **BASE_CONNECTION_KWARGS,
+            "login": "test_client_id",
+            "password": "test_client_secret",
+            "extra": {
+                "database": "db",
+                "account": "airflow",
+                "warehouse": "af_wh",
+                "region": "af_region",
+                "role": "af_role",
+                "authenticator": "oauth",
+                "token_endpoint": "https://www.example.com/oauth/token",
+                "grant_type": "client_credentials",
+            },
+        }
+        mock_get_conn_params.return_value = connection_kwargs
+        with mock.patch.dict("os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()):
+            hook = SnowflakeHook(snowflake_conn_id="test_conn")
+            conn_params = hook._get_conn_params
+
+        conn_params_keys = conn_params.keys()
+        conn_params_extra = conn_params.get("extra", {})
+        conn_params_extra_keys = conn_params_extra.keys()
+
+        assert "authenticator" in conn_params_extra_keys
+        assert conn_params_extra["authenticator"] == "oauth"
+        assert conn_params_extra["grant_type"] == "client_credentials"
+
+        assert "user" not in conn_params_keys
+        assert "password" in conn_params_keys
+        assert "refresh_token" not in conn_params_extra_keys
+        # Mandatory fields to generate account_identifier `https://<account>.<region>`
+        assert "region" in conn_params_extra_keys
+        assert "account" in conn_params_extra_keys
+
     def test_should_add_partner_info(self):
         with mock.patch.dict(
             "os.environ",
@@ -917,3 +1024,31 @@ class TestPytestSnowflakeHook:
             headers={"Content-Type": "application/x-www-form-urlencoded"},
             auth=basic_auth,
         )
+
+    @mock.patch("airflow.providers.snowflake.hooks.snowflake.HTTPBasicAuth")
+    @mock.patch("requests.post")
+    @mock.patch(
+        "airflow.providers.snowflake.hooks.snowflake.SnowflakeHook._get_conn_params",
+        new_callable=PropertyMock,
+    )
+    def test_get_oauth_token_with_token_endpoint(self, mock_conn_param, requests_post, mock_auth):
+        """Test get_oauth_token method makes the right http request"""
+        basic_auth = {"Authorization": "Basic usernamepassword"}
+        token_endpoint = "https://example.com/oauth/token"
+        mock_conn_param.return_value = CONN_PARAMS_OAUTH
+        requests_post.return_value.status_code = 200
+        mock_auth.return_value = basic_auth
+
+        hook = SnowflakeHook(snowflake_conn_id="mock_conn_id")
+        hook.get_oauth_token(conn_config=CONN_PARAMS_OAUTH, token_endpoint=token_endpoint)
+
+        requests_post.assert_called_once_with(
+            token_endpoint,
+            data={
+                "grant_type": "refresh_token",
+                "refresh_token": CONN_PARAMS_OAUTH["refresh_token"],
+                "redirect_uri": "https://localhost.com",
+            },
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            auth=basic_auth,
+        )

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/tests/unit/snowflake/hooks/test_snowflake_sql_api.py

@@ -203,6 +203,46 @@ class TestSnowflakeSqlApiHook:
         query_ids = hook.execute_query(sql, statement_count)
         assert query_ids == expected_query_ids
 
+    @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.requests")
+    @mock.patch(
+        "airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook._get_conn_params",
+        new_callable=PropertyMock,
+    )
+    @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_headers")
+    def test_execute_query_multiple_times_give_fresh_query_ids_each_time(
+        self, mock_get_header, mock_conn_param, mock_requests
+    ):
+        """Test execute_query method, run query by mocking post request method and return the query ids"""
+        sql, statement_count, expected_response, expected_query_ids = (
+            SQL_MULTIPLE_STMTS,
+            4,
+            {"statementHandles": ["uuid2", "uuid3"]},
+            ["uuid2", "uuid3"],
+        )
+
+        mock_requests.codes.ok = 200
+        mock_requests.post.side_effect = [
+            create_successful_response_mock(expected_response),
+        ]
+        status_code_mock = mock.PropertyMock(return_value=200)
+        type(mock_requests.post.return_value).status_code = status_code_mock
+
+        hook = SnowflakeSqlApiHook("mock_conn_id")
+        query_ids = hook.execute_query(sql, statement_count)
+        assert query_ids == expected_query_ids
+
+        sql, statement_count, expected_response, expected_query_ids = (
+            SINGLE_STMT,
+            1,
+            {"statementHandle": "uuid"},
+            ["uuid"],
+        )
+        mock_requests.post.side_effect = [
+            create_successful_response_mock(expected_response),
+        ]
+        query_ids = hook.execute_query(sql, statement_count)
+        assert query_ids == expected_query_ids
+
     @pytest.mark.parametrize(
         "sql,statement_count,expected_response, expected_query_ids",
         [(SINGLE_STMT, 1, {"statementHandle": "uuid"}, ["uuid"])],

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/tests/unit/snowflake/operators/test_snowflake.py

@@ -332,6 +332,48 @@ class TestSnowflakeSqlApiOperator:
             operator.execute_complete(context=None, event=mock_event)
         mock_log_info.assert_called_with("%s completed successfully.", TASK_ID)
 
+    @pytest.mark.parametrize(
+        "mock_event",
+        [
+            None,
+            ({"status": "success", "statement_query_ids": ["uuid", "uuid"]}),
+        ],
+    )
+    @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.check_query_output")
+    def test_snowflake_sql_api_execute_complete_reassigns_query_ids(self, mock_conn, mock_event):
+        """Tests execute_complete assert with successful message"""
+
+        operator = SnowflakeSqlApiOperator(
+            task_id=TASK_ID,
+            snowflake_conn_id=CONN_ID,
+            sql=SQL_MULTIPLE_STMTS,
+            statement_count=4,
+            deferrable=True,
+        )
+        expected_query_ids = mock_event["statement_query_ids"] if mock_event else []
+
+        assert operator.query_ids == []
+        assert operator._hook.query_ids == []
+
+        operator.execute_complete(context=None, event=mock_event)
+
+        assert operator.query_ids == expected_query_ids
+        assert operator._hook.query_ids == expected_query_ids
+
+    def test_snowflake_sql_api_caches_hook(self):
+        """Tests execute_complete assert with successful message"""
+
+        operator = SnowflakeSqlApiOperator(
+            task_id=TASK_ID,
+            snowflake_conn_id=CONN_ID,
+            sql=SQL_MULTIPLE_STMTS,
+            statement_count=4,
+            deferrable=True,
+        )
+        hook1 = operator._hook
+        hook2 = operator._hook
+        assert hook1 is hook2
+
     @mock.patch("airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiOperator.defer")
     def test_snowflake_sql_api_execute_operator_failed_before_defer(
         self, mock_defer, mock_execute_query, mock_get_sql_api_query_status

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/tests/unit/snowflake/transfers/test_copy_into_snowflake.py

@@ -549,3 +549,26 @@ class TestCopyFromExternalStageToSnowflake:
             },
             job_facets={"sql": SQLJobFacet(query=expected_sql)},
         )
+
+    def test_init_with_invalid_parameters(self):
+        # Test by passing invalid parameters to the operator
+        import re
+
+        with pytest.raises(ValueError, match=re.escape("Invalid table: semicolons (;) not allowed.")):
+            _ = CopyFromExternalStageToSnowflakeOperator(
+                task_id="test",
+                table="table; some_new_table",
+                stage="stage",
+                database="",
+                schema="schema",
+                file_format="CSV",
+            )
+        with pytest.raises(ValueError, match=re.escape("Invalid stage: semicolons (;) not allowed.")):
+            _ = CopyFromExternalStageToSnowflakeOperator(
+                task_id="test",
+                table="table",
+                stage="stage; some_new_stabge",
+                database="",
+                schema="schema",
+                file_format="CSV",
+            )

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/tests/unit/snowflake/triggers/test_snowflake.py

@@ -30,12 +30,13 @@ POLL_INTERVAL = 1.0
 LIFETIME = timedelta(minutes=59)
 RENEWAL_DELTA = timedelta(minutes=54)
 MODULE = "airflow.providers.snowflake"
+QUERY_IDS = ["uuid"]
 
 
 class TestSnowflakeSqlApiTrigger:
     TRIGGER = SnowflakeSqlApiTrigger(
         poll_interval=POLL_INTERVAL,
-        query_ids=["uuid"],
+        query_ids=QUERY_IDS,
         snowflake_conn_id="test_conn",
         token_life_time=LIFETIME,
         token_renewal_delta=RENEWAL_DELTA,
@@ -82,8 +83,8 @@ class TestSnowflakeSqlApiTrigger:
         Test SnowflakeSqlApiTrigger run method with success status and mock the get_sql_api_query_status
          result and  get_query_status to False.
         """
-        mock_get_query_status.return_value = {"status": "success", "statement_handles": ["uuid", "uuid1"]}
         statement_query_ids = ["uuid", "uuid1"]
+        mock_get_query_status.return_value = {"status": "success", "statement_handles": statement_query_ids}
         mock_get_sql_api_query_status_async.return_value = {
             "message": "Statement executed successfully.",
             "status": "success",
@@ -92,7 +93,7 @@ class TestSnowflakeSqlApiTrigger:
 
         generator = self.TRIGGER.run()
         actual = await generator.asend(None)
-        assert TriggerEvent({"status": "success", "statement_query_ids": statement_query_ids}) == actual
+        assert TriggerEvent({"status": "success", "statement_query_ids": QUERY_IDS}) == actual
 
     @pytest.mark.asyncio
     @mock.patch(f"{MODULE}.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")

{apache_airflow_providers_snowflake-6.3.1 → apache_airflow_providers_snowflake-6.4.0}/tests/unit/snowflake/utils/test_openlineage.py

@@ -93,6 +93,9 @@ def test_snowflake_sqlite_account_urls(source, target):
         ("xy12345", "xy12345.us-west-1.aws"),  # No '-' or '_' in name
         ("xy12345.us-west-1.aws", "xy12345.us-west-1.aws"),  # Already complete locator
         ("xy12345.us-west-2.gcp", "xy12345.us-west-2.gcp"),  # Already complete locator for GCP
+        ("xy12345.us-west-2.gcp.us-west-2.gcp", "xy12345.us-west-2.gcp"),  # Duplicated region
+        ("xy12345.us-west-2.gcp.us-west-2.gcp.us-west-2.gcp", "xy12345.us-west-2.gcp"),  # Triple region
+        ("xy12345.us-west-2.gcp.some_random_part", "xy12345.us-west-2.gcp"),  # Suffix to locator, ignored
         ("xy12345aws", "xy12345aws.us-west-1.aws"),  # AWS without '-' or '_'
         ("xy12345-aws", "xy12345-aws"),  # AWS with '-'
         ("xy12345_gcp-europe-west1", "xy12345.europe-west1.gcp"),  # GCP with '_'