apache-airflow-providers-databricks 7.4.0__py3-none-any.whl → 7.5.0__py3-none-any.whl

airflow/providers/databricks/__init__.py

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

airflow/providers/databricks/get_provider_info.py

@@ -107,6 +107,7 @@ def get_provider_info():
             {
                 "integration-name": "Databricks",
                 "python-modules": [
+                    "airflow.providers.databricks.sensors.databricks",
                     "airflow.providers.databricks.sensors.databricks_sql",
                     "airflow.providers.databricks.sensors.databricks_partition",
                 ],

airflow/providers/databricks/operators/databricks.py

@@ -34,7 +34,6 @@ from airflow.providers.databricks.hooks.databricks import (
     DatabricksHook,
     RunLifeCycleState,
     RunState,
-    SQLStatementState,
 )
 from airflow.providers.databricks.operators.databricks_workflow import (
     DatabricksWorkflowTaskGroup,
@@ -46,13 +45,14 @@ from airflow.providers.databricks.plugins.databricks_workflow import (
 )
 from airflow.providers.databricks.triggers.databricks import (
     DatabricksExecutionTrigger,
-    DatabricksSQLStatementExecutionTrigger,
 )
 from airflow.providers.databricks.utils.databricks import normalise_json_content, validate_trigger_event
+from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMixin
 from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
     from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
     from airflow.utils.task_group import TaskGroup
 
@@ -978,7 +978,7 @@ class DatabricksRunNowOperator(BaseOperator):
             self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
 
 
-class DatabricksSQLStatementsOperator(BaseOperator):
+class DatabricksSQLStatementsOperator(DatabricksSQLStatementsMixin, BaseOperator):
     """
     Submits a Databricks SQL Statement to Databricks using the api/2.0/sql/statements/ API endpoint.
 
@@ -1073,59 +1073,6 @@ class DatabricksSQLStatementsOperator(BaseOperator):
             caller=caller,
         )
 
-    def _handle_operator_execution(self) -> None:
-        end_time = time.time() + self.timeout
-        while end_time > time.time():
-            statement_state = self._hook.get_sql_statement_state(self.statement_id)
-            if statement_state.is_terminal:
-                if statement_state.is_successful:
-                    self.log.info("%s completed successfully.", self.task_id)
-                    return
-                error_message = (
-                    f"{self.task_id} failed with terminal state: {statement_state.state} "
-                    f"and with the error code {statement_state.error_code} "
-                    f"and error message {statement_state.error_message}"
-                )
-                raise AirflowException(error_message)
-
-            self.log.info("%s in run state: %s", self.task_id, statement_state.state)
-            self.log.info("Sleeping for %s seconds.", self.polling_period_seconds)
-            time.sleep(self.polling_period_seconds)
-
-        self._hook.cancel_sql_statement(self.statement_id)
-        raise AirflowException(
-            f"{self.task_id} timed out after {self.timeout} seconds with state: {statement_state.state}",
-        )
-
-    def _handle_deferrable_operator_execution(self) -> None:
-        statement_state = self._hook.get_sql_statement_state(self.statement_id)
-        end_time = time.time() + self.timeout
-        if not statement_state.is_terminal:
-            if not self.statement_id:
-                raise AirflowException("Failed to retrieve statement_id after submitting SQL statement.")
-            self.defer(
-                trigger=DatabricksSQLStatementExecutionTrigger(
-                    statement_id=self.statement_id,
-                    databricks_conn_id=self.databricks_conn_id,
-                    end_time=end_time,
-                    polling_period_seconds=self.polling_period_seconds,
-                    retry_limit=self.databricks_retry_limit,
-                    retry_delay=self.databricks_retry_delay,
-                    retry_args=self.databricks_retry_args,
-                ),
-                method_name=DEFER_METHOD_NAME,
-            )
-        else:
-            if statement_state.is_successful:
-                self.log.info("%s completed successfully.", self.task_id)
-            else:
-                error_message = (
-                    f"{self.task_id} failed with terminal state: {statement_state.state} "
-                    f"and with the error code {statement_state.error_code} "
-                    f"and error message {statement_state.error_message}"
-                )
-                raise AirflowException(error_message)
-
     def execute(self, context: Context):
         json = {
             "statement": self.statement,
@@ -1146,34 +1093,65 @@ class DatabricksSQLStatementsOperator(BaseOperator):
         if not self.wait_for_termination:
             return
         if self.deferrable:
-            self._handle_deferrable_operator_execution()
+            self._handle_deferrable_execution(defer_method_name=DEFER_METHOD_NAME)  # type: ignore[misc]
         else:
-            self._handle_operator_execution()
+            self._handle_execution()  # type: ignore[misc]
 
-    def on_kill(self):
+    def get_openlineage_facets_on_complete(self, _) -> OperatorLineage:
+        """Implement _on_complete because we use statement_id."""
+        from airflow.providers.common.compat.openlineage.facet import (
+            ExternalQueryRunFacet,
+            SQLJobFacet,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+        from airflow.providers.openlineage.sqlparser import DatabaseInfo, SQLParser
+
+        db_info = DatabaseInfo(
+            scheme="databricks",
+            authority=self._hook.host,
+            database=self.catalog,
+            is_uppercase_names=False,
+            # Other args will not be used as we'll not query DB for details, we only do sql parsing.
+        )
+
+        sql_parser = SQLParser(
+            dialect="databricks",
+            default_schema=self.schema or "default",
+        )
+
+        run_facets = {}
         if self.statement_id:
-            self._hook.cancel_sql_statement(self.statement_id)
-            self.log.info(
-                "Task: %s with statement ID: %s was requested to be cancelled.",
-                self.task_id,
-                self.statement_id,
+            run_facets["externalQuery"] = ExternalQueryRunFacet(
+                externalQueryId=self.statement_id, source=sql_parser.create_namespace(db_info)
             )
-        else:
-            self.log.error(
-                "Error: Task: %s with invalid statement_id was requested to be cancelled.", self.task_id
+        job_facets = {"sql": SQLJobFacet(query=SQLParser.normalize_sql(self.statement))}
+
+        query = f"{self.statement}"
+        if self.parameters:
+            # Catalog, schema or table can be parameterized, so it's crucial to fill them before parsing
+            for param in self.parameters:
+                query = query.replace(f":{param['name']}", param.get("value") or "null")
+
+        parser_result = None
+        try:
+            # Try performing offline sql parsing, without db access,
+            parser_result = sql_parser.generate_openlineage_metadata_from_sql(
+                sql=query,
+                database_info=db_info,
+                database=None,  # Provided in db_info
+                use_connection=False,  # Prevents DB call for table details, that will fail with API
+                sqlalchemy_engine=None,  # Not needed when use_connection is False
+                hook=None,  # type: ignore[arg-type] # Not needed when use_connection is False
             )
-
-    def execute_complete(self, context: dict | None, event: dict):
-        statement_state = SQLStatementState.from_json(event["state"])
-        error = event["error"]
-        statement_id = event["statement_id"]
-
-        if statement_state.is_successful:
-            self.log.info("SQL Statement with ID %s completed successfully.", statement_id)
-            return
-
-        error_message = f"SQL Statement execution failed with terminal state: {statement_state} and with the error {error}"
-        raise AirflowException(error_message)
+        except Exception as e:
+            self.log.debug("OpenLineage failed to parse query `%s` with error %s", query, e)
+
+        return OperatorLineage(
+            inputs=parser_result.inputs if parser_result else [],
+            outputs=parser_result.outputs if parser_result else [],
+            job_facets=parser_result.job_facets if parser_result else job_facets,
+            run_facets={**parser_result.run_facets, **run_facets} if parser_result else run_facets,
+        )
 
 
 class DatabricksTaskBaseOperator(BaseOperator, ABC):

airflow/providers/databricks/operators/databricks_sql.py

@@ -277,8 +277,13 @@ class DatabricksCopyIntoOperator(BaseOperator):
         self._client_parameters = client_parameters or {}
         if force_copy is not None:
             self._copy_options["force"] = "true" if force_copy else "false"
+        self._sql: str | None = None
 
     def _get_hook(self) -> DatabricksSqlHook:
+        return self._hook
+
+    @cached_property
+    def _hook(self) -> DatabricksSqlHook:
         return DatabricksSqlHook(
             self.databricks_conn_id,
             http_path=self._http_path,
@@ -354,12 +359,116 @@ FILEFORMAT = {self._file_format}
         return sql.strip()
 
     def execute(self, context: Context) -> Any:
-        sql = self._create_sql_query()
-        self.log.info("Executing: %s", sql)
+        self._sql = self._create_sql_query()
+        self.log.info("Executing: %s", self._sql)
         hook = self._get_hook()
-        hook.run(sql)
+        hook.run(self._sql)
 
     def on_kill(self) -> None:
         # NB: on_kill isn't required for this operator since query cancelling gets
         # handled in `DatabricksSqlHook.run()` method which is called in `execute()`
         ...
+
+    def _build_input_openlineage_dataset(self) -> tuple[Any, list[Any]]:
+        """Parse file_location to build the OpenLineage input dataset."""
+        from urllib.parse import urlparse
+
+        from airflow.providers.common.compat.openlineage.facet import Dataset, Error
+
+        try:
+            uri = urlparse(self.file_location)
+
+            # Only process schemes we know produce valid OL datasets with current implementation
+            if uri.scheme not in ("s3", "s3a", "s3n", "gs", "abfss", "wasbs"):
+                raise ValueError(f"Unsupported scheme: `{uri.scheme}` in `{self.file_location}`")
+
+            namespace = f"{uri.scheme}://{uri.netloc}"
+            name = uri.path.strip("/")
+            if name in ("", "."):
+                name = "/"
+            return Dataset(namespace=namespace, name=name), []
+        except Exception as e:
+            self.log.debug("Failed to parse file_location: `%s`, error: %s", self.file_location, str(e))
+            extraction_errors = [
+                Error(errorMessage=str(e), stackTrace=None, task=self.file_location, taskNumber=None)
+            ]
+            return None, extraction_errors
+
+    def _build_output_openlineage_dataset(self, namespace: str) -> tuple[Any, list[Any]]:
+        """Build output OpenLineage dataset from table information."""
+        from airflow.providers.common.compat.openlineage.facet import Dataset, Error
+
+        try:
+            table_parts = self.table_name.split(".")
+            if len(table_parts) == 3:  # catalog.schema.table
+                catalog, schema, table = table_parts
+            elif len(table_parts) == 2:  # schema.table
+                catalog = None
+                schema, table = table_parts
+            else:
+                catalog = None
+                schema = None
+                table = self.table_name
+
+            hook = self._get_hook()
+            schema = schema or hook.get_openlineage_default_schema()  # Fallback to default schema
+            catalog = catalog or hook.catalog  # Fallback to default catalog, if provided
+
+            # Combine schema/table with optional catalog for final dataset name
+            fq_name = table
+            if schema:
+                fq_name = f"{schema}.{fq_name}"
+            if catalog:
+                fq_name = f"{catalog}.{fq_name}"
+
+            return Dataset(namespace=namespace, name=fq_name), []
+        except Exception as e:
+            self.log.debug("Failed to construct output dataset: `%s`, error: %s", self.table_name, str(e))
+            extraction_errors = [
+                Error(errorMessage=str(e), stackTrace=None, task=self.table_name, taskNumber=None)
+            ]
+            return None, extraction_errors
+
+    def get_openlineage_facets_on_complete(self, _):
+        """Implement _on_complete as we are attaching query id."""
+        from airflow.providers.common.compat.openlineage.facet import (
+            ExternalQueryRunFacet,
+            ExtractionErrorRunFacet,
+            SQLJobFacet,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+        from airflow.providers.openlineage.sqlparser import SQLParser
+
+        if not self._sql:
+            self.log.warning("No SQL query found, returning empty OperatorLineage.")
+            return OperatorLineage()
+
+        hook = self._get_hook()
+        run_facets = {}
+
+        connection = hook.get_connection(self.databricks_conn_id)
+        database_info = hook.get_openlineage_database_info(connection)
+        dbx_namespace = SQLParser.create_namespace(database_info)
+
+        if hook.query_ids:
+            run_facets["externalQuery"] = ExternalQueryRunFacet(
+                externalQueryId=hook.query_ids[0], source=dbx_namespace
+            )
+
+        input_dataset, extraction_errors = self._build_input_openlineage_dataset()
+        output_dataset, output_errors = self._build_output_openlineage_dataset(dbx_namespace)
+        extraction_errors.extend(output_errors)
+
+        if extraction_errors:
+            run_facets["extractionError"] = ExtractionErrorRunFacet(
+                totalTasks=1,
+                failedTasks=len(extraction_errors),
+                errors=extraction_errors,
+            )
+
+        return OperatorLineage(
+            inputs=[input_dataset] if input_dataset else [],
+            outputs=[output_dataset] if output_dataset else [],
+            job_facets={"sql": SQLJobFacet(query=SQLParser.normalize_sql(self._sql))},
+            run_facets=run_facets,
+        )

airflow/providers/databricks/sensors/databricks.py

@@ -0,0 +1,162 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.databricks.hooks.databricks import DatabricksHook, SQLStatementState
+from airflow.providers.databricks.operators.databricks import DEFER_METHOD_NAME
+from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMixin
+from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+XCOM_STATEMENT_ID_KEY = "statement_id"
+
+
+class DatabricksSQLStatementsSensor(DatabricksSQLStatementsMixin, BaseSensorOperator):
+    """DatabricksSQLStatementsSensor."""
+
+    template_fields: Sequence[str] = (
+        "databricks_conn_id",
+        "statement",
+        "statement_id",
+    )
+    template_ext: Sequence[str] = (".json-tpl",)
+    ui_color = "#1CB1C2"
+    ui_fgcolor = "#fff"
+
+    def __init__(
+        self,
+        warehouse_id: str,
+        *,
+        statement: str | None = None,
+        statement_id: str | None = None,
+        catalog: str | None = None,
+        schema: str | None = None,
+        parameters: list[dict[str, Any]] | None = None,
+        databricks_conn_id: str = "databricks_default",
+        polling_period_seconds: int = 30,
+        databricks_retry_limit: int = 3,
+        databricks_retry_delay: int = 1,
+        databricks_retry_args: dict[Any, Any] | None = None,
+        do_xcom_push: bool = True,
+        wait_for_termination: bool = True,
+        timeout: float = 3600,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        # Handle the scenario where either both statement and statement_id are set/not set
+        if statement and statement_id:
+            raise AirflowException("Cannot provide both statement and statement_id.")
+        if not statement and not statement_id:
+            raise AirflowException("One of either statement or statement_id must be provided.")
+
+        if not warehouse_id:
+            raise AirflowException("warehouse_id must be provided.")
+
+        super().__init__(**kwargs)
+
+        self.statement = statement
+        self.statement_id = statement_id
+        self.warehouse_id = warehouse_id
+        self.catalog = catalog
+        self.schema = schema
+        self.parameters = parameters
+        self.databricks_conn_id = databricks_conn_id
+        self.polling_period_seconds = polling_period_seconds
+        self.databricks_retry_limit = databricks_retry_limit
+        self.databricks_retry_delay = databricks_retry_delay
+        self.databricks_retry_args = databricks_retry_args
+        self.wait_for_termination = wait_for_termination
+        self.deferrable = deferrable
+        self.timeout = timeout
+        self.do_xcom_push = do_xcom_push
+
+    @cached_property
+    def _hook(self):
+        return self._get_hook(caller="DatabricksSQLStatementsSensor")
+
+    def _get_hook(self, caller: str) -> DatabricksHook:
+        return DatabricksHook(
+            self.databricks_conn_id,
+            retry_limit=self.databricks_retry_limit,
+            retry_delay=self.databricks_retry_delay,
+            retry_args=self.databricks_retry_args,
+            caller=caller,
+        )
+
+    def execute(self, context: Context):
+        if not self.statement_id:
+            # Otherwise, we'll go ahead and "submit" the statement
+            json = {
+                "statement": self.statement,
+                "warehouse_id": self.warehouse_id,
+                "catalog": self.catalog,
+                "schema": self.schema,
+                "parameters": self.parameters,
+                "wait_timeout": "0s",
+            }
+
+            self.statement_id = self._hook.post_sql_statement(json)
+            self.log.info("SQL Statement submitted with statement_id: %s", self.statement_id)
+
+        if self.do_xcom_push and context is not None:
+            context["ti"].xcom_push(key=XCOM_STATEMENT_ID_KEY, value=self.statement_id)
+
+        # If we're not waiting for the query to complete execution, then we'll go ahead and return. However, a
+        # recommendation to use the DatabricksSQLStatementOperator is made in this case
+        if not self.wait_for_termination:
+            self.log.info(
+                "If setting wait_for_termination = False, consider using the DatabricksSQLStatementsOperator instead."
+            )
+            return
+
+        if self.deferrable:
+            self._handle_deferrable_execution(defer_method_name=DEFER_METHOD_NAME)  # type: ignore[misc]
+
+    def poke(self, context: Context):
+        """
+        Handle non-deferrable Sensor execution.
+
+        :param context: (Context)
+        :return: (bool)
+        """
+        # This is going to very closely mirror the execute_complete
+        statement_state: SQLStatementState = self._hook.get_sql_statement_state(self.statement_id)
+
+        if statement_state.is_running:
+            self.log.info("SQL Statement with ID %s is running", self.statement_id)
+            return False
+        if statement_state.is_successful:
+            self.log.info("SQL Statement with ID %s completed successfully.", self.statement_id)
+            return True
+        raise AirflowException(
+            f"SQL Statement with ID {statement_state} failed with error: {statement_state.error_message}"
+        )

airflow/providers/databricks/utils/mixins.py

@@ -0,0 +1,194 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from __future__ import annotations
+
+import time
+from logging import Logger
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Protocol,
+)
+
+from airflow.exceptions import AirflowException
+from airflow.providers.databricks.hooks.databricks import DatabricksHook, SQLStatementState
+from airflow.providers.databricks.triggers.databricks import DatabricksSQLStatementExecutionTrigger
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class GetHookHasFields(Protocol):
+    """Protocol for get_hook method."""
+
+    databricks_conn_id: str
+    databricks_retry_args: dict | None
+    databricks_retry_delay: int
+    databricks_retry_limit: int
+
+
+class HandleExecutionHasFields(Protocol):
+    """Protocol for _handle_execution method."""
+
+    _hook: DatabricksHook
+    log: Logger
+    polling_period_seconds: int
+    task_id: str
+    timeout: int
+    statement_id: str
+
+
+class HandleDeferrableExecutionHasFields(Protocol):
+    """Protocol for _handle_deferrable_execution method."""
+
+    _hook: DatabricksHook
+    databricks_conn_id: str
+    databricks_retry_args: dict[Any, Any] | None
+    databricks_retry_delay: int
+    databricks_retry_limit: int
+    defer: Any
+    log: Logger
+    polling_period_seconds: int
+    statement_id: str
+    task_id: str
+    timeout: int
+
+
+class ExecuteCompleteHasFields(Protocol):
+    """Protocol for execute_complete method."""
+
+    statement_id: str
+    _hook: DatabricksHook
+    log: Logger
+
+
+class OnKillHasFields(Protocol):
+    """Protocol for on_kill method."""
+
+    _hook: DatabricksHook
+    log: Logger
+    statement_id: str
+    task_id: str
+
+
+class DatabricksSQLStatementsMixin:
+    """
+    Mixin class to be used by both the DatabricksSqlStatementsOperator, and the DatabricksSqlStatementSensor.
+
+        - _handle_operator_execution (renamed to _handle_execution)
+        - _handle_deferrable_operator_execution (renamed to _handle_deferrable_execution)
+        - execute_complete
+        - on_kill
+    """
+
+    def _handle_execution(self: HandleExecutionHasFields) -> None:
+        """Execute a SQL statement in non-deferrable mode."""
+        # Determine the time at which the Task will timeout. The statement_state is defined here in the event
+        # the while-loop is never entered
+        end_time = time.time() + self.timeout
+
+        while end_time > time.time():
+            statement_state: SQLStatementState = self._hook.get_sql_statement_state(self.statement_id)
+
+            if statement_state.is_terminal:
+                if statement_state.is_successful:
+                    self.log.info("%s completed successfully.", self.task_id)
+                    return
+
+                error_message = (
+                    f"{self.task_id} failed with terminal state: {statement_state.state} "
+                    f"and with the error code {statement_state.error_code} "
+                    f"and error message {statement_state.error_message}"
+                )
+                raise AirflowException(error_message)
+
+            self.log.info("%s in run state: %s", self.task_id, statement_state.state)
+            self.log.info("Sleeping for %s seconds.", self.polling_period_seconds)
+            time.sleep(self.polling_period_seconds)
+
+        # Once the timeout is exceeded, the query is cancelled. This is an important steps; if a query takes
+        # to log, it needs to be killed. Otherwise, it may be the case that there are "zombie" queries running
+        # that are no longer being orchestrated
+        self._hook.cancel_sql_statement(self.statement_id)
+        raise AirflowException(
+            f"{self.task_id} timed out after {self.timeout} seconds with state: {statement_state}",
+        )
+
+    def _handle_deferrable_execution(
+        self: HandleDeferrableExecutionHasFields, defer_method_name: str = "execute_complete"
+    ) -> None:
+        """Execute a SQL statement in deferrable mode."""
+        statement_state: SQLStatementState = self._hook.get_sql_statement_state(self.statement_id)
+        end_time: float = time.time() + self.timeout
+
+        if not statement_state.is_terminal:
+            # If the query is still running and there is no statement_id, this is somewhat of a "zombie"
+            # query, and should throw an exception
+            if not self.statement_id:
+                raise AirflowException("Failed to retrieve statement_id after submitting SQL statement.")
+
+            self.defer(
+                trigger=DatabricksSQLStatementExecutionTrigger(
+                    statement_id=self.statement_id,
+                    databricks_conn_id=self.databricks_conn_id,
+                    end_time=end_time,
+                    polling_period_seconds=self.polling_period_seconds,
+                    retry_limit=self.databricks_retry_limit,
+                    retry_delay=self.databricks_retry_delay,
+                    retry_args=self.databricks_retry_args,
+                ),
+                method_name=defer_method_name,
+            )
+
+        else:
+            if statement_state.is_successful:
+                self.log.info("%s completed successfully.", self.task_id)
+            else:
+                error_message = (
+                    f"{self.task_id} failed with terminal state: {statement_state.state} "
+                    f"and with the error code {statement_state.error_code} "
+                    f"and error message {statement_state.error_message}"
+                )
+                raise AirflowException(error_message)
+
+    def execute_complete(self: ExecuteCompleteHasFields, context: Context, event: dict):
+        statement_state = SQLStatementState.from_json(event["state"])
+        error = event["error"]
+        # Save as instance attribute again after coming back from defer (e.g., for later use in listeners)
+        self.statement_id = event["statement_id"]
+
+        if statement_state.is_successful:
+            self.log.info("SQL Statement with ID %s completed successfully.", self.statement_id)
+            return
+
+        error_message = f"SQL Statement execution failed with terminal state: {statement_state} and with the error {error}"
+        raise AirflowException(error_message)
+
+    def on_kill(self: OnKillHasFields) -> None:
+        if self.statement_id:
+            self._hook.cancel_sql_statement(self.statement_id)
+            self.log.info(
+                "Task: %s with statement ID: %s was requested to be cancelled.",
+                self.task_id,
+                self.statement_id,
+            )
+        else:
+            self.log.error(
+                "Error: Task: %s with invalid statement_id was requested to be cancelled.", self.task_id
+            )

{apache_airflow_providers_databricks-7.4.0.dist-info → apache_airflow_providers_databricks-7.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-databricks
-Version: 7.4.0
+Version: 7.5.0
 Summary: Provider package apache-airflow-providers-databricks for Apache Airflow
 Keywords: airflow-provider,databricks,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -23,7 +23,7 @@ Classifier: Topic :: System :: Monitoring
 Requires-Dist: apache-airflow>=2.10.0
 Requires-Dist: apache-airflow-providers-common-compat>=1.6.0
 Requires-Dist: apache-airflow-providers-common-sql>=1.27.0
-Requires-Dist: requests>=2.31.0,<3
+Requires-Dist: requests>=2.32.0,<3
 Requires-Dist: databricks-sql-connector>=3.0.0
 Requires-Dist: databricks-sqlalchemy>=1.0.2
 Requires-Dist: aiohttp>=3.9.2, <4
@@ -36,8 +36,8 @@ Requires-Dist: apache-airflow-providers-openlineage>=2.3.0 ; extra == "openlinea
 Requires-Dist: databricks-sdk==0.10.0 ; extra == "sdk"
 Requires-Dist: apache-airflow-providers-standard ; extra == "standard"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.4.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.4.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.5.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.5.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
@@ -73,7 +73,7 @@ Provides-Extra: standard
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.4.0``
+Release: ``7.5.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -86,7 +86,7 @@ This is a provider package for ``databricks`` provider. All classes for this pro
 are in ``airflow.providers.databricks`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.5.0/>`_.
 
 Installation
 ------------
@@ -106,7 +106,7 @@ PIP package                                 Version required
 ``apache-airflow``                          ``>=2.10.0``
 ``apache-airflow-providers-common-compat``  ``>=1.6.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
-``requests``                                ``>=2.31.0,<3``
+``requests``                                ``>=2.32.0,<3``
 ``databricks-sql-connector``                ``>=3.0.0``
 ``databricks-sqlalchemy``                   ``>=1.0.2``
 ``aiohttp``                                 ``>=3.9.2,<4``
@@ -125,16 +125,18 @@ You can install such cross-provider dependencies when installing from PyPI. For
 
 .. code-block:: bash
 
-    pip install apache-airflow-providers-databricks[common.sql]
+    pip install apache-airflow-providers-databricks[common.compat]
 
 
-============================================================================================================  ==============
-Dependent package                                                                                             Extra
-============================================================================================================  ==============
-`apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_  ``common.sql``
-`apache-airflow-providers-fab <https://airflow.apache.org/docs/apache-airflow-providers-fab>`_                ``fab``
-============================================================================================================  ==============
+==================================================================================================================  =================
+Dependent package                                                                                                   Extra
+==================================================================================================================  =================
+`apache-airflow-providers-common-compat <https://airflow.apache.org/docs/apache-airflow-providers-common-compat>`_  ``common.compat``
+`apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_        ``common.sql``
+`apache-airflow-providers-fab <https://airflow.apache.org/docs/apache-airflow-providers-fab>`_                      ``fab``
+`apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_      ``openlineage``
+==================================================================================================================  =================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.5.0/changelog.html>`_.
 

{apache_airflow_providers_databricks-7.4.0.dist-info → apache_airflow_providers_databricks-7.5.0.dist-info}/RECORD

@@ -1,28 +1,30 @@
 airflow/providers/databricks/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
-airflow/providers/databricks/__init__.py,sha256=D4S1V7H6S0R_iJ75oqoynUbmOYkMET00uTSUyPzvkzM,1499
+airflow/providers/databricks/__init__.py,sha256=d1KmgGbpEA3KHmX7l-luO2NBwkPvv7tSFPf-KMbN2LE,1499
 airflow/providers/databricks/exceptions.py,sha256=85RklmLOI_PnTzfXNIUd5fAu2aMMUhelwumQAX0wANE,1261
-airflow/providers/databricks/get_provider_info.py,sha256=qNMX4Lft-NItPhFewFBSCi8n0_ISid_MQeETKQ67vdo,5573
+airflow/providers/databricks/get_provider_info.py,sha256=NZ-rY6k6ctDZN7rDngN7mAzq7RMhLag5NwfnuBNcKuw,5644
 airflow/providers/databricks/version_compat.py,sha256=j5PCtXvZ71aBjixu-EFTNtVDPsngzzs7os0ZQDgFVDk,1536
 airflow/providers/databricks/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
 airflow/providers/databricks/hooks/databricks.py,sha256=FIoiKWIc9AP3s8Av3Av9yleTg1kI0norwW5CAc6jTQc,28867
 airflow/providers/databricks/hooks/databricks_base.py,sha256=D7-_74QgQaZm1NfHKl_UOXbVAXRo2xjnOx_r1MI-rWI,34871
 airflow/providers/databricks/hooks/databricks_sql.py,sha256=xTdi0JN-ZdsGe2XnCa8yBi-AINZUlyIVlP-5nb2d2T0,16964
 airflow/providers/databricks/operators/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/operators/databricks.py,sha256=E8fgk3Z67uOTSvWvbF23Miv6EruSGOTdFvHn7pGVWp0,80138
+airflow/providers/databricks/operators/databricks.py,sha256=yDy_pBaAi_muP3NstpXOqBNxSP9WL0_X3fX2OmR1f3c,79235
 airflow/providers/databricks/operators/databricks_repos.py,sha256=m_72OnnU9df7UB-8SK2Tp5VjfNyjYeAnil3dCKs9SbA,13282
-airflow/providers/databricks/operators/databricks_sql.py,sha256=Ycp5mcb3uScQrognB2k8IeSR9oBx-Vnv6NEYGYuE800,17159
+airflow/providers/databricks/operators/databricks_sql.py,sha256=yrYZa9Hq8JDc-8F5DGfW2mkcaNwu4o09JZj_SQQnsrE,21807
 airflow/providers/databricks/operators/databricks_workflow.py,sha256=9WNQR9COa90fbqb9qSzut34K9Z1S_ZdpNHAfIcuH454,14227
 airflow/providers/databricks/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/databricks/plugins/databricks_workflow.py,sha256=1UpsodBLRrTah9zBGBzfM7n1pdkzTo7yilt6QxASspQ,17460
 airflow/providers/databricks/sensors/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/databricks/sensors/databricks.py,sha256=zG2rS7xemVUk-ztDSj0t90Ws47kqRgPm3NsBMQQR8bA,6389
 airflow/providers/databricks/sensors/databricks_partition.py,sha256=2zWdnqVaSSd7PFTZadfvtbsR7zOI4GwfZFOuEnXRLSM,10023
 airflow/providers/databricks/sensors/databricks_sql.py,sha256=jIA9oGBUCAlXzyrqigxlg7JQDsBFuNIF8ZUEJM8gPxg,5766
 airflow/providers/databricks/triggers/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
 airflow/providers/databricks/triggers/databricks.py,sha256=dSogx6GlcJfZ4CFhtlMeWs9sYFEYthP82S_U8-tM2Tk,9240
 airflow/providers/databricks/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/databricks/utils/databricks.py,sha256=s0qEr_DsFhKW4uUiq2VQbtqcj52isYIplPZsUcxGPrI,2862
+airflow/providers/databricks/utils/mixins.py,sha256=WUmkt3AmXalmV6zOUIJZWbTldxYunAZOstddDhKCC94,7407
 airflow/providers/databricks/utils/openlineage.py,sha256=7fR3CPcOruHapsz1DOZ38QN3ZcAGDADNHPY28CzYCbg,13194
-apache_airflow_providers_databricks-7.4.0.dist-info/entry_points.txt,sha256=hjmZm3ab2cteTR4t9eE28oKixHwNIKtLCThd6sx3XRQ,227
-apache_airflow_providers_databricks-7.4.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-apache_airflow_providers_databricks-7.4.0.dist-info/METADATA,sha256=6IAI-03VfbWg0tayX29gnzSeq5Em510la7J98MSQPd8,6446
-apache_airflow_providers_databricks-7.4.0.dist-info/RECORD,,
+apache_airflow_providers_databricks-7.5.0.dist-info/entry_points.txt,sha256=hjmZm3ab2cteTR4t9eE28oKixHwNIKtLCThd6sx3XRQ,227
+apache_airflow_providers_databricks-7.5.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_databricks-7.5.0.dist-info/METADATA,sha256=3Of4M9LHN0OsBur6wL1Nr3olZEGS_zAMIvrYOW6Rjaw,6760
+apache_airflow_providers_databricks-7.5.0.dist-info/RECORD,,