Lee-W commented on code in PR #39771:
URL: https://github.com/apache/airflow/pull/39771#discussion_r1618659959
##########
airflow/providers/databricks/operators/databricks.py:
##########
@@ -1039,6 +1123,17 @@ def launch_notebook_job(self) -> int:
self.log.info("Check the job run in Databricks: %s", url)
return self.databricks_run_id
+ def _handle_terminal_run_state(self, run_state):
Review Comment:
```suggestion
def _handle_terminal_run_state(self, run_state: RunState) -> None:
```
##########
airflow/providers/databricks/operators/databricks.py:
##########
@@ -1016,6 +1029,77 @@ 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('.', '__')}"
+ @property
+ def _databricks_workflow_task_group(self) -> DatabricksWorkflowTaskGroup |
None:
+ """
+ Traverse up parent TaskGroups until the `is_databricks` flag
associated with the root DatabricksWorkflowTaskGroup is found.
+
+ If found, returns the task group. Otherwise, return None.
+ """
+ parent_tg: TaskGroup | DatabricksWorkflowTaskGroup | None =
self.task_group
+
+ while parent_tg:
+ if getattr(parent_tg, "is_databricks", False):
+ return parent_tg # type: ignore[return-value]
+
+ if getattr(parent_tg, "task_group", None):
+ parent_tg = parent_tg.task_group
+ else:
+ return None
+
+ return None
+
+ def _extend_workflow_notebook_packages(self,
databricks_workflow_task_group: DatabricksWorkflowTaskGroup):
+ """Extend the task group packages into the notebook's packages,
without adding any duplicates."""
+ for task_group_package in
databricks_workflow_task_group.notebook_packages:
+ exists = any(
+ task_group_package == existing_package for existing_package in
self.notebook_packages
+ )
+ if not exists:
+ self.notebook_packages.append(task_group_package)
+
+ def _convert_to_databricks_workflow_task(
+ self, relevant_upstreams: list[BaseOperator], context: Context | None
= None
+ ):
Review Comment:
```suggestion
) -> None:
```
##########
airflow/providers/databricks/operators/databricks.py:
##########
@@ -1039,6 +1123,17 @@ def launch_notebook_job(self) -> int:
self.log.info("Check the job run in Databricks: %s", url)
return self.databricks_run_id
+ def _handle_terminal_run_state(self, run_state):
+ if run_state.life_cycle_state != "TERMINATED":
Review Comment:
Should we use `RunState.is_terminal`?
##########
airflow/providers/databricks/operators/databricks.py:
##########
@@ -1016,6 +1029,77 @@ 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('.', '__')}"
+ @property
+ def _databricks_workflow_task_group(self) -> DatabricksWorkflowTaskGroup |
None:
+ """
+ Traverse up parent TaskGroups until the `is_databricks` flag
associated with the root DatabricksWorkflowTaskGroup is found.
+
+ If found, returns the task group. Otherwise, return None.
+ """
+ parent_tg: TaskGroup | DatabricksWorkflowTaskGroup | None =
self.task_group
+
+ while parent_tg:
+ if getattr(parent_tg, "is_databricks", False):
+ return parent_tg # type: ignore[return-value]
+
+ if getattr(parent_tg, "task_group", None):
+ parent_tg = parent_tg.task_group
+ else:
+ return None
+
+ return None
+
+ def _extend_workflow_notebook_packages(self,
databricks_workflow_task_group: DatabricksWorkflowTaskGroup):
Review Comment:
```suggestion
def _extend_workflow_notebook_packages(self,
databricks_workflow_task_group: DatabricksWorkflowTaskGroup) -> None:
```
##########
tests/system/providers/databricks/example_databricks_workflow.py:
##########
@@ -0,0 +1,120 @@
+# 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.
+
+"""Example DAG for using the DatabricksWorkflowTaskGroup and
DatabricksNotebookOperator."""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+
+from airflow.models.dag import DAG
+from airflow.providers.databricks.operators.databricks import
DatabricksNotebookOperator
+from airflow.providers.databricks.operators.databricks_workflow import
DatabricksWorkflowTaskGroup
+from airflow.utils.timezone import datetime
+
+EXECUTION_TIMEOUT = int(os.getenv("EXECUTION_TIMEOUT", 6))
+
+DATABRICKS_CONN_ID = os.getenv("ASTRO_DATABRICKS_CONN_ID", "databricks_conn")
+DATABRICKS_NOTIFICATION_EMAIL = os.getenv(
+ "ASTRO_DATABRICKS_NOTIFICATION_EMAIL", "[email protected]"
Review Comment:
```suggestion
"DATABRICKS_NOTIFICATION_EMAIL", "[email protected]"
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
+ return f"{self.dag_id}.{self.task_group.group_id}"
+
+ def create_workflow_json(self, context: Context | None = None) ->
dict[str, object]:
+ """Create a workflow json to be used in the Databricks API."""
+ task_json = [
+ task._convert_to_databricks_workflow_task( # type:
ignore[attr-defined]
+ relevant_upstreams=self.relevant_upstreams, context=context
+ )
+ for task in self.tasks_to_convert
+ ]
+
+ default_json = {
+ "name": self.databricks_job_name,
+ "email_notifications": {"no_alert_for_skipped_runs": False},
+ "timeout_seconds": 0,
+ "tasks": task_json,
+ "format": "MULTI_TASK",
+ "job_clusters": self.job_clusters,
+ "max_concurrent_runs": self.max_concurrent_runs,
+ }
+ return merge(default_json, self.extra_job_params)
+
+ def _create_or_reset_job(self, context: Context) -> int:
+ job_spec = self.create_workflow_json(context=context)
+ existing_jobs = self._hook.list_jobs(job_name=self.databricks_job_name)
+ job_id = existing_jobs[0]["job_id"] if existing_jobs else None
+ if job_id:
+ self.log.info(
+ "Updating existing Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ self._hook.reset_job(job_id, job_spec)
+ else:
+ self.log.info(
+ "Creating new Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ job_id = self._hook.create_job(job_spec)
+ return job_id
+
+ def _wait_for_job_to_start(self, run_id: int):
Review Comment:
```suggestion
def _wait_for_job_to_start(self, run_id: int) -> None:
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
+ return f"{self.dag_id}.{self.task_group.group_id}"
+
+ def create_workflow_json(self, context: Context | None = None) ->
dict[str, object]:
Review Comment:
What's the difference between `dict[str, object]` and `dict[str, Any]` which
seems to be more commonly used
`
##########
tests/system/providers/databricks/example_databricks_workflow.py:
##########
@@ -0,0 +1,120 @@
+# 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.
+
+"""Example DAG for using the DatabricksWorkflowTaskGroup and
DatabricksNotebookOperator."""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+
+from airflow.models.dag import DAG
+from airflow.providers.databricks.operators.databricks import
DatabricksNotebookOperator
+from airflow.providers.databricks.operators.databricks_workflow import
DatabricksWorkflowTaskGroup
+from airflow.utils.timezone import datetime
+
+EXECUTION_TIMEOUT = int(os.getenv("EXECUTION_TIMEOUT", 6))
+
+DATABRICKS_CONN_ID = os.getenv("ASTRO_DATABRICKS_CONN_ID", "databricks_conn")
+DATABRICKS_NOTIFICATION_EMAIL = os.getenv(
+ "ASTRO_DATABRICKS_NOTIFICATION_EMAIL", "[email protected]"
+)
+
+GROUP_ID = os.getenv("DATABRICKS_GROUP_ID", "1234").replace(".", "_")
+USER = os.environ.get("USER")
+
+job_cluster_spec = [
+ {
+ "job_cluster_key": "Shared_job_cluster",
+ "new_cluster": {
+ "cluster_name": "",
+ "spark_version": "11.3.x-scala2.12",
+ "aws_attributes": {
+ "first_on_demand": 1,
+ "availability": "SPOT_WITH_FALLBACK",
+ "zone_id": "us-east-2b",
+ "spot_bid_price_percent": 100,
+ "ebs_volume_count": 0,
+ },
+ "node_type_id": "i3.xlarge",
+ "spark_env_vars": {"PYSPARK_PYTHON":
"/databricks/python3/bin/python3"},
+ "enable_elastic_disk": False,
+ "data_security_mode": "LEGACY_SINGLE_USER_STANDARD",
+ "runtime_engine": "STANDARD",
+ "num_workers": 8,
+ },
+ }
+]
+dag = DAG(
+ dag_id="example_databricks_workflow",
+ start_date=datetime(2022, 1, 1),
+ schedule_interval=None,
+ catchup=False,
+ tags=["example", "async", "databricks"],
Review Comment:
```suggestion
tags=["example", "databricks"],
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
Review Comment:
Does it make sense to name it as `job_name` since it's already in a
databricks related operator
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
Review Comment:
```suggestion
def databricks_job_name(self) -> str:
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
Review Comment:
```suggestion
task_id: str,
databricks_conn_id: str,
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
+ return f"{self.dag_id}.{self.task_group.group_id}"
+
+ def create_workflow_json(self, context: Context | None = None) ->
dict[str, object]:
+ """Create a workflow json to be used in the Databricks API."""
+ task_json = [
+ task._convert_to_databricks_workflow_task( # type:
ignore[attr-defined]
+ relevant_upstreams=self.relevant_upstreams, context=context
+ )
+ for task in self.tasks_to_convert
+ ]
+
+ default_json = {
+ "name": self.databricks_job_name,
+ "email_notifications": {"no_alert_for_skipped_runs": False},
+ "timeout_seconds": 0,
+ "tasks": task_json,
+ "format": "MULTI_TASK",
+ "job_clusters": self.job_clusters,
+ "max_concurrent_runs": self.max_concurrent_runs,
+ }
+ return merge(default_json, self.extra_job_params)
+
+ def _create_or_reset_job(self, context: Context) -> int:
+ job_spec = self.create_workflow_json(context=context)
+ existing_jobs = self._hook.list_jobs(job_name=self.databricks_job_name)
+ job_id = existing_jobs[0]["job_id"] if existing_jobs else None
+ if job_id:
+ self.log.info(
+ "Updating existing Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ self._hook.reset_job(job_id, job_spec)
+ else:
+ self.log.info(
+ "Creating new Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ job_id = self._hook.create_job(job_spec)
+ return job_id
+
+ def _wait_for_job_to_start(self, run_id: int):
+ run_url = self._hook.get_run_page_url(run_id)
+ self.log.info("Check the progress of the Databricks job at %s",
run_url)
+ life_cycle_state = self._hook.get_run_state(run_id).life_cycle_state
+ if life_cycle_state not in ("PENDING", "RUNNING", "BLOCKED"):
+ raise AirflowException(f"Could not start the workflow job. State:
{life_cycle_state}")
+ while life_cycle_state in ("PENDING", "BLOCKED"):
+ self.log.info("Waiting for the Databricks job to start running")
+ time.sleep(5)
+ life_cycle_state =
self._hook.get_run_state(run_id).life_cycle_state
+ self.log.info("Databricks job started. State: %s", life_cycle_state)
+
+ def execute(self, context: Context) -> Any:
+ if not isinstance(self.task_group, DatabricksWorkflowTaskGroup):
+ raise AirflowException("Task group must be a
DatabricksWorkflowTaskGroup")
+
+ job_id = self._create_or_reset_job(context)
+
+ run_id = self._hook.run_now(
+ {
+ "job_id": job_id,
+ "jar_params": self.task_group.jar_params,
+ "notebook_params": self.notebook_params,
+ "python_params": self.task_group.python_params,
+ "spark_submit_params": self.task_group.spark_submit_params,
+ }
+ )
+
+ self._wait_for_job_to_start(run_id)
+
+ return {
+ "conn_id": self.databricks_conn_id,
+ "job_id": job_id,
+ "run_id": run_id,
+ }
+
+
+class DatabricksWorkflowTaskGroup(TaskGroup):
+ """
+ A task group that takes a list of tasks and creates a databricks workflow.
+
+ The DatabricksWorkflowTaskGroup takes a list of tasks and creates a
databricks workflow
+ based on the metadata produced by those tasks. For a task to be eligible
for this
+ TaskGroup, it must contain the ``_convert_to_databricks_workflow_task``
method. If any tasks
+ do not contain this method then the Taskgroup will raise an error at parse
time.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DatabricksWorkflowTaskGroup`
+
+ :param databricks_conn_id: The name of the databricks connection to use.
+ :param existing_clusters: A list of existing clusters to use for this
workflow.
+ :param extra_job_params: A dictionary containing properties which will
override the default
+ Databricks Workflow Job definitions.
+ :param jar_params: A list of jar parameters to pass to the workflow. These
parameters will be passed to all jar
+ tasks in the workflow.
+ :param job_clusters: A list of job clusters to use for this workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for this
workflow.
+ :param notebook_packages: A list of dictionary of Python packages to be
installed. Packages defined
+ at the workflow task group level are installed for each of the
notebook tasks under it. And
+ packages defined at the notebook task level are installed specific for
the notebook task.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebook tasks in the workflow.
+ :param python_params: A list of python parameters to pass to the workflow.
These parameters will be passed to
+ all python tasks in the workflow.
+ :param spark_submit_params: A list of spark submit parameters to pass to
the workflow. These parameters
+ will be passed to all spark submit tasks.
+ """
+
+ is_databricks = True
+
+ def __init__(
+ self,
+ databricks_conn_id,
+ existing_clusters=None,
+ extra_job_params: dict[str, Any] | None = None,
+ jar_params=None,
+ job_clusters=None,
+ max_concurrent_runs: int = 1,
+ notebook_packages: list[dict[str, Any]] | None = None,
+ notebook_params: dict | None = None,
+ python_params: list | None = None,
+ spark_submit_params: list | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.jar_params = jar_params or []
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_packages = notebook_packages or []
+ self.notebook_params = notebook_params or {}
+ self.python_params = python_params or []
+ self.spark_submit_params = spark_submit_params or []
+ super().__init__(**kwargs)
+
+ def __exit__(self, _type, _value, _tb):
Review Comment:
```suggestion
def __exit__(self, _type: type[BaseException] | None, _value:
BaseException | None, _tb: TracebackType | None) -> None:
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
+ return f"{self.dag_id}.{self.task_group.group_id}"
+
+ def create_workflow_json(self, context: Context | None = None) ->
dict[str, object]:
+ """Create a workflow json to be used in the Databricks API."""
+ task_json = [
+ task._convert_to_databricks_workflow_task( # type:
ignore[attr-defined]
+ relevant_upstreams=self.relevant_upstreams, context=context
+ )
+ for task in self.tasks_to_convert
+ ]
+
+ default_json = {
+ "name": self.databricks_job_name,
+ "email_notifications": {"no_alert_for_skipped_runs": False},
+ "timeout_seconds": 0,
+ "tasks": task_json,
+ "format": "MULTI_TASK",
+ "job_clusters": self.job_clusters,
+ "max_concurrent_runs": self.max_concurrent_runs,
+ }
+ return merge(default_json, self.extra_job_params)
+
+ def _create_or_reset_job(self, context: Context) -> int:
+ job_spec = self.create_workflow_json(context=context)
+ existing_jobs = self._hook.list_jobs(job_name=self.databricks_job_name)
+ job_id = existing_jobs[0]["job_id"] if existing_jobs else None
+ if job_id:
+ self.log.info(
+ "Updating existing Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ self._hook.reset_job(job_id, job_spec)
+ else:
+ self.log.info(
+ "Creating new Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ job_id = self._hook.create_job(job_spec)
+ return job_id
+
+ def _wait_for_job_to_start(self, run_id: int):
+ run_url = self._hook.get_run_page_url(run_id)
+ self.log.info("Check the progress of the Databricks job at %s",
run_url)
+ life_cycle_state = self._hook.get_run_state(run_id).life_cycle_state
+ if life_cycle_state not in ("PENDING", "RUNNING", "BLOCKED"):
+ raise AirflowException(f"Could not start the workflow job. State:
{life_cycle_state}")
+ while life_cycle_state in ("PENDING", "BLOCKED"):
+ self.log.info("Waiting for the Databricks job to start running")
+ time.sleep(5)
+ life_cycle_state =
self._hook.get_run_state(run_id).life_cycle_state
+ self.log.info("Databricks job started. State: %s", life_cycle_state)
+
+ def execute(self, context: Context) -> Any:
+ if not isinstance(self.task_group, DatabricksWorkflowTaskGroup):
+ raise AirflowException("Task group must be a
DatabricksWorkflowTaskGroup")
+
+ job_id = self._create_or_reset_job(context)
+
+ run_id = self._hook.run_now(
+ {
+ "job_id": job_id,
+ "jar_params": self.task_group.jar_params,
+ "notebook_params": self.notebook_params,
+ "python_params": self.task_group.python_params,
+ "spark_submit_params": self.task_group.spark_submit_params,
+ }
+ )
+
+ self._wait_for_job_to_start(run_id)
+
+ return {
+ "conn_id": self.databricks_conn_id,
+ "job_id": job_id,
+ "run_id": run_id,
+ }
+
+
+class DatabricksWorkflowTaskGroup(TaskGroup):
+ """
+ A task group that takes a list of tasks and creates a databricks workflow.
+
+ The DatabricksWorkflowTaskGroup takes a list of tasks and creates a
databricks workflow
+ based on the metadata produced by those tasks. For a task to be eligible
for this
+ TaskGroup, it must contain the ``_convert_to_databricks_workflow_task``
method. If any tasks
+ do not contain this method then the Taskgroup will raise an error at parse
time.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DatabricksWorkflowTaskGroup`
+
+ :param databricks_conn_id: The name of the databricks connection to use.
+ :param existing_clusters: A list of existing clusters to use for this
workflow.
+ :param extra_job_params: A dictionary containing properties which will
override the default
+ Databricks Workflow Job definitions.
+ :param jar_params: A list of jar parameters to pass to the workflow. These
parameters will be passed to all jar
+ tasks in the workflow.
+ :param job_clusters: A list of job clusters to use for this workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for this
workflow.
+ :param notebook_packages: A list of dictionary of Python packages to be
installed. Packages defined
+ at the workflow task group level are installed for each of the
notebook tasks under it. And
+ packages defined at the notebook task level are installed specific for
the notebook task.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebook tasks in the workflow.
+ :param python_params: A list of python parameters to pass to the workflow.
These parameters will be passed to
+ all python tasks in the workflow.
+ :param spark_submit_params: A list of spark submit parameters to pass to
the workflow. These parameters
+ will be passed to all spark submit tasks.
+ """
+
+ is_databricks = True
+
+ def __init__(
+ self,
+ databricks_conn_id,
+ existing_clusters=None,
+ extra_job_params: dict[str, Any] | None = None,
+ jar_params=None,
+ job_clusters=None,
Review Comment:
```suggestion
job_clusters: list[dict] | None = None,
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
Review Comment:
```suggestion
def add_task(self, task: BaseOperator) -> None:
```
##########
tests/system/providers/databricks/example_databricks_workflow.py:
##########
@@ -0,0 +1,120 @@
+# 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.
+
+"""Example DAG for using the DatabricksWorkflowTaskGroup and
DatabricksNotebookOperator."""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+
+from airflow.models.dag import DAG
+from airflow.providers.databricks.operators.databricks import
DatabricksNotebookOperator
+from airflow.providers.databricks.operators.databricks_workflow import
DatabricksWorkflowTaskGroup
+from airflow.utils.timezone import datetime
+
+EXECUTION_TIMEOUT = int(os.getenv("EXECUTION_TIMEOUT", 6))
+
+DATABRICKS_CONN_ID = os.getenv("ASTRO_DATABRICKS_CONN_ID", "databricks_conn")
Review Comment:
```suggestion
DATABRICKS_CONN_ID = os.getenv("DATABRICKS_CONN_ID", "databricks_conn")
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
+ return f"{self.dag_id}.{self.task_group.group_id}"
+
+ def create_workflow_json(self, context: Context | None = None) ->
dict[str, object]:
+ """Create a workflow json to be used in the Databricks API."""
+ task_json = [
+ task._convert_to_databricks_workflow_task( # type:
ignore[attr-defined]
+ relevant_upstreams=self.relevant_upstreams, context=context
+ )
+ for task in self.tasks_to_convert
+ ]
+
+ default_json = {
+ "name": self.databricks_job_name,
+ "email_notifications": {"no_alert_for_skipped_runs": False},
+ "timeout_seconds": 0,
+ "tasks": task_json,
+ "format": "MULTI_TASK",
+ "job_clusters": self.job_clusters,
+ "max_concurrent_runs": self.max_concurrent_runs,
+ }
+ return merge(default_json, self.extra_job_params)
+
+ def _create_or_reset_job(self, context: Context) -> int:
+ job_spec = self.create_workflow_json(context=context)
+ existing_jobs = self._hook.list_jobs(job_name=self.databricks_job_name)
+ job_id = existing_jobs[0]["job_id"] if existing_jobs else None
+ if job_id:
+ self.log.info(
+ "Updating existing Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ self._hook.reset_job(job_id, job_spec)
+ else:
+ self.log.info(
+ "Creating new Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ job_id = self._hook.create_job(job_spec)
+ return job_id
+
+ def _wait_for_job_to_start(self, run_id: int):
+ run_url = self._hook.get_run_page_url(run_id)
+ self.log.info("Check the progress of the Databricks job at %s",
run_url)
+ life_cycle_state = self._hook.get_run_state(run_id).life_cycle_state
+ if life_cycle_state not in ("PENDING", "RUNNING", "BLOCKED"):
+ raise AirflowException(f"Could not start the workflow job. State:
{life_cycle_state}")
+ while life_cycle_state in ("PENDING", "BLOCKED"):
+ self.log.info("Waiting for the Databricks job to start running")
+ time.sleep(5)
+ life_cycle_state =
self._hook.get_run_state(run_id).life_cycle_state
+ self.log.info("Databricks job started. State: %s", life_cycle_state)
+
+ def execute(self, context: Context) -> Any:
+ if not isinstance(self.task_group, DatabricksWorkflowTaskGroup):
+ raise AirflowException("Task group must be a
DatabricksWorkflowTaskGroup")
+
+ job_id = self._create_or_reset_job(context)
+
+ run_id = self._hook.run_now(
+ {
+ "job_id": job_id,
+ "jar_params": self.task_group.jar_params,
+ "notebook_params": self.notebook_params,
+ "python_params": self.task_group.python_params,
+ "spark_submit_params": self.task_group.spark_submit_params,
+ }
+ )
+
+ self._wait_for_job_to_start(run_id)
+
+ return {
+ "conn_id": self.databricks_conn_id,
+ "job_id": job_id,
+ "run_id": run_id,
+ }
+
+
+class DatabricksWorkflowTaskGroup(TaskGroup):
+ """
+ A task group that takes a list of tasks and creates a databricks workflow.
+
+ The DatabricksWorkflowTaskGroup takes a list of tasks and creates a
databricks workflow
+ based on the metadata produced by those tasks. For a task to be eligible
for this
+ TaskGroup, it must contain the ``_convert_to_databricks_workflow_task``
method. If any tasks
+ do not contain this method then the Taskgroup will raise an error at parse
time.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DatabricksWorkflowTaskGroup`
+
+ :param databricks_conn_id: The name of the databricks connection to use.
+ :param existing_clusters: A list of existing clusters to use for this
workflow.
+ :param extra_job_params: A dictionary containing properties which will
override the default
+ Databricks Workflow Job definitions.
+ :param jar_params: A list of jar parameters to pass to the workflow. These
parameters will be passed to all jar
+ tasks in the workflow.
+ :param job_clusters: A list of job clusters to use for this workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for this
workflow.
+ :param notebook_packages: A list of dictionary of Python packages to be
installed. Packages defined
+ at the workflow task group level are installed for each of the
notebook tasks under it. And
+ packages defined at the notebook task level are installed specific for
the notebook task.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebook tasks in the workflow.
+ :param python_params: A list of python parameters to pass to the workflow.
These parameters will be passed to
+ all python tasks in the workflow.
+ :param spark_submit_params: A list of spark submit parameters to pass to
the workflow. These parameters
+ will be passed to all spark submit tasks.
+ """
+
+ is_databricks = True
+
+ def __init__(
+ self,
+ databricks_conn_id,
+ existing_clusters=None,
Review Comment:
```suggestion
existing_clusters: list[str] | None = None,
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
+ return f"{self.dag_id}.{self.task_group.group_id}"
+
+ def create_workflow_json(self, context: Context | None = None) ->
dict[str, object]:
+ """Create a workflow json to be used in the Databricks API."""
+ task_json = [
+ task._convert_to_databricks_workflow_task( # type:
ignore[attr-defined]
+ relevant_upstreams=self.relevant_upstreams, context=context
+ )
+ for task in self.tasks_to_convert
+ ]
+
+ default_json = {
+ "name": self.databricks_job_name,
+ "email_notifications": {"no_alert_for_skipped_runs": False},
+ "timeout_seconds": 0,
+ "tasks": task_json,
+ "format": "MULTI_TASK",
+ "job_clusters": self.job_clusters,
+ "max_concurrent_runs": self.max_concurrent_runs,
+ }
+ return merge(default_json, self.extra_job_params)
+
+ def _create_or_reset_job(self, context: Context) -> int:
+ job_spec = self.create_workflow_json(context=context)
+ existing_jobs = self._hook.list_jobs(job_name=self.databricks_job_name)
+ job_id = existing_jobs[0]["job_id"] if existing_jobs else None
+ if job_id:
+ self.log.info(
+ "Updating existing Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ self._hook.reset_job(job_id, job_spec)
+ else:
+ self.log.info(
+ "Creating new Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ job_id = self._hook.create_job(job_spec)
+ return job_id
+
+ def _wait_for_job_to_start(self, run_id: int):
+ run_url = self._hook.get_run_page_url(run_id)
+ self.log.info("Check the progress of the Databricks job at %s",
run_url)
+ life_cycle_state = self._hook.get_run_state(run_id).life_cycle_state
+ if life_cycle_state not in ("PENDING", "RUNNING", "BLOCKED"):
+ raise AirflowException(f"Could not start the workflow job. State:
{life_cycle_state}")
+ while life_cycle_state in ("PENDING", "BLOCKED"):
+ self.log.info("Waiting for the Databricks job to start running")
+ time.sleep(5)
+ life_cycle_state =
self._hook.get_run_state(run_id).life_cycle_state
+ self.log.info("Databricks job started. State: %s", life_cycle_state)
+
+ def execute(self, context: Context) -> Any:
+ if not isinstance(self.task_group, DatabricksWorkflowTaskGroup):
+ raise AirflowException("Task group must be a
DatabricksWorkflowTaskGroup")
+
+ job_id = self._create_or_reset_job(context)
+
+ run_id = self._hook.run_now(
+ {
+ "job_id": job_id,
+ "jar_params": self.task_group.jar_params,
+ "notebook_params": self.notebook_params,
+ "python_params": self.task_group.python_params,
+ "spark_submit_params": self.task_group.spark_submit_params,
+ }
+ )
+
+ self._wait_for_job_to_start(run_id)
+
+ return {
+ "conn_id": self.databricks_conn_id,
+ "job_id": job_id,
+ "run_id": run_id,
+ }
+
+
+class DatabricksWorkflowTaskGroup(TaskGroup):
+ """
+ A task group that takes a list of tasks and creates a databricks workflow.
+
+ The DatabricksWorkflowTaskGroup takes a list of tasks and creates a
databricks workflow
+ based on the metadata produced by those tasks. For a task to be eligible
for this
+ TaskGroup, it must contain the ``_convert_to_databricks_workflow_task``
method. If any tasks
+ do not contain this method then the Taskgroup will raise an error at parse
time.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DatabricksWorkflowTaskGroup`
+
+ :param databricks_conn_id: The name of the databricks connection to use.
+ :param existing_clusters: A list of existing clusters to use for this
workflow.
+ :param extra_job_params: A dictionary containing properties which will
override the default
+ Databricks Workflow Job definitions.
+ :param jar_params: A list of jar parameters to pass to the workflow. These
parameters will be passed to all jar
+ tasks in the workflow.
+ :param job_clusters: A list of job clusters to use for this workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for this
workflow.
+ :param notebook_packages: A list of dictionary of Python packages to be
installed. Packages defined
+ at the workflow task group level are installed for each of the
notebook tasks under it. And
+ packages defined at the notebook task level are installed specific for
the notebook task.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebook tasks in the workflow.
+ :param python_params: A list of python parameters to pass to the workflow.
These parameters will be passed to
+ all python tasks in the workflow.
+ :param spark_submit_params: A list of spark submit parameters to pass to
the workflow. These parameters
+ will be passed to all spark submit tasks.
+ """
+
+ is_databricks = True
+
+ def __init__(
+ self,
+ databricks_conn_id,
Review Comment:
```suggestion
databricks_conn_id: str,
```
##########
airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -0,0 +1,302 @@
+# 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 json
+import time
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from mergedeep import merge
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+ from airflow.models.taskmixin import DAGNode
+ from airflow.utils.context import Context
+
+
+@dataclass
+class WorkflowRunMetadata:
+ """
+ Metadata for a Databricks workflow run.
+
+ :param run_id: The ID of the Databricks workflow run.
+ :param job_id: The ID of the Databricks workflow job.
+ :param conn_id: The connection ID used to connect to Databricks.
+ """
+
+ conn_id: str
+ job_id: str
+ run_id: int
+
+
+def _flatten_node(
+ node: TaskGroup | BaseOperator | DAGNode, tasks: list[BaseOperator] | None
= None
+) -> list[BaseOperator]:
+ """Flatten a node (either a TaskGroup or Operator) to a list of nodes."""
+ if tasks is None:
+ tasks = []
+ if isinstance(node, BaseOperator):
+ return [node]
+
+ if isinstance(node, TaskGroup):
+ new_tasks = []
+ for _, child in node.children.items():
+ new_tasks += _flatten_node(child, tasks)
+
+ return tasks + new_tasks
+
+ return tasks
+
+
+class _CreateDatabricksWorkflowOperator(BaseOperator):
+ """
+ Creates a Databricks workflow from a DatabricksWorkflowTaskGroup specified
in a DAG.
+
+ :param task_id: The task_id of the operator
+ :param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param existing_clusters: A list of existing clusters to use for the
workflow.
+ :param extra_job_params: A dictionary of extra properties which will
override the default Databricks
+ Workflow Job definitions.
+ :param job_clusters: A list of job clusters to use for the workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for the
workflow.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebooks in the workflow.
+ :param tasks_to_convert: A list of tasks to convert to a Databricks
workflow. This list can also be
+ populated after instantiation using the `add_task` method.
+ """
+
+ template_fields = ("notebook_params",)
+ caller = "_CreateDatabricksWorkflowOperator"
+
+ def __init__(
+ self,
+ task_id,
+ databricks_conn_id,
+ existing_clusters: list[str] | None = None,
+ extra_job_params: dict[str, Any] | None = None,
+ job_clusters: list[dict[str, object]] | None = None,
+ max_concurrent_runs: int = 1,
+ notebook_params: dict | None = None,
+ tasks_to_convert: list[BaseOperator] | None = None,
+ **kwargs,
+ ):
+ self.databricks_conn_id = databricks_conn_id
+ self.existing_clusters = existing_clusters or []
+ self.extra_job_params = extra_job_params or {}
+ self.job_clusters = job_clusters or []
+ self.max_concurrent_runs = max_concurrent_runs
+ self.notebook_params = notebook_params or {}
+ self.tasks_to_convert = tasks_to_convert or []
+ self.relevant_upstreams = [task_id]
+ super().__init__(task_id=task_id, **kwargs)
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ caller=caller,
+ )
+
+ @cached_property
+ def _hook(self) -> DatabricksHook:
+ return self._get_hook(caller=self.caller)
+
+ def add_task(self, task: BaseOperator):
+ """Add a task to the list of tasks to convert to a Databricks
workflow."""
+ self.tasks_to_convert.append(task)
+
+ @property
+ def databricks_job_name(self):
+ return f"{self.dag_id}.{self.task_group.group_id}"
+
+ def create_workflow_json(self, context: Context | None = None) ->
dict[str, object]:
+ """Create a workflow json to be used in the Databricks API."""
+ task_json = [
+ task._convert_to_databricks_workflow_task( # type:
ignore[attr-defined]
+ relevant_upstreams=self.relevant_upstreams, context=context
+ )
+ for task in self.tasks_to_convert
+ ]
+
+ default_json = {
+ "name": self.databricks_job_name,
+ "email_notifications": {"no_alert_for_skipped_runs": False},
+ "timeout_seconds": 0,
+ "tasks": task_json,
+ "format": "MULTI_TASK",
+ "job_clusters": self.job_clusters,
+ "max_concurrent_runs": self.max_concurrent_runs,
+ }
+ return merge(default_json, self.extra_job_params)
+
+ def _create_or_reset_job(self, context: Context) -> int:
+ job_spec = self.create_workflow_json(context=context)
+ existing_jobs = self._hook.list_jobs(job_name=self.databricks_job_name)
+ job_id = existing_jobs[0]["job_id"] if existing_jobs else None
+ if job_id:
+ self.log.info(
+ "Updating existing Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ self._hook.reset_job(job_id, job_spec)
+ else:
+ self.log.info(
+ "Creating new Databricks workflow job %s with spec %s",
+ self.databricks_job_name,
+ json.dumps(job_spec, indent=2),
+ )
+ job_id = self._hook.create_job(job_spec)
+ return job_id
+
+ def _wait_for_job_to_start(self, run_id: int):
+ run_url = self._hook.get_run_page_url(run_id)
+ self.log.info("Check the progress of the Databricks job at %s",
run_url)
+ life_cycle_state = self._hook.get_run_state(run_id).life_cycle_state
+ if life_cycle_state not in ("PENDING", "RUNNING", "BLOCKED"):
+ raise AirflowException(f"Could not start the workflow job. State:
{life_cycle_state}")
+ while life_cycle_state in ("PENDING", "BLOCKED"):
+ self.log.info("Waiting for the Databricks job to start running")
+ time.sleep(5)
+ life_cycle_state =
self._hook.get_run_state(run_id).life_cycle_state
+ self.log.info("Databricks job started. State: %s", life_cycle_state)
+
+ def execute(self, context: Context) -> Any:
+ if not isinstance(self.task_group, DatabricksWorkflowTaskGroup):
+ raise AirflowException("Task group must be a
DatabricksWorkflowTaskGroup")
+
+ job_id = self._create_or_reset_job(context)
+
+ run_id = self._hook.run_now(
+ {
+ "job_id": job_id,
+ "jar_params": self.task_group.jar_params,
+ "notebook_params": self.notebook_params,
+ "python_params": self.task_group.python_params,
+ "spark_submit_params": self.task_group.spark_submit_params,
+ }
+ )
+
+ self._wait_for_job_to_start(run_id)
+
+ return {
+ "conn_id": self.databricks_conn_id,
+ "job_id": job_id,
+ "run_id": run_id,
+ }
+
+
+class DatabricksWorkflowTaskGroup(TaskGroup):
+ """
+ A task group that takes a list of tasks and creates a databricks workflow.
+
+ The DatabricksWorkflowTaskGroup takes a list of tasks and creates a
databricks workflow
+ based on the metadata produced by those tasks. For a task to be eligible
for this
+ TaskGroup, it must contain the ``_convert_to_databricks_workflow_task``
method. If any tasks
+ do not contain this method then the Taskgroup will raise an error at parse
time.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DatabricksWorkflowTaskGroup`
+
+ :param databricks_conn_id: The name of the databricks connection to use.
+ :param existing_clusters: A list of existing clusters to use for this
workflow.
+ :param extra_job_params: A dictionary containing properties which will
override the default
+ Databricks Workflow Job definitions.
+ :param jar_params: A list of jar parameters to pass to the workflow. These
parameters will be passed to all jar
+ tasks in the workflow.
+ :param job_clusters: A list of job clusters to use for this workflow.
+ :param max_concurrent_runs: The maximum number of concurrent runs for this
workflow.
+ :param notebook_packages: A list of dictionary of Python packages to be
installed. Packages defined
+ at the workflow task group level are installed for each of the
notebook tasks under it. And
+ packages defined at the notebook task level are installed specific for
the notebook task.
+ :param notebook_params: A dictionary of notebook parameters to pass to the
workflow. These parameters
+ will be passed to all notebook tasks in the workflow.
+ :param python_params: A list of python parameters to pass to the workflow.
These parameters will be passed to
+ all python tasks in the workflow.
+ :param spark_submit_params: A list of spark submit parameters to pass to
the workflow. These parameters
+ will be passed to all spark submit tasks.
+ """
+
+ is_databricks = True
+
+ def __init__(
+ self,
+ databricks_conn_id,
+ existing_clusters=None,
+ extra_job_params: dict[str, Any] | None = None,
+ jar_params=None,
Review Comment:
```suggestion
jar_params: list[str] | None = None,
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
