apache-airflow-providers-snowflake 6.3.0__py3-none-any.whl → 6.8.0rc1__py3-none-any.whl

airflow/providers/snowflake/__init__.py

@@ -29,11 +29,11 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "6.3.0"
+__version__ = "6.8.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.10.0"
+    "2.11.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-snowflake:{__version__}` needs Apache Airflow 2.10.0+"
+        f"The package `apache-airflow-providers-snowflake:{__version__}` needs Apache Airflow 2.11.0+"
     )

airflow/providers/snowflake/decorators/snowpark.py

@@ -17,22 +17,12 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Callable
-
-from airflow.providers.snowflake.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk.bases.decorator import DecoratedOperator, task_decorator_factory
-else:
-    from airflow.decorators.base import DecoratedOperator, task_decorator_factory  # type: ignore[no-redef]
+from collections.abc import Callable, Sequence
 
+from airflow.providers.common.compat.sdk import DecoratedOperator, TaskDecorator, task_decorator_factory
 from airflow.providers.snowflake.operators.snowpark import SnowparkOperator
 from airflow.providers.snowflake.utils.snowpark import inject_session_into_op_kwargs
 
-if TYPE_CHECKING:
-    from airflow.sdk.bases.decorator import TaskDecorator
-
 
 class _SnowparkDecoratedOperator(DecoratedOperator, SnowparkOperator):
     """

airflow/providers/snowflake/get_provider_info.py

@@ -73,11 +73,13 @@ def get_provider_info():
                 "source-integration-name": "Google Cloud Storage (GCS)",
                 "target-integration-name": "Snowflake",
                 "python-module": "airflow.providers.snowflake.transfers.copy_into_snowflake",
+                "how-to-guide": "/docs/apache-airflow-providers-snowflake/operators/copy_into_snowflake.rst",
             },
             {
                 "source-integration-name": "Microsoft Azure Blob Storage",
                 "target-integration-name": "Snowflake",
                 "python-module": "airflow.providers.snowflake.transfers.copy_into_snowflake",
+                "how-to-guide": "/docs/apache-airflow-providers-snowflake/operators/copy_into_snowflake.rst",
             },
         ],
         "connection-types": [
@@ -92,4 +94,18 @@ def get_provider_info():
                 "python-modules": ["airflow.providers.snowflake.triggers.snowflake_trigger"],
             }
         ],
+        "config": {
+            "snowflake": {
+                "description": "Configuration for Snowflake hooks and operators.\n",
+                "options": {
+                    "azure_oauth_scope": {
+                        "description": "The scope to use while retrieving OAuth token for Snowflake from Azure Entra authentication.\n",
+                        "version_added": "6.6.0",
+                        "type": "string",
+                        "example": None,
+                        "default": "api://snowflake_oauth_server/.default",
+                    }
+                },
+            }
+        },
     }

airflow/providers/snowflake/hooks/snowflake.py

@@ -19,12 +19,12 @@ from __future__ import annotations
 
 import base64
 import os
-from collections.abc import Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 from contextlib import closing, contextmanager
 from functools import cached_property
 from io import StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, overload
+from typing import TYPE_CHECKING, Any, TypeVar, overload
 from urllib.parse import urlparse
 
 import requests
@@ -36,7 +36,9 @@ from snowflake.connector import DictCursor, SnowflakeConnection, util_text
 from snowflake.sqlalchemy import URL
 from sqlalchemy import create_engine
 
-from airflow.exceptions import AirflowException
+from airflow.configuration import conf
+from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.common.compat.sdk import AirflowException, Connection
 from airflow.providers.common.sql.hooks.handlers import return_single_query_results
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.snowflake.utils.openlineage import fix_snowflake_sqlalchemy_uri
@@ -94,6 +96,7 @@ class SnowflakeHook(DbApiHook):
     hook_name = "Snowflake"
     supports_autocommit = True
     _test_connection_sql = "select 1"
+    default_azure_oauth_scope = "api://snowflake_oauth_server/.default"
 
     @classmethod
     def get_connection_form_widgets(cls) -> dict[str, Any]:
@@ -136,6 +139,10 @@ 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",
+                        "scope": "scope",
                     },
                     indent=1,
                 ),
@@ -200,18 +207,37 @@ 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"),
         }
+
+        scope = conn_config.get("scope")
+
+        if scope:
+            data["scope"] = scope
+
+        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 +252,40 @@ 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
+
+    def get_azure_oauth_token(self, azure_conn_id: str) -> str:
+        """
+        Generate OAuth access token using Azure connection id.
+
+        This uses AzureBaseHook on the connection id to retrieve the token. Scope for the OAuth token can be
+        set in the config option ``azure_oauth_scope`` under the section ``[snowflake]``.
+
+        :param azure_conn_id: The connection id for the Azure connection that will be used to fetch the token.
+        :raises AttributeError: If AzureBaseHook does not have a get_token method which happens when
+            package apache-airflow-providers-microsoft-azure<12.8.0.
+        :returns: The OAuth access token string.
+        """
+        if TYPE_CHECKING:
+            from airflow.providers.microsoft.azure.hooks.azure_base import AzureBaseHook
+
+        try:
+            azure_conn = Connection.get(azure_conn_id)
+        except AttributeError:
+            azure_conn = Connection.get_connection_from_secrets(azure_conn_id)  # type: ignore[attr-defined]
+        try:
+            azure_base_hook: AzureBaseHook = azure_conn.get_hook()
+        except TypeError as e:
+            if "required positional argument: 'sdk_client'" in str(e):
+                raise AirflowOptionalProviderFeatureException(
+                    "Getting azure token is not supported by current version of 'AzureBaseHook'. "
+                    "Please upgrade apache-airflow-providers-microsoft-azure>=12.8.0"
+                ) from e
+            raise
+        scope = conf.get("snowflake", "azure_oauth_scope", fallback=self.default_azure_oauth_scope)
+        token = azure_base_hook.get_token(scope).token
+        return token
 
     @cached_property
     def _get_conn_params(self) -> dict[str, str | None]:
@@ -235,7 +294,7 @@ class SnowflakeHook(DbApiHook):
 
         This is used in ``get_uri()`` and ``get_connection()``.
         """
