apache-airflow-providers-databricks 6.3.0rc3__tar.gz → 6.4.0__tar.gz

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 6.3.0rc3
+Version: 6.4.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>
@@ -22,15 +22,15 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: aiohttp>=3.9.2, <4
-Requires-Dist: apache-airflow-providers-common-sql>=1.10.0rc0
-Requires-Dist: apache-airflow>=2.6.0rc0
+Requires-Dist: apache-airflow-providers-common-sql>=1.10.0
+Requires-Dist: apache-airflow>=2.7.0
 Requires-Dist: databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0
 Requires-Dist: requests>=2.27.0,<3
 Requires-Dist: apache-airflow-providers-common-sql ; extra == "common.sql"
 Requires-Dist: databricks-sdk==0.10.0 ; extra == "sdk"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.3.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.3.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.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
@@ -82,7 +82,7 @@ Provides-Extra: sdk
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.3.0.rc3``
+Release: ``6.4.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -95,7 +95,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/6.3.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0/>`_.
 
 Installation
 ------------
@@ -112,7 +112,7 @@ Requirements
 =======================================  ==========================
 PIP package                              Version required
 =======================================  ==========================
-``apache-airflow``                       ``>=2.6.0``
+``apache-airflow``                       ``>=2.7.0``
 ``apache-airflow-providers-common-sql``  ``>=1.10.0``
 ``requests``                             ``>=2.27.0,<3``
 ``databricks-sql-connector``             ``>=2.0.0,!=2.9.0,<3.0.0``
@@ -139,4 +139,4 @@ Dependent package
 ============================================================================================================  ==============
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.3.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.3.0.rc3``
+Release: ``6.4.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -55,7 +55,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/6.3.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0/>`_.
 
 Installation
 ------------
@@ -72,7 +72,7 @@ Requirements
 =======================================  ==========================
 PIP package                              Version required
 =======================================  ==========================
-``apache-airflow``                       ``>=2.6.0``
+``apache-airflow``                       ``>=2.7.0``
 ``apache-airflow-providers-common-sql``  ``>=1.10.0``
 ``requests``                             ``>=2.27.0,<3``
 ``databricks-sql-connector``             ``>=2.0.0,!=2.9.0,<3.0.0``
@@ -99,4 +99,4 @@ Dependent package
 ============================================================================================================  ==============
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.3.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/airflow/providers/databricks/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "6.3.0"
+__version__ = "6.4.0"
 
 try:
     from airflow import __version__ as airflow_version
@@ -35,8 +35,8 @@ except ImportError:
     from airflow.version import version as airflow_version
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.6.0"
+    "2.7.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-databricks:{__version__}` needs Apache Airflow 2.6.0+"
+        f"The package `apache-airflow-providers-databricks:{__version__}` needs Apache Airflow 2.7.0+"
     )

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/airflow/providers/databricks/get_provider_info.py

@@ -28,8 +28,9 @@ def get_provider_info():
         "name": "Databricks",
         "description": "`Databricks <https://databricks.com/>`__\n",
         "state": "ready",
