apache-airflow-providers-snowflake 5.2.1rc1__py3-none-any.whl → 5.3.0__py3-none-any.whl

airflow/providers/snowflake/__init__.py

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

airflow/providers/snowflake/get_provider_info.py

@@ -27,9 +27,10 @@ def get_provider_info():
         "package-name": "apache-airflow-providers-snowflake",
         "name": "Snowflake",
         "description": "`Snowflake <https://www.snowflake.com/>`__\n",
-        "suspended": False,
-        "source-date-epoch": 1703288173,
+        "state": "ready",
+        "source-date-epoch": 1705912272,
         "versions": [
+            "5.3.0",
             "5.2.1",
             "5.2.0",
             "5.1.2",

airflow/providers/snowflake/hooks/snowflake.py

@@ -170,7 +170,9 @@ class SnowflakeHook(DbApiHook):
                 warnings.warn(
                     f"Conflicting params `{field_name}` and `{backcompat_key}` found in extras. "
                     f"Using value for `{field_name}`.  Please ensure this is the correct "
-                    f"value and remove the backcompat key `{backcompat_key}`."
+                    f"value and remove the backcompat key `{backcompat_key}`.",
+                    UserWarning,
+                    stacklevel=2,
                 )
             return extra_dict[field_name] or None
         return extra_dict.get(backcompat_key) or None
@@ -300,7 +302,7 @@ class SnowflakeHook(DbApiHook):
     def get_autocommit(self, conn):
         return getattr(conn, "autocommit_mode", False)
 
-    @overload
+    @overload  # type: ignore[override]
     def run(
         self,
         sql: str | Iterable[str],
@@ -385,10 +387,10 @@ class SnowflakeHook(DbApiHook):
             with self._get_cursor(conn, return_dictionaries) as cur:
                 results = []
                 for sql_statement in sql_list:
-                    self._run_command(cur, sql_statement, parameters)
+                    self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
 
                     if handler is not None:
-                        result = self._make_common_data_structure(handler(cur))
+                        result = self._make_common_data_structure(handler(cur))  # type: ignore[attr-defined]
                         if return_single_query_results(sql, return_last, split_statements):
                             _last_result = result
                             _last_description = cur.description

airflow/providers/snowflake/operators/snowflake.py

@@ -514,6 +514,21 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
         if self.do_xcom_push:
             context["ti"].xcom_push(key="query_ids", value=self.query_ids)
 
+        succeeded_query_ids = []
+        for query_id in self.query_ids:
+            self.log.info("Retrieving status for query id %s", query_id)
+            statement_status = self._hook.get_sql_api_query_status(query_id)
+            if statement_status.get("status") == "running":
+                break
+            elif statement_status.get("status") == "success":
+                succeeded_query_ids.append(query_id)
+            else:
+                raise AirflowException(f"{statement_status.get('status')}: {statement_status.get('message')}")
+
+        if len(self.query_ids) == len(succeeded_query_ids):
+            self.log.info("%s completed successfully.", self.task_id)
+            return
+
         if self.deferrable:
             self.defer(
                 timeout=self.execution_timeout,

airflow/providers/snowflake/transfers/copy_into_snowflake.py

@@ -108,8 +108,12 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
         self.copy_options = copy_options
         self.validation_mode = validation_mode
 
+        self.hook: SnowflakeHook | None = None
+        self._sql: str | None = None
+        self._result: list[dict[str, Any]] = []
+
     def execute(self, context: Any) -> None:
-        snowflake_hook = SnowflakeHook(
+        self.hook = SnowflakeHook(
             snowflake_conn_id=self.snowflake_conn_id,
             warehouse=self.warehouse,
             database=self.database,
@@ -127,7 +131,7 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
         if self.columns_array:
             into = f"{into}({', '.join(self.columns_array)})"
 
-        sql = f"""
+        self._sql = f"""
         COPY INTO {into}
              FROM  @{self.stage}/{self.prefix or ""}
         {"FILES=(" + ",".join(map(enclose_param, self.files)) + ")" if self.files else ""}
@@ -137,5 +141,158 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
         {self.validation_mode or ""}
         """
         self.log.info("Executing COPY command...")
-        snowflake_hook.run(sql=sql, autocommit=self.autocommit)
+        self._result = self.hook.run(  # type: ignore # mypy does not work well with return_dictionaries=True
+            sql=self._sql,
+            autocommit=self.autocommit,
+            handler=lambda x: x.fetchall(),
+            return_dictionaries=True,
+        )
         self.log.info("COPY command completed")
+
+    @staticmethod
+    def _extract_openlineage_unique_dataset_paths(
+        query_result: list[dict[str, Any]],
+    ) -> tuple[list[tuple[str, str]], list[str]]:
+        """Extracts and returns unique OpenLineage dataset paths and file paths that failed to be parsed.
+
+        Each row in the results is expected to have a 'file' field, which is a URI.
+        The function parses these URIs and constructs a set of unique OpenLineage (namespace, name) tuples.
+        Additionally, it captures any URIs that cannot be parsed or processed
+        and returns them in a separate error list.
+
+        For Azure, Snowflake has a unique way of representing URI:
+            azure://<account_name>.blob.core.windows.net/<container_name>/path/to/file.csv
+        that is transformed by this function to a Dataset with more universal naming convention:
+            Dataset(namespace="wasbs://container_name@account_name", name="path/to"), as described at
+        https://github.com/OpenLineage/OpenLineage/blob/main/spec/Naming.md#wasbs-azure-blob-storage
+
+        :param query_result: A list of dictionaries, each containing a 'file' key with a URI value.
+        :return: Two lists - the first is a sorted list of tuples, each representing a unique dataset path,
+         and the second contains any URIs that cannot be parsed or processed correctly.
+
+        >>> method = CopyFromExternalStageToSnowflakeOperator._extract_openlineage_unique_dataset_paths
+
+        >>> results = [{"file": "azure://my_account.blob.core.windows.net/azure_container/dir3/file.csv"}]
+        >>> method(results)
+        ([('wasbs://azure_container@my_account', 'dir3')], [])
+
+        >>> results = [{"file": "azure://my_account.blob.core.windows.net/azure_container"}]
+        >>> method(results)
+        ([('wasbs://azure_container@my_account', '/')], [])
+
+        >>> results = [{"file": "s3://bucket"}, {"file": "gcs://bucket/"}, {"file": "s3://bucket/a.csv"}]
+        >>> method(results)
+        ([('gcs://bucket', '/'), ('s3://bucket', '/')], [])
+
+        >>> results = [{"file": "s3://bucket/dir/file.csv"}, {"file": "gcs://bucket/dir/dir2/a.txt"}]
+        >>> method(results)
+        ([('gcs://bucket', 'dir/dir2'), ('s3://bucket', 'dir')], [])
+
+        >>> results = [
+        ...     {"file": "s3://bucket/dir/file.csv"},
+        ...     {"file": "azure://my_account.something_new.windows.net/azure_container"},
+        ... ]
+        >>> method(results)
+        ([('s3://bucket', 'dir')], ['azure://my_account.something_new.windows.net/azure_container'])
+        """
+        import re
+        from pathlib import Path
+        from urllib.parse import urlparse
+
+        azure_regex = r"azure:\/\/(\w+)?\.blob.core.windows.net\/(\w+)\/?(.*)?"
+        extraction_error_files = []
+        unique_dataset_paths = set()
+
+        for row in query_result:
+            uri = urlparse(row["file"])
+            if uri.scheme == "azure":
+                match = re.fullmatch(azure_regex, row["file"])
+                if not match:
+                    extraction_error_files.append(row["file"])
+                    continue
+                account_name, container_name, name = match.groups()
+                namespace = f"wasbs://{container_name}@{account_name}"
+            else:
+                namespace = f"{uri.scheme}://{uri.netloc}"
+                name = uri.path.lstrip("/")
+
+            name = Path(name).parent.as_posix()
+            if name in ("", "."):
+                name = "/"
+
+            unique_dataset_paths.add((namespace, name))
+
+        return sorted(unique_dataset_paths), sorted(extraction_error_files)
+
+    def get_openlineage_facets_on_complete(self, task_instance):
+        """Implement _on_complete because we rely on return value of a query."""
+        import re
+
+        from openlineage.client.facet import (
+            ExternalQueryRunFacet,
+            ExtractionError,
+            ExtractionErrorRunFacet,
+            SqlJobFacet,
+        )
+        from openlineage.client.run import Dataset
+
+        from airflow.providers.openlineage.extractors import OperatorLineage
+        from airflow.providers.openlineage.sqlparser import SQLParser
+
+        if not self._sql:
+            return OperatorLineage()
+
+        query_results = self._result or []
+        # If no files were uploaded we get [{"status": "0 files were uploaded..."}]
+        if len(query_results) == 1 and query_results[0].get("status"):
+            query_results = []
+        unique_dataset_paths, extraction_error_files = self._extract_openlineage_unique_dataset_paths(
+            query_results
+        )
+        input_datasets = [Dataset(namespace=namespace, name=name) for namespace, name in unique_dataset_paths]
+
+        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}`."
+            )
+            run_facets["extractionError"] = ExtractionErrorRunFacet(
+                totalTasks=len(query_results),
+                failedTasks=len(extraction_error_files),
+                errors=[
+                    ExtractionError(
+                        errorMessage="Unable to extract Dataset namespace and name.",
+                        stackTrace=None,
+                        task=file_uri,
+                        taskNumber=None,
+                    )
+                    for file_uri in extraction_error_files
+                ],
+            )
+
+        connection = self.hook.get_connection(getattr(self.hook, str(self.hook.conn_name_attr)))
+        database_info = self.hook.get_openlineage_database_info(connection)
+
+        dest_name = self.table
+        schema = self.hook.get_openlineage_default_schema()
+        database = database_info.database
+        if schema:
+            dest_name = f"{schema}.{dest_name}"
+            if database:
+                dest_name = f"{database}.{dest_name}"
+
+        snowflake_namespace = SQLParser.create_namespace(database_info)
+        query = SQLParser.normalize_sql(self._sql)
+        query = re.sub(r"\n+", "\n", re.sub(r" +", " ", query))
+
+        run_facets["externalQuery"] = ExternalQueryRunFacet(
+            externalQueryId=self.hook.query_ids[0], source=snowflake_namespace
+        )
+
+        return OperatorLineage(
+            inputs=input_datasets,
+            outputs=[Dataset(namespace=snowflake_namespace, name=dest_name)],
+            job_facets={"sql": SqlJobFacet(query=query)},
+            run_facets=run_facets,
+        )

{apache_airflow_providers_snowflake-5.2.1rc1.dist-info → apache_airflow_providers_snowflake-5.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-snowflake
-Version: 5.2.1rc1
+Version: 5.3.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>
@@ -20,15 +20,15 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 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.2.1/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.2.1
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.3.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.3.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 +80,7 @@ Provides-Extra: openlineage
 
 Package ``apache-airflow-providers-snowflake``
 
-Release: ``5.2.1.rc1``
+Release: ``5.3.0``
 
 
 `Snowflake <https://www.snowflake.com/>`__
@@ -93,7 +93,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.2.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.3.0/>`_.
 
 Installation
 ------------
@@ -137,4 +137,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.2.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-snowflake/5.3.0/changelog.html>`_.

{apache_airflow_providers_snowflake-5.2.1rc1.dist-info → apache_airflow_providers_snowflake-5.3.0.dist-info}/RECORD

@@ -1,19 +1,19 @@
 airflow/providers/snowflake/LICENSE,sha256=ywUBpKZc7Jb96rVt5I3IDbg7dIJAbUSHkuoDcF3jbH4,13569
-airflow/providers/snowflake/__init__.py,sha256=N1ToUM81znAFMFd1wrRK0qvhAnkr9XI8cWGn4piL7XU,1584
-airflow/providers/snowflake/get_provider_info.py,sha256=tVuJUNRgl7RUJmFB50ibcYus3VyZGJ7fqE1nWu9m4A4,4660
+airflow/providers/snowflake/__init__.py,sha256=uni5RXXyJ7ikG3ilrlEdQ5XaDkSiQD_RdMg01GlbUnc,1584
+airflow/providers/snowflake/get_provider_info.py,sha256=69aKCJgzdB03_e0BobdDmCp4CcVA3XrDwr6u8Ozja1c,4679
 airflow/providers/snowflake/hooks/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/hooks/snowflake.py,sha256=e8hmCQ0_yWj48leljR5pJgWgk9g4bL8IvbBkzW_usDA,21026
+airflow/providers/snowflake/hooks/snowflake.py,sha256=hPJfFjTWDVAuSqp8GhDrNpQ0kh59t4hLMQAUVioTa1k,21180
 airflow/providers/snowflake/hooks/snowflake_sql_api.py,sha256=RCakX8FjmTDWfDZJkxkHnZB9VJDxORZ0W6SHfS-Oh5M,12703
 airflow/providers/snowflake/operators/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/operators/snowflake.py,sha256=Pyp-SS1knrJm6GziOjHwfEpcAY96BubCy_-utTrZOg8,25686
+airflow/providers/snowflake/operators/snowflake.py,sha256=yZhA9A2nYdIiFVEW4iDC6OtpKx7Qw4DJ7MMq1me4M08,26386
 airflow/providers/snowflake/transfers/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/snowflake/transfers/copy_into_snowflake.py,sha256=rdhgkggjEEeZ7Afi3PIyTC-XP0a2NgCIqzSWe1l262w,5765
+airflow/providers/snowflake/transfers/copy_into_snowflake.py,sha256=YB7E31SsgCKFHfp7GYYSm-0dM5iCxSPMGE9B10-SwT8,12634
 airflow/providers/snowflake/triggers/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/snowflake/triggers/snowflake_trigger.py,sha256=HgQGOpAhrCfh0bn6kDxakIi2afq1oRr_d6Av5yxRx5c,4224
 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=4LueR3FySbo99cza8nZBJIGy9Wcrs9jFiWlVUHiK42w,6763
-apache_airflow_providers_snowflake-5.2.1rc1.dist-info/entry_points.txt,sha256=bCrl5J1PXUMzbgnrKYho61rkbL2gHRT4I6f_1jlxAX4,105
-apache_airflow_providers_snowflake-5.2.1rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_snowflake-5.2.1rc1.dist-info/METADATA,sha256=S0-P_jN7VqjJx2lWd36lEvDRW5hrlYPZhhJ3uFc2FdQ,6469
-apache_airflow_providers_snowflake-5.2.1rc1.dist-info/RECORD,,
+apache_airflow_providers_snowflake-5.3.0.dist-info/entry_points.txt,sha256=bCrl5J1PXUMzbgnrKYho61rkbL2gHRT4I6f_1jlxAX4,105
+apache_airflow_providers_snowflake-5.3.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_snowflake-5.3.0.dist-info/METADATA,sha256=7cHE3hWFtVw1GByGlGpBgpoGmymPjGkKaAm6O7OFSAE,6452
+apache_airflow_providers_snowflake-5.3.0.dist-info/RECORD,,