apache-airflow-providers-snowflake 5.3.1rc1__py3-none-any.whl → 5.4.0__py3-none-any.whl

airflow/providers/snowflake/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "5.3.1"
+__version__ = "5.4.0"
 
 try:
     from airflow import __version__ as airflow_version

airflow/providers/snowflake/get_provider_info.py

@@ -28,8 +28,9 @@ def get_provider_info():
         "name": "Snowflake",
         "description": "`Snowflake <https://www.snowflake.com/>`__\n",
         "state": "ready",
-        "source-date-epoch": 1707636525,
+        "source-date-epoch": 1712666514,
         "versions": [
+            "5.4.0",
             "5.3.1",
             "5.3.0",
             "5.2.1",

airflow/providers/snowflake/hooks/snowflake.py

@@ -256,6 +256,15 @@ class SnowflakeHook(DbApiHook):
             conn_config["private_key"] = pkb
             conn_config.pop("password", None)
 
+        refresh_token = self._get_field(extra_dict, "refresh_token") or ""
+        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
+            conn_config.pop("login", None)
+            conn_config.pop("password", None)
+
         return conn_config
 
     def get_uri(self) -> str:
@@ -312,8 +321,7 @@ class SnowflakeHook(DbApiHook):
         split_statements: bool = ...,
         return_last: bool = ...,
         return_dictionaries: bool = ...,
-    ) -> None:
-        ...
+    ) -> None: ...
 
     @overload
     def run(
@@ -325,8 +333,7 @@ class SnowflakeHook(DbApiHook):
         split_statements: bool = ...,
         return_last: bool = ...,
         return_dictionaries: bool = ...,
-    ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
-        ...
+    ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ...
 
     def run(
         self,
@@ -341,10 +348,8 @@ class SnowflakeHook(DbApiHook):
         """Run a command or list of commands.
 
         Pass a list of SQL statements to the SQL parameter to get them to
-        execute sequentially. The variable ``execution_info`` is returned so
-        that it can be used in the Operators to modify the behavior depending on
-        the result of the query (i.e fail the operator if the copy has processed
-        0 files).
+        execute sequentially. The result of the queries is returned if the
+        ``handler`` callable is set.
 
         :param sql: The SQL string to be executed with possibly multiple
             statements, or a list of sql statements to execute

airflow/providers/snowflake/hooks/snowflake_sql_api.py

@@ -25,6 +25,7 @@ import aiohttp
 import requests
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
+from requests.auth import HTTPBasicAuth
 
 from airflow.exceptions import AirflowException
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
@@ -39,8 +40,9 @@ class SnowflakeSqlApiHook(SnowflakeHook):
     poll to check the status of the execution of a statement. Fetch query results asynchronously.
 
     This hook requires the snowflake_conn_id connection. This hooks mainly uses account, schema, database,
-     warehouse, private_key_file or private_key_content field must be setup in the connection. Other inputs
-      can be defined in the connection or hook instantiation.
+    warehouse, and an authentication mechanism from one of below:
+    1. JWT Token generated from private_key_file or private_key_content. Other inputs can be defined in the connection or hook instantiation.
+    2. OAuth Token generated from the refresh_token, client_id and client_secret specified in the connection
 
     :param snowflake_conn_id: Reference to
         :ref:`Snowflake connection id<howto/connection:snowflake>`
@@ -81,6 +83,17 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         super().__init__(snowflake_conn_id, *args, **kwargs)
         self.private_key: Any = None
 
+    @property
+    def account_identifier(self) -> str:
+        """Returns snowflake account identifier."""
+        conn_config = self._get_conn_params()
+        account_identifier = f"https://{conn_config['account']}"
+
+        if conn_config["region"]:
+            account_identifier += f".{conn_config['region']}"
+
+        return account_identifier
+
     def get_private_key(self) -> None:
         """Get the private key from snowflake connection."""
         conn = self.get_connection(self.snowflake_conn_id)
@@ -137,10 +150,7 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         conn_config = self._get_conn_params()
 
         req_id = uuid.uuid4()
-        url = (
-            f"https://{conn_config['account']}.{conn_config['region']}"
-            f".snowflakecomputing.com/api/v2/statements"
-        )
+        url = f"{self.account_identifier}.snowflakecomputing.com/api/v2/statements"
         params: dict[str, Any] | None = {"requestId": str(req_id), "async": True, "pageSize": 10}
         headers = self.get_headers()
         if bindings is None:
@@ -175,12 +185,27 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         return self.query_ids
 
     def get_headers(self) -> dict[str, Any]:
-        """Form JWT Token and header based on the private key, and connection details."""
+        """Form auth headers based on either OAuth token or JWT token from private key."""
+        conn_config = self._get_conn_params()
+
+        # Use OAuth if refresh_token and client_id and client_secret are provided
+        if all(
+            [conn_config.get("refresh_token"), conn_config.get("client_id"), conn_config.get("client_secret")]
+        ):
+            oauth_token = self.get_oauth_token()
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {oauth_token}",
+                "Accept": "application/json",
+                "User-Agent": "snowflakeSQLAPI/1.0",
+                "X-Snowflake-Authorization-Token-Type": "OAUTH",
+            }
+            return headers
+
+        # Alternatively, get the JWT token from the connection details and the private key
         if not self.private_key:
             self.get_private_key()
-        conn_config = self._get_conn_params()
 
-        # Get the JWT token from the connection details and the private key
         token = JWTGenerator(
             conn_config["account"],  # type: ignore[arg-type]
             conn_config["user"],  # type: ignore[arg-type]
@@ -198,20 +223,41 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         }
         return headers
 
+    def get_oauth_token(self) -> str:
+        """Generate temporary OAuth access token using refresh token in connection details."""
+        conn_config = self._get_conn_params()
+        url = f"{self.account_identifier}.snowflakecomputing.com/oauth/token-request"
+        data = {
+            "grant_type": "refresh_token",
+            "refresh_token": conn_config["refresh_token"],
+            "redirect_uri": conn_config.get("redirect_uri", "https://localhost.com"),
+        }
+        response = requests.post(
+            url,
+            data=data,
+            headers={
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            auth=HTTPBasicAuth(conn_config["client_id"], conn_config["client_secret"]),  # type: ignore[arg-type]
+        )
+
+        try:
+            response.raise_for_status()
+        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"]
+
     def get_request_url_header_params(self, query_id: str) -> tuple[dict[str, Any], dict[str, Any], str]:
         """
         Build the request header Url with account name identifier and query id from the connection params.
 
         :param query_id: statement handles query ids for the individual statements.
         """
-        conn_config = self._get_conn_params()
         req_id = uuid.uuid4()
         header = self.get_headers()
         params = {"requestId": str(req_id)}
-        url = (
-            f"https://{conn_config['account']}.{conn_config['region']}"
-            f".snowflakecomputing.com/api/v2/statements/{query_id}"
-        )
+        url = f"{self.account_identifier}.snowflakecomputing.com/api/v2/statements/{query_id}"
         return header, params, url
 
     def check_query_output(self, query_ids: list[str]) -> None:

airflow/providers/snowflake/operators/snowflake.py

@@ -449,7 +449,7 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
             When executing the statement, Snowflake replaces placeholders (? and :name) in
             the statement with these specified values.
     :param deferrable: Run operator in the deferrable mode.
-    """  # noqa
+    """  # noqa: D205, D400
 
     LIFETIME = timedelta(minutes=59)  # The tokens will have a 59 minutes lifetime
     RENEWAL_DELTA = timedelta(minutes=54)  # Tokens will be renewed after 54 minutes

airflow/providers/snowflake/transfers/copy_into_snowflake.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Abstract operator that child classes implement ``COPY INTO <TABLE> SQL in Snowflake``."""
+
 from __future__ import annotations
 
 from typing import Any, Sequence
@@ -254,8 +255,8 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
         run_facets = {}
         if extraction_error_files:
             self.log.debug(
-                f"Unable to extract Dataset namespace and name "
-                f"for the following files: `{extraction_error_files}`."
+                "Unable to extract Dataset namespace and name for the following files: `%s`.",
+                extraction_error_files,
             )
             run_facets["extractionError"] = ExtractionErrorRunFacet(
                 totalTasks=len(query_results),

{apache_airflow_providers_snowflake-5.3.1rc1.dist-info → apache_airflow_providers_snowflake-5.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-snowflake
-Version: 5.3.1rc1
+Version: 5.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>
@@ -19,16 +19,17 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
-Requires-Dist: apache-airflow-providers-common-sql>=1.10.0.dev0
-Requires-Dist: apache-airflow>=2.6.0.dev0
+Requires-Dist: apache-airflow-providers-common-sql>=1.10.0
+Requires-Dist: apache-airflow>=2.6.0
 Requires-Dist: snowflake-connector-python>=2.7.8
 Requires-Dist: snowflake-sqlalchemy>=1.1.0
 Requires-Dist: apache-airflow-providers-common-sql ; extra == "common.sql"
 Requires-Dist: apache-airflow-providers-openlineage ; extra == "openlineage"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.3.1/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.3.1
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.4.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.4.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -80,7 +81,7 @@ Provides-Extra: openlineage
 
 Package ``apache-airflow-providers-snowflake``
 
-Release: ``5.3.1.rc1``
+Release: ``5.4.0``
 
 
 `Snowflake <https://www.snowflake.com/>`__
@@ -93,7 +94,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/5.3.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.4.0/>`_.
 
 Installation
 ------------
@@ -102,7 +103,7 @@ You can install this package on top of an existing Airflow 2 installation (see `
 for the minimum Airflow version supported) via
 ``pip install apache-airflow-providers-snowflake``
 
-The package supports the following python versions: 3.8,3.9,3.10,3.11
+The package supports the following python versions: 3.8,3.9,3.10,3.11,3.12
 
 Requirements
 ------------
@@ -137,4 +138,4 @@ Dependent package
 ==============================================================================================================  ===============
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.3.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.4.0/changelog.html>`_.

{apache_airflow_providers_snowflake-5.3.1rc1.dist-info → apache_airflow_providers_snowflake-5.4.0.dist-info}/RECORD

@@ -1,19 +1,19 @@
 airflow/providers/snowflake/LICENSE,sha256=ywUBpKZc7Jb96rVt5I3IDbg7dIJAbUSHkuoDcF3jbH4,13569
-airflow/providers/snowflake/__init__.py,sha256=0Ltb9XwXTO8bBn-KCUd84KRldnXPXmVy1JHTie04zmc,1584
-airflow/providers/snowflake/get_provider_info.py,sha256=EDKGs3pfZhTomDdR-heB3J6wCUZJRhtF7ws2cIfLysM,4700
+airflow/providers/snowflake/__init__.py,sha256=TaU_oNeaJZa6OMgMaNZR2zEfkbRylW_ZWaJ_UIo78-s,1584
+airflow/providers/snowflake/get_provider_info.py,sha256=sVxscfHltpajNtnDlTjTIOnQ20SZU91TbztOHdKpHG0,4721
 airflow/providers/snowflake/hooks/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/hooks/snowflake.py,sha256=p8En-0cyi0JFXVv20-z4ZlBESrnpmkd7pY1Ccd4hXJQ,21173
-airflow/providers/snowflake/hooks/snowflake_sql_api.py,sha256=b8RNT4cdBaWPO31OB6J0VDj-bEYl4TDgjq-Py3B9n7U,12701
+airflow/providers/snowflake/hooks/snowflake.py,sha256=h0581Pn8Hc16CXKYOxg4EsnAkuJQrEA-Zx_MkVLdfSU,21418
+airflow/providers/snowflake/hooks/snowflake_sql_api.py,sha256=sviTTshDIzqlCQCpw_J3OPQJjO8wa1_6JX8-U_goAw8,14773
 airflow/providers/snowflake/operators/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/operators/snowflake.py,sha256=fX9rot7gXxAZ05h0dfvGcJt256_KtfnJLOLEOO4nTnU,26371
+airflow/providers/snowflake/operators/snowflake.py,sha256=Evn8RfBTjLaSqrj7Dkhvu2-ewJoY6miTUFMiTTVd_dA,26383
 airflow/providers/snowflake/transfers/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/transfers/copy_into_snowflake.py,sha256=-luuj72W3R7SzPXEjrX6waFGExenvma_GsrQ9KWezMA,12632
+airflow/providers/snowflake/transfers/copy_into_snowflake.py,sha256=3ll6N0UhtjnwneSKiBRb9DVUDMX-uODftYccM8b5Bhs,12631
 airflow/providers/snowflake/triggers/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/snowflake/triggers/snowflake_trigger.py,sha256=YfMA7IXq3T7voUishTi171-8nIotPFzhYeFboSrHHAg,4223
 airflow/providers/snowflake/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/snowflake/utils/common.py,sha256=DG-KLy2KpZWAqZqm_XIECm8lmdoUlzwkXv9onmkQThc,1644
 airflow/providers/snowflake/utils/sql_api_generate_jwt.py,sha256=9mR-vHIquv60tfAni87f6FAjKsiRHUDDrsVhzw4M9vM,6762
-apache_airflow_providers_snowflake-5.3.1rc1.dist-info/entry_points.txt,sha256=bCrl5J1PXUMzbgnrKYho61rkbL2gHRT4I6f_1jlxAX4,105
-apache_airflow_providers_snowflake-5.3.1rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_snowflake-5.3.1rc1.dist-info/METADATA,sha256=cOymtTavLMwzJ_sHvCQZO4S95ZCd4x_xE2iCUjFQP68,6469
-apache_airflow_providers_snowflake-5.3.1rc1.dist-info/RECORD,,
+apache_airflow_providers_snowflake-5.4.0.dist-info/entry_points.txt,sha256=bCrl5J1PXUMzbgnrKYho61rkbL2gHRT4I6f_1jlxAX4,105
+apache_airflow_providers_snowflake-5.4.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_snowflake-5.4.0.dist-info/METADATA,sha256=QF3LKAo1GQ7Bp-aDGafZJv2dve-tsQp_h8-TTFXMDSc,6508
+apache_airflow_providers_snowflake-5.4.0.dist-info/RECORD,,