-        "source-date-epoch": 1712665557,
+        "source-date-epoch": 1714476154,
         "versions": [
+            "6.4.0",
             "6.3.0",
             "6.2.0",
             "6.1.0",
@@ -67,7 +68,7 @@ def get_provider_info():
             "1.0.0",
         ],
         "dependencies": [
-            "apache-airflow>=2.6.0",
+            "apache-airflow>=2.7.0",
             "apache-airflow-providers-common-sql>=1.10.0",
             "requests>=2.27.0,<3",
             "databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0",
@@ -87,6 +88,7 @@ def get_provider_info():
                 "external-doc-url": "https://databricks.com/",
                 "how-to-guide": [
                     "/docs/apache-airflow-providers-databricks/operators/jobs_create.rst",
+                    "/docs/apache-airflow-providers-databricks/operators/notebook.rst",
                     "/docs/apache-airflow-providers-databricks/operators/submit_run.rst",
                     "/docs/apache-airflow-providers-databricks/operators/run_now.rst",
                 ],

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/airflow/providers/databricks/hooks/databricks.py

@@ -491,6 +491,17 @@ class DatabricksHook(BaseDatabricksHook):
         run_output = self._do_api_call(OUTPUT_RUNS_JOB_ENDPOINT, json)
         return run_output
 
+    async def a_get_run_output(self, run_id: int) -> dict:
+        """
+        Async version of `get_run_output()`.
+
+        :param run_id: id of the run
+        :return: output of the run
+        """
+        json = {"run_id": run_id}
+        run_output = await self._a_do_api_call(OUTPUT_RUNS_JOB_ENDPOINT, json)
+        return run_output
+
     def cancel_run(self, run_id: int) -> None:
         """
         Cancel the run.

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/airflow/providers/databricks/operators/databricks.py

@@ -70,10 +70,9 @@ def _handle_databricks_operator_execution(operator, hook, log, context) -> None:
 
                 if run_state.result_state == "FAILED":
                     task_run_id = None
-                    if "tasks" in run_info:
-                        for task in run_info["tasks"]:
-                            if task.get("state", {}).get("result_state", "") == "FAILED":
-                                task_run_id = task["run_id"]
+                    for task in run_info.get("tasks", []):
+                        if task.get("state", {}).get("result_state", "") == "FAILED":
+                            task_run_id = task["run_id"]
                     if task_run_id is not None:
                         run_output = hook.get_run_output(task_run_id)
                         if "error" in run_output:
@@ -160,13 +159,15 @@ def _handle_deferrable_databricks_operator_completion(event: dict, log: Logger)
     validate_trigger_event(event)
     run_state = RunState.from_json(event["run_state"])
     run_page_url = event["run_page_url"]
+    errors = event["errors"]
     log.info("View run status, Spark UI, and logs at %s", run_page_url)
 
     if run_state.is_successful:
         log.info("Job run completed successfully.")
         return
 
-    error_message = f"Job run failed with terminal state: {run_state}"
+    error_message = f"Job run failed with terminal state: {run_state} and with the errors {errors}"
+
     if event["repair_run"]:
         log.warning(
             "%s but since repair run is set, repairing the run with all failed tasks",
@@ -207,6 +208,7 @@ class DatabricksCreateJobsOperator(BaseOperator):
         .. seealso::
             For more information about templating see :ref:`concepts:jinja-templating`.
     :param name: An optional name for the job.
+    :param description: An optional description for the job.
     :param tags: A map of tags associated with the job.
     :param tasks: A list of task specifications to be executed by this job.
         Array of objects (JobTaskSettings).
@@ -214,6 +216,7 @@ class DatabricksCreateJobsOperator(BaseOperator):
         tasks of this job. Array of objects (JobCluster).
     :param email_notifications: Object (JobEmailNotifications).
     :param webhook_notifications: Object (WebhookNotifications).
+    :param notification_settings: Optional notification settings.
     :param timeout_seconds: An optional timeout applied to each run of this job.
     :param schedule: Object (CronSchedule).
     :param max_concurrent_runs: An optional maximum allowed number of concurrent runs of the job.
@@ -249,11 +252,13 @@ class DatabricksCreateJobsOperator(BaseOperator):
         *,
         json: Any | None = None,
         name: str | None = None,
+        description: str | None = None,
         tags: dict[str, str] | None = None,
         tasks: list[dict] | None = None,
         job_clusters: list[dict] | None = None,
         email_notifications: dict | None = None,
         webhook_notifications: dict | None = None,
+        notification_settings: dict | None = None,
         timeout_seconds: int | None = None,
         schedule: dict | None = None,
         max_concurrent_runs: int | None = None,
@@ -276,6 +281,8 @@ class DatabricksCreateJobsOperator(BaseOperator):
         self.databricks_retry_args = databricks_retry_args
         if name is not None:
             self.json["name"] = name
+        if description is not None:
+            self.json["description"] = description
         if tags is not None:
             self.json["tags"] = tags
         if tasks is not None:
@@ -286,6 +293,8 @@ class DatabricksCreateJobsOperator(BaseOperator):
             self.json["email_notifications"] = email_notifications
         if webhook_notifications is not None:
             self.json["webhook_notifications"] = webhook_notifications
+        if notification_settings is not None:
+            self.json["notification_settings"] = notification_settings
         if timeout_seconds is not None:
             self.json["timeout_seconds"] = timeout_seconds
         if schedule is not None:
@@ -884,3 +893,177 @@ class DatabricksRunNowDeferrableOperator(DatabricksRunNowOperator):
 
     def __init__(self, *args, **kwargs):
         super().__init__(deferrable=True, *args, **kwargs)
+
+
+class DatabricksNotebookOperator(BaseOperator):
+    """
+    Runs a notebook on Databricks using an Airflow operator.
+
+    The DatabricksNotebookOperator allows users to launch and monitor notebook
+    job runs on Databricks as Airflow tasks.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:DatabricksNotebookOperator`
+
+    :param notebook_path: The path to the notebook in Databricks.
+    :param source: Optional location type of the notebook. When set to WORKSPACE, the notebook will be retrieved
+            from the local Databricks workspace. When set to GIT, the notebook will be retrieved from a Git repository
+            defined in git_source. If the value is empty, the task will use GIT if git_source is defined
+            and WORKSPACE otherwise. For more information please visit
+            https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsCreate
+    :param notebook_params: A dict of key-value pairs to be passed as optional params to the notebook task.
+    :param notebook_packages: A list of the Python libraries to be installed on the cluster running the
+        notebook.
+    :param new_cluster: Specs for a new cluster on which this task will be run.
+    :param existing_cluster_id: ID for existing cluster on which to run this task.
+    :param job_cluster_key: The key for the job cluster.
+    :param polling_period_seconds: Controls the rate which we poll for the result of this notebook job run.
+    :param databricks_retry_limit: Amount of times to retry if the Databricks backend is unreachable.
+    :param databricks_retry_delay: Number of seconds to wait between retries.
+    :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+    :param wait_for_termination: if we should wait for termination of the job run. ``True`` by default.
+    :param databricks_conn_id: The name of the Airflow connection to use.
+    """
+
+    template_fields = ("notebook_params",)
+
+    def __init__(
+        self,
+        notebook_path: str,
+        source: str,
+        notebook_params: dict | None = None,
+        notebook_packages: list[dict[str, Any]] | None = None,
+        new_cluster: dict[str, Any] | None = None,
+        existing_cluster_id: str = "",
+        job_cluster_key: str = "",
+        polling_period_seconds: int = 5,
+        databricks_retry_limit: int = 3,
+        databricks_retry_delay: int = 1,
+        databricks_retry_args: dict[Any, Any] | None = None,
+        wait_for_termination: bool = True,
+        databricks_conn_id: str = "databricks_default",
+        **kwargs: Any,
+    ):
+        self.notebook_path = notebook_path
+        self.source = source
+        self.notebook_params = notebook_params or {}
+        self.notebook_packages = notebook_packages or []
+        self.new_cluster = new_cluster or {}
+        self.existing_cluster_id = existing_cluster_id
+        self.job_cluster_key = job_cluster_key
+        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.databricks_conn_id = databricks_conn_id
+        self.databricks_run_id: int | None = None
+        super().__init__(**kwargs)
+
+    @cached_property
+    def _hook(self) -> DatabricksHook:
+        return self._get_hook(caller="DatabricksNotebookOperator")
+
+    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 _get_task_timeout_seconds(self) -> int:
+        """
+        Get the timeout seconds value for the Databricks job based on the execution timeout value provided for the Airflow task.
+
+        By default, tasks in Airflow have an execution_timeout set to None. In Airflow, when
+        execution_timeout is not defined, the task continues to run indefinitely. Therefore,
+        to mirror this behavior in the Databricks Jobs API, we set the timeout to 0, indicating
+        that the job should run indefinitely. This aligns with the default behavior of Databricks jobs,
+        where a timeout seconds value of 0 signifies an indefinite run duration.
+        More details can be found in the Databricks documentation:
+        See https://docs.databricks.com/api/workspace/jobs/submit#timeout_seconds
+        """
+        if self.execution_timeout is None:
+            return 0
+        execution_timeout_seconds = int(self.execution_timeout.total_seconds())
+        if execution_timeout_seconds == 0:
+            raise ValueError(
+                "If you've set an `execution_timeout` for the task, ensure it's not `0`. Set it instead to "
+                "`None` if you desire the task to run indefinitely."
+            )
+        return execution_timeout_seconds
+
+    def _get_task_base_json(self) -> dict[str, Any]:
+        """Get task base json to be used for task submissions."""
+        return {
+            "timeout_seconds": self._get_task_timeout_seconds(),
+            "email_notifications": {},
+            "notebook_task": {
+                "notebook_path": self.notebook_path,
+                "source": self.source,
+                "base_parameters": self.notebook_params,
+            },
+            "libraries": self.notebook_packages,
+        }
+
+    def _get_databricks_task_id(self, task_id: str) -> str:
+        """Get the databricks task ID using dag_id and task_id. Removes illegal characters."""
+        return f"{self.dag_id}__{task_id.replace('.', '__')}"
+
+    def _get_run_json(self) -> dict[str, Any]:
+        """Get run json to be used for task submissions."""
+        run_json = {
+            "run_name": self._get_databricks_task_id(self.task_id),
+            **self._get_task_base_json(),
+        }
+        if self.new_cluster and self.existing_cluster_id:
+            raise ValueError("Both new_cluster and existing_cluster_id are set. Only one should be set.")
+        if self.new_cluster:
+            run_json["new_cluster"] = self.new_cluster
+        elif self.existing_cluster_id:
+            run_json["existing_cluster_id"] = self.existing_cluster_id
+        else:
+            raise ValueError("Must specify either existing_cluster_id or new_cluster.")
+        return run_json
+
+    def launch_notebook_job(self) -> int:
+        run_json = self._get_run_json()
+        self.databricks_run_id = self._hook.submit_run(run_json)
+        url = self._hook.get_run_page_url(self.databricks_run_id)
+        self.log.info("Check the job run in Databricks: %s", url)
+        return self.databricks_run_id
+
+    def monitor_databricks_job(self) -> None:
+        if self.databricks_run_id is None:
+            raise ValueError("Databricks job not yet launched. Please run launch_notebook_job first.")
+        run = self._hook.get_run(self.databricks_run_id)
+        run_state = RunState(**run["state"])
+        self.log.info("Current state of the job: %s", run_state.life_cycle_state)
+        while not run_state.is_terminal:
+            time.sleep(self.polling_period_seconds)
+            run = self._hook.get_run(self.databricks_run_id)
+            run_state = RunState(**run["state"])
+            self.log.info(
+                "task %s %s", self._get_databricks_task_id(self.task_id), run_state.life_cycle_state
+            )
+            self.log.info("Current state of the job: %s", run_state.life_cycle_state)
+        if run_state.life_cycle_state != "TERMINATED":
+            raise AirflowException(
+                f"Databricks job failed with state {run_state.life_cycle_state}. "
+                f"Message: {run_state.state_message}"
+            )
+        if not run_state.is_successful:
+            raise AirflowException(
+                "Task failed. Final state %s. Reason: %s",
+                run_state.result_state,
+                run_state.state_message,
+            )
+        self.log.info("Task succeeded. Final state %s.", run_state.result_state)
+
+    def execute(self, context: Context) -> None:
+        self.launch_notebook_job()
+        if self.wait_for_termination:
+            self.monitor_databricks_job()

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/airflow/providers/databricks/triggers/databricks.py

@@ -84,21 +84,36 @@ class DatabricksExecutionTrigger(BaseTrigger):
         async with self.hook:
             while True:
                 run_state = await self.hook.a_get_run_state(self.run_id)
-                if run_state.is_terminal:
-                    yield TriggerEvent(
-                        {
-                            "run_id": self.run_id,
-                            "run_page_url": self.run_page_url,
-                            "run_state": run_state.to_json(),
-                            "repair_run": self.repair_run,
-                        }
+                if not run_state.is_terminal:
+                    self.log.info(
+                        "run-id %s in run state %s. sleeping for %s seconds",
+                        self.run_id,
+                        run_state,
+                        self.polling_period_seconds,
                     )
-                    return
+                    await asyncio.sleep(self.polling_period_seconds)
+                    continue
 
-                self.log.info(
-                    "run-id %s in run state %s. sleeping for %s seconds",
-                    self.run_id,
-                    run_state,
-                    self.polling_period_seconds,
+                failed_tasks = []
+                if run_state.result_state == "FAILED":
+                    run_info = await self.hook.a_get_run(self.run_id)
+                    for task in run_info.get("tasks", []):
+                        if task.get("state", {}).get("result_state", "") == "FAILED":
+                            task_run_id = task["run_id"]
+                            task_key = task["task_key"]
+                            run_output = await self.hook.a_get_run_output(task_run_id)
+                            if "error" in run_output:
+                                error = run_output["error"]
+                            else:
+                                error = run_state.state_message
+                            failed_tasks.append({"task_key": task_key, "run_id": task_run_id, "error": error})
+                yield TriggerEvent(
+                    {
+                        "run_id": self.run_id,
+                        "run_page_url": self.run_page_url,
+                        "run_state": run_state.to_json(),
+                        "repair_run": self.repair_run,
+                        "errors": failed_tasks,
+                    }
                 )
-                await asyncio.sleep(self.polling_period_seconds)
+                return

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/airflow/providers/databricks/utils/databricks.py

@@ -55,7 +55,7 @@ def validate_trigger_event(event: dict):
 
     See: :class:`~airflow.providers.databricks.triggers.databricks.DatabricksExecutionTrigger`.
     """
-    keys_to_check = ["run_id", "run_page_url", "run_state"]
+    keys_to_check = ["run_id", "run_page_url", "run_state", "errors"]
     for key in keys_to_check:
         if key not in event:
             raise AirflowException(f"Could not find `{key}` in the event: {event}")

{apache_airflow_providers_databricks-6.3.0rc3 → apache_airflow_providers_databricks-6.4.0}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-databricks"
-version = "6.3.0.rc3"
+version = "6.4.0"
 description = "Provider package apache-airflow-providers-databricks for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -57,15 +57,15 @@ classifiers = [
 requires-python = "~=3.8"
 dependencies = [
     "aiohttp>=3.9.2, <4",
-    "apache-airflow-providers-common-sql>=1.10.0rc0",
-    "apache-airflow>=2.6.0rc0",
+    "apache-airflow-providers-common-sql>=1.10.0",
+    "apache-airflow>=2.7.0",
     "databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0",
     "requests>=2.27.0,<3",
 ]
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.3.0"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.3.0/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"