-        conn = self.get_connection(self.snowflake_conn_id)  # type: ignore[attr-defined]
+        conn = self.get_connection(self.get_conn_id())
         extra_dict = conn.extra_dejson
         account = self._get_field(extra_dict, "account") or ""
         warehouse = self._get_field(extra_dict, "warehouse") or ""
@@ -329,14 +388,26 @@ class SnowflakeHook(DbApiHook):
         if refresh_token:
             conn_config["refresh_token"] = refresh_token
             conn_config["authenticator"] = "oauth"
-            conn_config["client_id"] = conn.login
-            conn_config["client_secret"] = conn.password
+
+        if conn_config.get("authenticator") == "oauth":
+            if extra_dict.get("azure_conn_id"):
+                conn_config["token"] = self.get_azure_oauth_token(extra_dict["azure_conn_id"])
+            else:
+                token_endpoint = self._get_field(extra_dict, "token_endpoint") or ""
+                conn_config["scope"] = self._get_field(extra_dict, "scope")
+                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")
@@ -436,7 +507,7 @@ class SnowflakeHook(DbApiHook):
     def get_autocommit(self, conn):
         return getattr(conn, "autocommit_mode", False)
 
-    @overload  # type: ignore[override]
+    @overload
     def run(
         self,
         sql: str | Iterable[str],
@@ -519,16 +590,16 @@ class SnowflakeHook(DbApiHook):
                 results = []
                 for sql_statement in sql_list:
                     self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters)
-                    self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
+                    self._run_command(cur, sql_statement, parameters)
 
                     if handler is not None:
-                        result = self._make_common_data_structure(handler(cur))  # type: ignore[attr-defined]
+                        result = self._make_common_data_structure(handler(cur))
                         if return_single_query_results(sql, return_last, split_statements):
                             _last_result = result
                             _last_description = cur.description
                         else:
                             results.append(result)
-                            self.descriptions.append(cur.description)  # type: ignore[has-type]
+                            self.descriptions.append(cur.description)
 
                     query_id = cur.sfqid
                     self.log.info("Rows affected: %s", cur.rowcount)
@@ -592,10 +663,9 @@ class SnowflakeHook(DbApiHook):
 
     def get_openlineage_database_specific_lineage(self, task_instance) -> OperatorLineage | None:
         """
-        Generate OpenLineage metadata for a Snowflake task instance based on executed query IDs.
+        Emit separate OpenLineage events for each Snowflake query, based on executed query IDs.
 
-        If a single query ID is present, attach an `ExternalQueryRunFacet` to the lineage metadata.
-        If multiple query IDs are present, emits separate OpenLineage events for each query.
+        If a single query ID is present, also add an `ExternalQueryRunFacet` to the returned lineage metadata.
 
         Note that `get_openlineage_database_specific_lineage` is usually called after task's execution,
         so if multiple query IDs are present, both START and COMPLETE event for each query will be emitted
@@ -616,13 +686,22 @@ class SnowflakeHook(DbApiHook):
         )
 
         if not self.query_ids:
-            self.log.debug("openlineage: no snowflake query ids found.")
+            self.log.info("OpenLineage could not find snowflake query ids.")
             return None
 
         self.log.debug("openlineage: getting connection to get database info")
         connection = self.get_connection(self.get_conn_id())
         namespace = SQLParser.create_namespace(self.get_openlineage_database_info(connection))
 
+        self.log.info("Separate OpenLineage events will be emitted for each query_id.")
+        emit_openlineage_events_for_snowflake_queries(
+            task_instance=task_instance,
+            hook=self,
+            query_ids=self.query_ids,
+            query_for_extra_metadata=True,
+            query_source_namespace=namespace,
+        )
+
         if len(self.query_ids) == 1:
             self.log.debug("Attaching ExternalQueryRunFacet with single query_id to OpenLineage event.")
             return OperatorLineage(
@@ -633,20 +712,4 @@ class SnowflakeHook(DbApiHook):
                 }
             )
 
-        self.log.info("Multiple query_ids found. Separate OpenLineage event will be emitted for each query.")
-        try:
-            from airflow.providers.openlineage.utils.utils import should_use_external_connection
-
-            use_external_connection = should_use_external_connection(self)
-        except ImportError:
-            # OpenLineage provider release < 1.8.0 - we always use connection
-            use_external_connection = True
-
-        emit_openlineage_events_for_snowflake_queries(
-            query_ids=self.query_ids,
-            query_source_namespace=namespace,
-            task_instance=task_instance,
-            hook=self if use_external_connection else None,
-        )
-
         return None