Databricks Operator set it owns task_key in depends_on instead of parent task key

[pacmora]

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

[boring-cyborg]

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]

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.

#create_databricks_workflow_task.relevant_upstreams.append(task.task_id)
create_databricks_workflow_task.relevant_upstreams.append(task)

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:

"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
  ],

Fixed code:

"depends_on": [
  {"task_key": upstream_task.databricks_task_key}
    for upstream_task in relevant_upstreams
    if isinstance(upstream_task, DatabricksTaskBaseOperator) and upstream_task.task_id in self.upstream_task_ids
  ],

KR,
Pedro.

[potiuk]

PR to fix it is the way to go @pacmora

[pacmora]

PR to fix it is the way to go @pacmora

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!!!!