apache-airflow-providers-databricks 4.4.0__py3-none-any.whl → 4.5.0__py3-none-any.whl

airflow/providers/databricks/__init__.py

@@ -28,7 +28,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "4.4.0"
+__version__ = "4.5.0"
 
 try:
     from airflow import __version__ as airflow_version

airflow/providers/databricks/get_provider_info.py

@@ -29,6 +29,7 @@ def get_provider_info():
         "description": "`Databricks <https://databricks.com/>`__\n",
         "suspended": False,
         "versions": [
+            "4.5.0",
             "4.4.0",
             "4.3.3",
             "4.3.2",

airflow/providers/databricks/hooks/databricks.py

@@ -52,10 +52,19 @@ INSTALL_LIBS_ENDPOINT = ("POST", "api/2.0/libraries/install")
 UNINSTALL_LIBS_ENDPOINT = ("POST", "api/2.0/libraries/uninstall")
 
 LIST_JOBS_ENDPOINT = ("GET", "api/2.1/jobs/list")
+LIST_PIPELINES_ENDPOINT = ("GET", "/api/2.0/pipelines")
 
 WORKSPACE_GET_STATUS_ENDPOINT = ("GET", "api/2.0/workspace/get-status")
 
-RUN_LIFE_CYCLE_STATES = ["PENDING", "RUNNING", "TERMINATING", "TERMINATED", "SKIPPED", "INTERNAL_ERROR"]
+RUN_LIFE_CYCLE_STATES = [
+    "PENDING",
+    "RUNNING",
+    "TERMINATING",
+    "TERMINATED",
+    "SKIPPED",
+    "INTERNAL_ERROR",
+    "QUEUED",
+]
 
 SPARK_VERSIONS_ENDPOINT = ("GET", "api/2.0/clusters/spark-versions")
 
@@ -75,11 +84,9 @@ class RunState:
         """True if the current state is a terminal state."""
         if self.life_cycle_state not in RUN_LIFE_CYCLE_STATES:
             raise AirflowException(
-                (
-                    "Unexpected life cycle state: {}: If the state has "
-                    "been introduced recently, please check the Databricks user "
-                    "guide for troubleshooting information"
-                ).format(self.life_cycle_state)
+                f"Unexpected life cycle state: {self.life_cycle_state}: If the state has "
+                "been introduced recently, please check the Databricks user "
+                "guide for troubleshooting information"
             )
         return self.life_cycle_state in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR")
 
@@ -209,6 +216,67 @@ class DatabricksHook(BaseDatabricksHook):
         else:
             return matching_jobs[0]["job_id"]
 
+    def list_pipelines(
+        self, batch_size: int = 25, pipeline_name: str | None = None, notebook_path: str | None = None
+    ) -> list[dict[str, Any]]:
+        """
+        Lists the pipelines in Databricks Delta Live Tables.
+
+        :param batch_size: The limit/batch size used to retrieve pipelines.
+        :param pipeline_name: Optional name of a pipeline to search. Cannot be combined with path.
+        :param notebook_path: Optional notebook of a pipeline to search. Cannot be combined with name.
+        :return: A list of pipelines.
+        """
+        has_more = True
+        next_token = None
+        all_pipelines = []
+        filter = None
+        if pipeline_name and notebook_path:
+            raise AirflowException("Cannot combine pipeline_name and notebook_path in one request")
+
+        if notebook_path:
+            filter = f"notebook='{notebook_path}'"
+        elif pipeline_name:
+            filter = f"name LIKE '{pipeline_name}'"
+        payload: dict[str, Any] = {
+            "max_results": batch_size,
+        }
+        if filter:
+            payload["filter"] = filter
+
+        while has_more:
+            if next_token:
+                payload["page_token"] = next_token
+            response = self._do_api_call(LIST_PIPELINES_ENDPOINT, payload)
+            pipelines = response.get("statuses", [])
+            all_pipelines += pipelines
+            if "next_page_token" in response:
+                next_token = response["next_page_token"]
+            else:
+                has_more = False
+
+        return all_pipelines
+
+    def find_pipeline_id_by_name(self, pipeline_name: str) -> str | None:
+        """
+        Finds pipeline id by its name. If multiple pipelines with the same name, raises AirflowException.
+
+        :param pipeline_name: The name of the pipeline to look up.
+        :return: The pipeline_id as a GUID string or None if no pipeline was found.
+        """
+        matching_pipelines = self.list_pipelines(pipeline_name=pipeline_name)
+
+        if len(matching_pipelines) > 1:
+            raise AirflowException(
+                f"There are more than one job with name {pipeline_name}. "
+                "Please delete duplicated pipelines first"
+            )
+
+        if not pipeline_name:
+            return None
+        else:
+            return matching_pipelines[0]["pipeline_id"]
+
     def get_run_page_url(self, run_id: int) -> str:
         """
         Retrieves run_page_url.

airflow/providers/databricks/hooks/databricks_base.py

@@ -28,7 +28,7 @@ import copy
 import platform
 import time
 from functools import cached_property
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urlsplit
 
 import aiohttp
@@ -48,9 +48,11 @@ from tenacity import (
 from airflow import __version__
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
-from airflow.models import Connection
 from airflow.providers_manager import ProvidersManager
 
+if TYPE_CHECKING:
+    from airflow.models import Connection
+
 # https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/aad/service-prin-aad-token#--get-an-azure-active-directory-access-token
 # https://docs.microsoft.com/en-us/graph/deployments#app-registration-and-token-service-root-endpoints
 AZURE_DEFAULT_AD_ENDPOINT = "https://login.microsoftonline.com"
@@ -121,12 +123,12 @@ class BaseDatabricksHook(BaseHook):
             self.retry_args["retry"] = retry_if_exception(self._retryable_error)
             self.retry_args["after"] = my_after_func
         else:
-            self.retry_args = dict(
-                stop=stop_after_attempt(self.retry_limit),
-                wait=wait_exponential(min=self.retry_delay, max=(2**retry_limit)),
-                retry=retry_if_exception(self._retryable_error),
-                after=my_after_func,
-            )
+            self.retry_args = {
+                "stop": stop_after_attempt(self.retry_limit),
+                "wait": wait_exponential(min=self.retry_delay, max=(2**retry_limit)),
+                "retry": retry_if_exception(self._retryable_error),
+                "after": my_after_func,
+            }
 
     @cached_property
     def databricks_conn(self) -> Connection:

airflow/providers/databricks/hooks/databricks_sql.py

@@ -18,15 +18,17 @@ from __future__ import annotations
 
 from contextlib import closing
 from copy import copy
-from typing import Any, Callable, Iterable, Mapping, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, TypeVar, overload
 
 from databricks import sql  # type: ignore[attr-defined]
-from databricks.sql.client import Connection  # type: ignore[attr-defined]
 
 from airflow.exceptions import AirflowException
 from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results
 from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook
 
+if TYPE_CHECKING:
+    from databricks.sql.client import Connection
+
 LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "api/2.0/sql/endpoints")
 
 
@@ -81,7 +83,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
 
     def _get_extra_config(self) -> dict[str, Any | None]:
         extra_params = copy(self.databricks_conn.extra_dejson)
-        for arg in ["http_path", "session_configuration"] + self.extra_parameters:
+        for arg in ["http_path", "session_configuration", *self.extra_parameters]:
             if arg in extra_params:
                 del extra_params[arg]
 

airflow/providers/databricks/operators/databricks.py

@@ -21,7 +21,6 @@ from __future__ import annotations
 import time
 import warnings
 from functools import cached_property
-from logging import Logger
 from typing import TYPE_CHECKING, Any, Sequence
 
 from airflow.configuration import conf
@@ -32,6 +31,8 @@ from airflow.providers.databricks.triggers.databricks import DatabricksExecution
 from airflow.providers.databricks.utils.databricks import normalise_json_content, validate_trigger_event
 
 if TYPE_CHECKING:
+    from logging import Logger
+
     from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.utils.context import Context
 
@@ -364,6 +365,8 @@ class DatabricksSubmitRunOperator(BaseOperator):
 
         if "dbt_task" in self.json and "git_source" not in self.json:
             raise AirflowException("git_source is required for dbt_task")
+        if pipeline_task is not None and "pipeline_id" in pipeline_task and "pipeline_name" in pipeline_task:
+            raise AirflowException("'pipeline_name' is not allowed in conjunction with 'pipeline_id'")
 
         # This variable will be used in case our task gets killed.
         self.run_id: int | None = None
@@ -383,6 +386,15 @@ class DatabricksSubmitRunOperator(BaseOperator):
         )
 
     def execute(self, context: Context):
+        if (
+            "pipeline_task" in self.json
+            and self.json["pipeline_task"].get("pipeline_id") is None
+            and self.json["pipeline_task"].get("pipeline_name")
+        ):
+            # If pipeline_id is not provided, we need to fetch it from the pipeline_name
+            pipeline_name = self.json["pipeline_task"]["pipeline_name"]
+            self.json["pipeline_task"]["pipeline_id"] = self._hook.get_pipeline_id(pipeline_name)
+            del self.json["pipeline_task"]["pipeline_name"]
         json_normalised = normalise_json_content(self.json)
         self.run_id = self._hook.submit_run(json_normalised)
         if self.deferrable:

airflow/providers/databricks/operators/databricks_sql.py

@@ -22,7 +22,6 @@ import csv
 import json
 from typing import TYPE_CHECKING, Any, Sequence
 
-from databricks.sql.types import Row
 from databricks.sql.utils import ParamEscaper
 
 from airflow.exceptions import AirflowException
@@ -31,6 +30,8 @@ from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook
 
 if TYPE_CHECKING:
+    from databricks.sql.types import Row
+
     from airflow.utils.context import Context
 
 

{apache_airflow_providers_databricks-4.4.0.dist-info → apache_airflow_providers_databricks-4.5.0.dist-info}/METADATA

@@ -1,14 +1,14 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 4.4.0
+Version: 4.5.0
 Summary: Provider for Apache Airflow. Implements apache-airflow-providers-databricks package
 Home-page: https://airflow.apache.org/
 Download-URL: https://archive.apache.org/dist/airflow/providers
 Author: Apache Software Foundation
 Author-email: dev@airflow.apache.org
 License: Apache License 2.0
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.4.0/
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.4.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.5.0/
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.5.0/changelog.html
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
@@ -77,7 +77,7 @@ Requires-Dist: apache-airflow-providers-common-sql ; extra == 'common.sql'
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``4.4.0``
+Release: ``4.5.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -90,7 +90,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/4.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.5.0/>`_.
 
 
 Installation
@@ -135,4 +135,4 @@ Dependent package
 ============================================================================================================  ==============
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.5.0/changelog.html>`_.

{apache_airflow_providers_databricks-4.4.0.dist-info → apache_airflow_providers_databricks-4.5.0.dist-info}/RECORD

@@ -1,13 +1,13 @@
-airflow/providers/databricks/__init__.py,sha256=eBT7V41nReeyh7F3ZOola64x2XQdoHS4XttOxrZsVro,1579
-airflow/providers/databricks/get_provider_info.py,sha256=-2a0F4ZpFSh-1rqVmqxklhaO4Ds7AHBNwi032qdDqXc,5599
+airflow/providers/databricks/__init__.py,sha256=f4tIs0EPxcFXuETkMe32PU1hiGP-frvdB9iD2QB-FI8,1579
+airflow/providers/databricks/get_provider_info.py,sha256=QKheBxRexRHoc7JFTQ3Vrx_3RrCesJNvmJoy3DfSFL4,5620
 airflow/providers/databricks/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/hooks/databricks.py,sha256=0gZtO9mlBGOQ-f249wF1vFvf4Q10It0gcKrIsS3Zb68,16773
-airflow/providers/databricks/hooks/databricks_base.py,sha256=PvWp1RcQQR5riJ0SWyTVUu2jY-Si2ohb6xYTaaHsiUg,30535
-airflow/providers/databricks/hooks/databricks_sql.py,sha256=1mJJG7eBXYAInPPGhuWFgo6pG3ykphoLCZNOGZVloQY,10083
+airflow/providers/databricks/hooks/databricks.py,sha256=2m0BPmhYbMwI_pKBmmvYL636ciYX1qYfkFfgkcTpk_c,19158
+airflow/providers/databricks/hooks/databricks_base.py,sha256=R8NCyXWPhji9lo0X0zM9JY9geprl95XpAJQXwfFjFcM,30581
+airflow/providers/databricks/hooks/databricks_sql.py,sha256=hCDJNousYOvxs3n08las3pz7ATo5bRXOoYI8ujDPmmI,10091
 airflow/providers/databricks/operators/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/operators/databricks.py,sha256=RdDQzT2yAYxrRYK56ZJ8RB1o6TxNjZqPUjoBQkFNpN0,33338
+airflow/providers/databricks/operators/databricks.py,sha256=gLrqiIxJ2s0P_MgPikbCQtU2iGD_K9PiBbJBpi8jjUI,34074
 airflow/providers/databricks/operators/databricks_repos.py,sha256=Z7bAB0HiQajfZghmwK9RRGqqumx5S5JJyOB86r6e23s,13096
-airflow/providers/databricks/operators/databricks_sql.py,sha256=kcDRJ9ljTmWtN5c7eoM1JlydQ1nlNp30KdqBrFDRmLo,17021
+airflow/providers/databricks/operators/databricks_sql.py,sha256=-_txO63CR7nbjNR4lkq2o7nosw19qIM5qBpxkOjh_K0,17026
 airflow/providers/databricks/sensors/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/databricks/sensors/databricks_partition.py,sha256=bmLZWn8hsAZp4Uooan5mP1H4uR1Dek7TgcNx1ffmssI,9787
 airflow/providers/databricks/sensors/databricks_sql.py,sha256=ESLhcaWHbW9xjsw2IwLHa3QAVnBs3gmsb9SYjs2UgIk,5543
@@ -15,10 +15,10 @@ airflow/providers/databricks/triggers/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvR
 airflow/providers/databricks/triggers/databricks.py,sha256=Qj9mB0bNYRY_toPEU17gxbxmPkkT3P789kCHu_T64BA,3997
 airflow/providers/databricks/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/databricks/utils/databricks.py,sha256=6d8zhuaKD-0G4ef-m8duzlZBrChcCa9hvxWb8pKjaqI,2881
-apache_airflow_providers_databricks-4.4.0.dist-info/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
-apache_airflow_providers_databricks-4.4.0.dist-info/METADATA,sha256=fs326d1zZYUjpLg4u8Ec3lnecKOr0K7SLPwTk8UE1AA,6143
-apache_airflow_providers_databricks-4.4.0.dist-info/NOTICE,sha256=m-6s2XynUxVSUIxO4rVablAZCvFq-wmLrqV91DotRBw,240
-apache_airflow_providers_databricks-4.4.0.dist-info/WHEEL,sha256=5sUXSg9e4bi7lTLOHcm6QEYwO5TIF1TNbTSVFVjcJcc,92
-apache_airflow_providers_databricks-4.4.0.dist-info/entry_points.txt,sha256=8r3YBg2Qr0qeOALxzlooH5pXM6QmoPQuyQ75cQrkY5A,107
-apache_airflow_providers_databricks-4.4.0.dist-info/top_level.txt,sha256=OeMVH5md7fr2QQWpnZoOWWxWO-0WH1IP70lpTVwopPg,8
-apache_airflow_providers_databricks-4.4.0.dist-info/RECORD,,
+apache_airflow_providers_databricks-4.5.0.dist-info/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
+apache_airflow_providers_databricks-4.5.0.dist-info/METADATA,sha256=NytwRUe06fdVp9k1XvajC_BLbxMmFksxKEryLIbff_g,6143
+apache_airflow_providers_databricks-4.5.0.dist-info/NOTICE,sha256=m-6s2XynUxVSUIxO4rVablAZCvFq-wmLrqV91DotRBw,240
+apache_airflow_providers_databricks-4.5.0.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+apache_airflow_providers_databricks-4.5.0.dist-info/entry_points.txt,sha256=8r3YBg2Qr0qeOALxzlooH5pXM6QmoPQuyQ75cQrkY5A,107
+apache_airflow_providers_databricks-4.5.0.dist-info/top_level.txt,sha256=OeMVH5md7fr2QQWpnZoOWWxWO-0WH1IP70lpTVwopPg,8
+apache_airflow_providers_databricks-4.5.0.dist-info/RECORD,,

{apache_airflow_providers_databricks-4.4.0.dist-info → apache_airflow_providers_databricks-4.5.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.1)
+Generator: bdist_wheel (0.41.2)
 Root-Is-Purelib: true
 Tag: py3-none-any