Repository navigation
Databricks Operator set it owns task_key in depends_on instead of parent task key #47614
Description
Apache Airflow Provider(s)
databricks
Versions of Apache Airflow Providers
apache-airflow-providers-databricks==7.2.0
Apache Airflow version
2.10.5
Operating System
Debian Bookworm
Deployment
Official Apache Airflow Helm Chart
Deployment details
No response
What happened
When DatabricksWorkflowTaskGroup contains at least one task that has a child task, launch task generates wrong json payload. Tasks with parents are setting it owns task_key in depends_on field.
What you think should happen instead
Tasks with parent/s should set parent task_key instead of it owns task_key
How to reproduce
Create DatabricksWorkflowTaskGroup. Then add a task_A >> task_B,
Task_B will add it's owns task_key in "depends_on" field instead of task_A task_key
Anything else
I think the issue is located in file airflow/providers/databricks/operators/databricks.py.py
The first time _generated_databricks_task_key is executed, even if task_id is provided or not, self._databricks_task_key will be set. It means no matter how many times "_generate_databricks_task_key", even passing differents task_id param, is called, it always will return the same value, the one returned in the first call.
def _generate_databricks_task_key(self, task_id: str | None = None) -> str:
"""Create a databricks task key using the hash of dag_id and task_id."""
if not self._databricks_task_key or len(self._databricks_task_key) > 100:
self.log.info(
"databricks_task_key has not be provided or the provided one exceeds 100 characters and will be truncated by the Databricks API. This will cause failure when trying to monitor the task. A task_key will be generated using the hash value of dag_id+task_id"
)
task_id = task_id or self.task_id
task_key = f"{self.dag_id}__{task_id}".encode()
self._databricks_task_key = hashlib.md5(task_key).hexdigest()
self.log.info("Generated databricks task_key: %s", self._databricks_task_key)
return self._databricks_task_key
If we check block of code that converts a task into databricks_task, is setting "task_key" with self.databricks_task_key. At this point if _generate_databricks_task_key is called again, with or without task_id param, it will return always the same value due self._databricks_task_key is not "None" anymore.
def _convert_to_databricks_workflow_task(
self, relevant_upstreams: list[BaseOperator], context: Context | None = None
) -> dict[str, object]:
"""Convert the operator to a Databricks workflow task that can be a task in a workflow."""
base_task_json = self._get_task_base_json()
result = {
"task_key": self.databricks_task_key,
"depends_on": [
{"task_key": self._generate_databricks_task_key(task_id)}
for task_id in self.upstream_task_ids
if task_id in relevant_upstreams
],
**base_task_json,
}
if self.existing_cluster_id and self.job_cluster_key:
raise ValueError(
"Both existing_cluster_id and job_cluster_key are set. Only one can be set per task."
)
if self.existing_cluster_id:
result["existing_cluster_id"] = self.existing_cluster_id
elif self.job_cluster_key:
result["job_cluster_key"] = self.job_cluster_key
return result
Are you willing to submit PR?
- Yes I am willing to submit a PR!
Code of Conduct
- I agree to follow this project's Code of Conduct
Activity
boring-cyborg commentedon Mar 11, 2025
Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.
pacmora commentedon Mar 25, 2025
Hi,
I tested a fix and it works for me.
I checked DatabricksTaskBaseOperator, the method "_convert_to_databricks_workflow_task" should received a list of BaseOperator but it receive a list of string that contains task_id.
Method "_generate_databricks_task_key" cannot generate the task_id for parent task, it was designed to generate the task_id for the object that execute this method.
So I made this modifications:
In file operators/databricks_workflow.py:
Modify class _CreateDatabricksWorkflowOperator changing list type for relevant_upstreams in init methid, BaseOperator instead of string.
self.relevant_upstreams: list[BaseOperator] = [self]Then modify class DatabricksWorkflowTaskGroup, in the exit method, instead of append task_id to relevant_upstreams, append the task itself.
In file operators/databricks.py:
Class DatabricksTaskBaseOperator, _convert_to_databricks_workflow_task is receiving now a list of BaseOperator instead of list of task_id's. When depends_on array is generate, instead of set task_key value with the return
of self._generate_databricks_task_key(task_id), must be set using the BaseOperator object from relevant_upstreams casted as DatabricksTaskBaseOperator.
Previous code:
Fixed code:
KR,
Pedro.
potiuk commentedon Mar 30, 2025
PR to fix it is the way to go @pacmora
pacmora commentedon Apr 1, 2025
Yes, sorry I had some issue with my laptop.
I just created the PR
It's my first PR, if there is something wrong so if there is something wrong sorry for the inconveniences.
Thanks!!!!
potiuk commentedon Mar 2, 2026
@pacmora We are unassigning you from this issue as part of our updated assignment policy.
This is not meant to discourage your contribution — quite the opposite! You are still very welcome to work on this issue and submit a PR for it. Simply comment that you are working on it and open a PR when ready.
We found that formal assignments were not working well, as they often prevented others from contributing when the assignee was not actively working on the issue.
Lock in Databricks workflow depends_on parent-key behavior (#47614) (#…