josh-fell commented on code in PR #32221:
URL: https://github.com/apache/airflow/pull/32221#discussion_r1311753571
##########
airflow/providers/databricks/operators/databricks.py:
##########
@@ -161,6 +163,139 @@ def get_link(
return XCom.get_value(key=XCOM_RUN_PAGE_URL_KEY, ti_key=ti_key)
+class DatabricksCreateJobsOperator(BaseOperator):
+ """
+ Creates (or resets) a Databricks job using the
+ `api/2.1/jobs/create
+
<https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsCreate>`_
+ (or `api/2.1/jobs/reset
+
<https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsReset>`_)
+ API endpoint.
+
+ .. seealso::
+
https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsCreate
+
+ :param json: A JSON object containing API parameters which will be passed
+ directly to the ``api/2.1/jobs/create`` endpoint. The other named
parameters
+ (i.e. ``name``, ``tags``, ``tasks``, etc.) to this operator will
+ be merged with this json dictionary if they are provided.
+ If there are conflicts during the merge, the named parameters will
+ take precedence and override the top level json keys. (templated)
+
+ .. seealso::
+ For more information about templating see
:ref:`concepts:jinja-templating`.
+ :param name: An optional name for the job.
+ :param tags: A map of tags associated with the job.
+ :param tasks: A list of task specifications to be executed by this job.
+ Array of objects (JobTaskSettings).
+ :param job_clusters: A list of job cluster specifications that can be
shared and reused by
+ tasks of this job. Array of objects (JobCluster).
+ :param email_notifications: Object (JobEmailNotifications).
+ :param webhook_notifications: Object (WebhookNotifications).
+ :param timeout_seconds: An optional timeout applied to each run of this
job.
+ :param schedule: Object (CronSchedule).
+ :param max_concurrent_runs: An optional maximum allowed number of
concurrent runs of the job.
+ :param git_source: An optional specification for a remote repository
containing the notebooks
+ used by this job's notebook tasks. Object (GitSource).
+ :param access_control_list: List of permissions to set on the job. Array
of object
+ (AccessControlRequestForUser) or object (AccessControlRequestForGroup)
or object
+ (AccessControlRequestForServicePrincipal).
+
+ .. seealso::
+ This will only be used on create. In order to reset ACL consider
using the Databricks
+ UI.
+ :param databricks_conn_id: Reference to the
+ :ref:`Databricks connection <howto/connection:databricks>`. (templated)
+ :param polling_period_seconds: Controls the rate which we poll for the
result of
+ this run. By default the operator will poll every 30 seconds.
+ :param databricks_retry_limit: Amount of times retry if the Databricks
backend is
+ unreachable. Its value must be greater than or equal to 1.
+ :param databricks_retry_delay: Number of seconds to wait between retries
(it
+ might be a floating point number).
+ :param databricks_retry_args: An optional dictionary with arguments passed
to ``tenacity.Retrying`` class.
+ """
+
+ # Used in airflow.models.BaseOperator
+ template_fields: Sequence[str] = ("json", "databricks_conn_id")
+ # Databricks brand color (blue) under white text
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ *,
+ json: dict | None = None,
+ name: str | None = None,
+ tags: dict[str, str] | None = None,
+ tasks: list[jobs.JobTaskSettings] | None = None,
+ job_clusters: list[jobs.JobCluster] | None = None,
+ email_notifications: jobs.JobEmailNotifications | None = None,
+ webhook_notifications: jobs.JobWebhookNotifications | None = None,
+ timeout_seconds: int | None = None,
+ schedule: jobs.CronSchedule | None = None,
+ max_concurrent_runs: int | None = None,
+ git_source: jobs.GitSource | None = None,
+ access_control_list: list[jobs.AccessControlRequest] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ **kwargs,
+ ) -> None:
+ """Creates a new ``DatabricksCreateJobsOperator``."""
+ super().__init__(**kwargs)
+ self.json = json or {}
+ self.databricks_conn_id = databricks_conn_id
+ self.polling_period_seconds = polling_period_seconds
+ self.databricks_retry_limit = databricks_retry_limit
+ self.databricks_retry_delay = databricks_retry_delay
+ self.databricks_retry_args = databricks_retry_args
+ if name is not None:
+ self.json["name"] = name
+ if tags is not None:
+ self.json["tags"] = tags
Review Comment:
To illustrate the point, let's use this example DAG where we define the
`json` arg in a previous task and use its output:
```py
from __future__ import annotations
from pendulum import datetime
from typing import TYPE_CHECKING, Sequence
from airflow.decorators import dag, task
from airflow.models.baseoperator import BaseOperator
if TYPE_CHECKING:
from airflow.utils.context import Context
class DatabricksCreateJobsOperator(BaseOperator):
template_fields: Sequence[str] = ("json", "databricks_conn_id")
def __init__(
self,
*,
json: dict | None = None,
name: str | None = None,
tags: dict[str, str] | None = None,
**kwargs
):
super().__init__(**kwargs)
self.json = json or {}
if name is not None:
self.json["name"] = name
if tags is not None:
self.json["tags"] = tags
def execute(context: Context) -> None:
pass
@dag(start_date=datetime(2023, 1, 1), schedule=None)
def derived_template_fields():
@task
def push_json() -> dict[str, str]:
return {"key1": "val1", "key2": "val2"}
json = push_json()
DatabricksCreateJobsOperator(
task_id="create_job_w_json", json=json, name="some_name",
tags={"key3": "value3"}
)
derived_template_fields()
```
DAG parsing fails with:
```
Running: airflow dags reserialize
[2023-08-31T14:29:57.796+0000] {utils.py:430} WARNING - No module named
'paramiko'
[2023-08-31T14:29:57.816+0000] {utils.py:430} WARNING - No module named
'airflow.providers.dbt'
[2023-08-31T14:29:58.533+0000] {dagbag.py:539} INFO - Filling up the DagBag
from /usr/local/airflow/dags
[2023-08-31T14:29:58.615+0000] {dagbag.py:347} ERROR - Failed to import:
/usr/local/airflow/dags/derived_template_fields.py
Traceback (most recent call last):
File "/usr/local/lib/python3.11/site-packages/airflow/models/dagbag.py",
line 343, in parse
loader.exec_module(new_module)
File "<frozen importlib._bootstrap_external>", line 940, in exec_module
File "<frozen importlib._bootstrap>", line 241, in
_call_with_frames_removed
File "/usr/local/airflow/dags/derived_template_fields.py", line 50, in
<module>
derived_template_fields()
File "/usr/local/lib/python3.11/site-packages/airflow/models/dag.py", line
3798, in factory
f(**f_kwargs)
File "/usr/local/airflow/dags/derived_template_fields.py", line 41, in
derived_template_fields
DatabricksCreateJobsOperator(
File
"/usr/local/lib/python3.11/site-packages/airflow/models/baseoperator.py", line
436, in apply_defaults
result = func(self, **kwargs, default_args=default_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/airflow/dags/derived_template_fields.py", line 27, in
__init__
self.json["name"] = name
~~~~~~~~~^^^^^^^^
TypeError: 'PlainXComArg' object does not support item assignment
```
Even if we change the `json` arg assignment to use the classic XCom Jinja
template approach (i.e. `json = "{{ ti.xcom_pull(task_ids='push_json') }}"`),
the DAG fails to parse:
```
Running: airflow dags reserialize
[2023-08-31T14:32:01.553+0000] {utils.py:430} WARNING - No module named
'paramiko'
[2023-08-31T14:32:01.574+0000] {utils.py:430} WARNING - No module named
'airflow.providers.dbt'
[2023-08-31T14:32:02.341+0000] {dagbag.py:539} INFO - Filling up the DagBag
from /usr/local/airflow/dags
[2023-08-31T14:32:02.415+0000] {dagbag.py:347} ERROR - Failed to import:
/usr/local/airflow/dags/derived_template_fields.py
Traceback (most recent call last):
File "/usr/local/lib/python3.11/site-packages/airflow/models/dagbag.py",
line 343, in parse
loader.exec_module(new_module)
File "<frozen importlib._bootstrap_external>", line 940, in exec_module
File "<frozen importlib._bootstrap>", line 241, in
_call_with_frames_removed
File "/usr/local/airflow/dags/derived_template_fields.py", line 51, in
<module>
derived_template_fields()
File "/usr/local/lib/python3.11/site-packages/airflow/models/dag.py", line
3798, in factory
f(**f_kwargs)
File "/usr/local/airflow/dags/derived_template_fields.py", line 42, in
derived_template_fields
DatabricksCreateJobsOperator(
File
"/usr/local/lib/python3.11/site-packages/airflow/models/baseoperator.py", line
436, in apply_defaults
result = func(self, **kwargs, default_args=default_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/airflow/dags/derived_template_fields.py", line 27, in
__init__
self.json["name"] = name
~~~~~~~~~^^^^^^^^
TypeError: 'str' object does not support item assignment
```
Perhaps it's possible users haven't needed a use case for predefining a
`json` arg from a previous task, Airflow Variable, DAG Param, etc. (accessed by
a Jinja templates).
Seems like there is now some movement on addressing #29069 to prevent this
in the future too.
--
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]
