This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4082687dfc5 Add `access_control_list` to `DatabricksWorkflowTaskGroup`
(#64538)
4082687dfc5 is described below
commit 4082687dfc57a68c9321dc5be8c2289c76c14c17
Author: Jake McGrath <[email protected]>
AuthorDate: Wed Apr 29 17:34:43 2026 -0400
Add `access_control_list` to `DatabricksWorkflowTaskGroup` (#64538)
* issue-64527: Adding access_control_list to DatabricksWorkflowTaskGroup
* Delete uv.lock
* issue-64527: Re-adding uv.lock
* issue-64527: Restoring uv.lock
* issue-64527: Removing TODO
* issue-64527: Implementing Copilot suggestions
* Update
providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py
Co-authored-by: Copilot <[email protected]>
* issue-64527: Resolving Copilot suggestions
* issue-64527: Updating docstring, adding new test
---------
Co-authored-by: Jarek Potiuk <[email protected]>
Co-authored-by: Copilot <[email protected]>
---
.../databricks/operators/databricks_workflow.py | 27 ++++++++++++
.../operators/test_databricks_workflow.py | 51 ++++++++++++++++++++++
2 files changed, 78 insertions(+)
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py
b/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py
index f2f0b670aab..4eeb41babed 100644
---
a/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py
+++
b/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py
@@ -81,6 +81,15 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
:param task_id: The task_id of the operator
:param databricks_conn_id: The connection ID to use when connecting to
Databricks.
+ :param access_control_list: List of permissions to set on the job. Array
of object
+ (AccessControlRequestForUser) or object (AccessControlRequestForGroup)
or object
+ (AccessControlRequestForServicePrincipal).
+
+ .. seealso::
+ The access control list is applied both when creating a new job
and when resetting
+ an existing job. If access_control_list != [], the supplied ACL
replaces the
+ job's current permissions. Otherwise, the prior ACL stays in
place. You can also manage
+ job permissions directly in the Databricks UI.
: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.
@@ -100,6 +109,7 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
template_fields = (
"databricks_conn_id",
+ "access_control_list",
"existing_clusters",
"extra_job_params",
"jar_params",
@@ -125,6 +135,7 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
self,
task_id: str,
databricks_conn_id: str,
+ access_control_list: list[dict] | None = None,
existing_clusters: list[str] | None = None,
extra_job_params: dict[str, Any] | None = None,
jar_params: list[str] | None = None,
@@ -137,6 +148,7 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
**kwargs,
):
self.databricks_conn_id = databricks_conn_id
+ self.access_control_list = access_control_list
self.existing_clusters = existing_clusters or []
self.extra_job_params = extra_job_params or {}
self.jar_params = jar_params or []
@@ -188,6 +200,10 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
"job_clusters": self.job_clusters,
"max_concurrent_runs": self.max_concurrent_runs,
}
+
+ if self.access_control_list is not None:
+ default_json["access_control_list"] = self.access_control_list
+
return merge(default_json, self.extra_job_params)
def _create_or_reset_job(self, context: Context) -> int:
@@ -299,6 +315,14 @@ class DatabricksWorkflowTaskGroup(TaskGroup):
: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 access_control_list: List of permissions to set on the job. Array
of object
+ (AccessControlRequestForUser) or object (AccessControlRequestForGroup)
or object
+ (AccessControlRequestForServicePrincipal).
+
+ .. seealso::
+ When provided, this is applied when the job is created and when an
existing job is
+ reset. Setting this replaces any existing job permissions, unless
its value is
+ an empty list. If access_control_list == [], the prior ACL remains
in place.
: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
@@ -321,6 +345,7 @@ class DatabricksWorkflowTaskGroup(TaskGroup):
def __init__(
self,
databricks_conn_id: str,
+ access_control_list: list[dict] | None = None,
existing_clusters: list[str] | None = None,
extra_job_params: dict[str, Any] | None = None,
jar_params: list[str] | None = None,
@@ -333,6 +358,7 @@ class DatabricksWorkflowTaskGroup(TaskGroup):
**kwargs,
):
self.databricks_conn_id = databricks_conn_id
+ self.access_control_list = access_control_list
self.existing_clusters = existing_clusters or []
self.extra_job_params = extra_job_params or {}
self.jar_params = jar_params or []
@@ -356,6 +382,7 @@ class DatabricksWorkflowTaskGroup(TaskGroup):
task_group=self,
task_id="launch",
databricks_conn_id=self.databricks_conn_id,
+ access_control_list=self.access_control_list,
existing_clusters=self.existing_clusters,
extra_job_params=self.extra_job_params,
jar_params=self.jar_params,
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_databricks_workflow.py
b/providers/databricks/tests/unit/databricks/operators/test_databricks_workflow.py
index 43b5e47a81a..518bac7b04c 100644
---
a/providers/databricks/tests/unit/databricks/operators/test_databricks_workflow.py
+++
b/providers/databricks/tests/unit/databricks/operators/test_databricks_workflow.py
@@ -39,6 +39,13 @@ from airflow.utils import timezone
DEFAULT_DATE = timezone.datetime(2021, 1, 1)
+ACCESS_CONTROL_LIST = [
+ {
+ "user_name": "[email protected]",
+ "permission_level": "CAN_MANAGE",
+ }
+]
+
@pytest.fixture
def mock_databricks_hook():
@@ -98,6 +105,49 @@ def test_create_workflow_json(mock_databricks_hook,
context, mock_task_group):
assert workflow_json["max_concurrent_runs"] == 1
assert workflow_json["timeout_seconds"] == 0
+ assert "access_control_list" not in workflow_json
+
+
[email protected]("access_control_list", [ACCESS_CONTROL_LIST, []])
+def test_create_workflow_json_access_control_list(
+ mock_databricks_hook, context, mock_task_group, access_control_list
+):
+ """Test that _CreateDatabricksWorkflowOperator.create_workflow_json
includes access_control_list."""
+ operator = _CreateDatabricksWorkflowOperator(
+ task_id="test_task",
+ databricks_conn_id="databricks_default",
+ access_control_list=access_control_list,
+ )
+ operator.task_group = mock_task_group
+
+ task = MagicMock(spec=BaseOperator, task_id="task_1")
+ task._convert_to_databricks_workflow_task = MagicMock(return_value={})
+ operator.add_task(task.task_id, task)
+
+ workflow_json = operator.create_workflow_json(context=context)
+
+ # Only validate the access_control_list parameter; everything else has
been tested above
+ assert workflow_json["access_control_list"] == access_control_list
+
+
+def test_create_or_reset_job_empty_access_control_list(mock_databricks_hook,
context, mock_task_group):
+ """Test that access_control_list=[] reaches reset_job unchanged."""
+ operator = _CreateDatabricksWorkflowOperator(
+ task_id="test_task",
+ databricks_conn_id="databricks_default",
+ access_control_list=[],
+ )
+ operator.task_group = mock_task_group
+ operator._hook.list_jobs.return_value = [{"job_id": 123}]
+
+ operator._create_or_reset_job(context)
+
+ operator._hook.reset_job.assert_called_once()
+ _, job_spec = operator._hook.reset_job.call_args.args
+
+ assert "access_control_list" in job_spec
+ assert job_spec["access_control_list"] == []
+
def test_create_or_reset_job_existing(mock_databricks_hook, context,
mock_task_group):
"""Test that _CreateDatabricksWorkflowOperator._create_or_reset_job resets
the job if it already exists."""
@@ -214,6 +264,7 @@ def
test_task_group_exit_creates_operator(mock_databricks_workflow_operator):
task_group=task_group,
task_id="launch",
databricks_conn_id="databricks_conn",
+ access_control_list=None,
existing_clusters=[],
extra_job_params={},
jar_params=[],
