apache-airflow-providers-snowflake 6.3.1rc1__py3-none-any.whl → 6.4.0rc1__py3-none-any.whl

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"

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")

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]:
         """

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)

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)})"

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:

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.1rc1.dist-info → apache_airflow_providers_snowflake-6.4.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-snowflake
-Version: 6.3.1rc1
+Version: 6.4.0rc1
 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.0rc1 ; extra == "openlineage"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-snowflake/6.3.1/changelog.html
-Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-snowflake/6.3.1
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-snowflake/6.4.0/changelog.html
+Project-URL: Documentation, https://airflow.staged.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.1rc1.dist-info → apache_airflow_providers_snowflake-6.4.0rc1.dist-info}/RECORD

@@ -1,25 +1,25 @@
 airflow/providers/snowflake/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
-airflow/providers/snowflake/__init__.py,sha256=QsK4vg_MkprBynxoW0TPaDLqVXgU_0-k0VW6Zabo7gk,1498
+airflow/providers/snowflake/__init__.py,sha256=hfjXA59cpm_yb00p5Y7jNTIBCe3BPNoYYXkAVAcF_c4,1498
 airflow/providers/snowflake/get_provider_info.py,sha256=NdNRMfulBbpD-I4yFRr8U533m9djD18ijEMvuxOp4_g,3875
 airflow/providers/snowflake/version_compat.py,sha256=j5PCtXvZ71aBjixu-EFTNtVDPsngzzs7os0ZQDgFVDk,1536
 airflow/providers/snowflake/decorators/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/snowflake/decorators/snowpark.py,sha256=tKXOjP8m8SEIu0jx2KSrd0n3jGMaIKDOwG2lMkvk3cI,5523
 airflow/providers/snowflake/hooks/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/hooks/snowflake.py,sha256=9OH16CYnnJ0-ayAg1D7OdZusEf5lSGjQurWifptp97k,28025
-airflow/providers/snowflake/hooks/snowflake_sql_api.py,sha256=-J0mPcdDc9wbB7DcnZfnXJN7H62nbR_NK5WQJxeKZjE,14532
+airflow/providers/snowflake/hooks/snowflake.py,sha256=t-sukzbQ1OYMeyZBDrJ9s9DuJbnDZFJKhBMZn_mQLgY,28933
+airflow/providers/snowflake/hooks/snowflake_sql_api.py,sha256=mREZ0nHc6L-9YSHZARYpgqrETqzVJ3Q6EfbWtEy5TV4,14745
 airflow/providers/snowflake/operators/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/operators/snowflake.py,sha256=5MisB-bKqUFM9t5Ky913UqewoHlq3k3mCv4bnc-VY7g,22657
+airflow/providers/snowflake/operators/snowflake.py,sha256=Abu0MVsUPwVxfDVNYn5OtwVoUOhQanUp-YSFtLdcn6c,22915
 airflow/providers/snowflake/operators/snowpark.py,sha256=Wt3wzcsja0ed4q2KE9WyL74XH6mUVSPNZvcCHWEHQtc,5815
 airflow/providers/snowflake/transfers/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/transfers/copy_into_snowflake.py,sha256=UjbznjbK-QWN071ZFMvBHZXoFddMo0vQFK-7VLv3amo,13191
+airflow/providers/snowflake/transfers/copy_into_snowflake.py,sha256=bXmqkNwthJqUo65DsI_pC3mwk_V_Iikwi646oRCyWus,13590
 airflow/providers/snowflake/triggers/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/triggers/snowflake_trigger.py,sha256=38tkByMyjbVbSt-69YL8EzRBQT4rhwuOKHgbwHfULL0,4250
+airflow/providers/snowflake/triggers/snowflake_trigger.py,sha256=QXNLijmtZI7NIdPtOwbvS-4ohgrm8RV_jaBKvekosHQ,4051
 airflow/providers/snowflake/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/snowflake/utils/common.py,sha256=DG-KLy2KpZWAqZqm_XIECm8lmdoUlzwkXv9onmkQThc,1644
-airflow/providers/snowflake/utils/openlineage.py,sha256=QjbN76qjboTvpQZtoi0g7s3R9LwutRtD7HZ3DEVLbyY,14372
+airflow/providers/snowflake/utils/openlineage.py,sha256=dr57b0fidPo7A451UE6s6d3PnOf4dOv3iM5aoyE4oBI,15067
 airflow/providers/snowflake/utils/snowpark.py,sha256=9kzWRkdgoNQ8f3Wnr92LdZylMpcpRasxefpOXrM30Cw,1602
 airflow/providers/snowflake/utils/sql_api_generate_jwt.py,sha256=9mR-vHIquv60tfAni87f6FAjKsiRHUDDrsVhzw4M9vM,6762
-apache_airflow_providers_snowflake-6.3.1rc1.dist-info/entry_points.txt,sha256=bCrl5J1PXUMzbgnrKYho61rkbL2gHRT4I6f_1jlxAX4,105
-apache_airflow_providers_snowflake-6.3.1rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-apache_airflow_providers_snowflake-6.3.1rc1.dist-info/METADATA,sha256=0pkDl6pLL8SvxNML46vEzwwQ1IP7TzRKYlbJ2MtTkzI,6242
-apache_airflow_providers_snowflake-6.3.1rc1.dist-info/RECORD,,
+apache_airflow_providers_snowflake-6.4.0rc1.dist-info/entry_points.txt,sha256=bCrl5J1PXUMzbgnrKYho61rkbL2gHRT4I6f_1jlxAX4,105
+apache_airflow_providers_snowflake-6.4.0rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_snowflake-6.4.0rc1.dist-info/METADATA,sha256=mDqBewdnQi0Pa9w3DNrTvh2i0csPs95pMlAsZDp1kX8,6242
+apache_airflow_providers_snowflake-6.4.0rc1.dist-info/RECORD,,