Lock in Databricks workflow depends_on parent-key behavior (#47614)

[Vamsi-klu]

The runtime fix for issue #47614 shipped in #48492 (commit 9dcce2f).
Closes: #47614
The GitHub issue stayed open because:

1. The existing unit test (test_convert_to_databricks_workflow_task) passed
   relevant_upstreams = [MagicMock(task_id="upstream_task")] — a list of
   mock objects, not strings — so the task_id in relevant_upstreams filter
   was silently always-False and the depends_on branch was never exercised.
   The fix could regress without anyone noticing.
2. A few small follow-on quality issues in the same area were worth tightening.
3. The issue was not linked to the merged PR, so it wasn't auto-closed.

This PR is therefore a lock-in / regression / hardening change, not a new fix.

Changes

• databricks.py: corrected relevant_upstreams annotation from
  list[BaseOperator] to list[str] on both DatabricksTaskBaseOperator
  and DatabricksNotebookOperator overrides.
• databricks_workflow.py: changed self.relevant_upstreams = [task_id]
  to self.relevant_upstreams: list[str] = []. Behavior is unchanged
  (the launch task's raw, unprefixed task_id was previously filtered out
  only by an accidental prefix mismatch); the new initializer makes the
  intent explicit.
• test_databricks_workflow.py: new TestWorkflowDependsOn class with 7
  end-to-end tests that build a real DAG + DatabricksWorkflowTaskGroup
  with real DatabricksNotebookOperator tasks and assert on the
  tasks[*]['depends_on'] payload returned by create_workflow_json().
  Covers default keys, custom parent key, >100-char parent key, diamond,
  fan-out, root tasks (launch task is filtered out), and external
  upstream filtering.
• test_databricks.py: fixed the existing
  test_convert_to_databricks_workflow_task to pass strings instead of
  mocks and to assert the correct parent-keyed depends_on. Added
  test_generate_databricks_task_key_requires_task_dict_when_task_id_passed.
• changelog.rst: Bug Fixes entry under 7.14.0.

closes: #47614

---

Was generative AI tooling used to co-author this PR?

• Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

[eladkal]

cc @moomindani for review

[moomindani]

The diff itself looks good to me — the type annotation is the right one, the relevant_upstreams = [task_id] → [] initializer is genuinely dead code being cleaned up (the bare "launch" could never match the child tasks' prefixed "<group>.launch" upstreams, so the original list was always semantically empty), the existing MagicMock-based test fix actually exercises the depends_on branch now, and TestWorkflowDependsOn covers default/custom/oversize key, diamond, fan-out, root-task filtering, and external-upstream filtering. Changelog revert in 025e8a4 is clean.

A couple of things before I'd be comfortable approving:

1. CI is red. The remaining failure (Compat 3.0.6:P3.10) is gh run download artifact infra, not your code, but we shouldn't merge on red — could you re-trigger it (or push an empty commit) so we get a clean build?

2. Real-workspace validation. The runtime fix has been live since #48492 in April 2025, so the user-visible behavior change here is small, but this PR does touch production code (relevant_upstreams = [task_id] → []). Have you (or anyone) actually run a DatabricksWorkflowTaskGroup with a parent/child task pair against a Databricks workspace and confirmed the rendered Jobs API JSON has depends_on: [{task_key: <parent_key>}]? The unit tests build real DAG objects but assert on the in-process create_workflow_json() payload, not what arrives at Databricks. An airflow dags test run + a screenshot of the resulting Databricks job graph would close the loop (similar pattern to what I used in #66613).

Once CI is green and there's some evidence the wire payload is right, LGTM.

[moomindani]

Real-workspace validation

Drove the PR branch end-to-end through airflow dags test against a Databricks workspace with a minimal DatabricksWorkflowTaskGroup (task_a >> task_b, both DatabricksNotebookOperator), then queried the created Databricks job back via 2.2/jobs/list. The Jobs API payload is correctly parent-keyed:

{
  "tasks": [
    {
      "task_key": "d13f300f3f3b28c910b3710198314589"
      // md5("pr66681_realenv__wf.task_a"); no depends_on
    },
    {
      "task_key": "dcf56d6d083da88df804ce506620a7a2",  // md5("pr66681_realenv__wf.task_b")
      "depends_on": [
        {"task_key": "d13f300f3f3b28c910b3710198314589"}  // parent task_a's key, not its own
      ]
    }
  ]
}

Closes the wire-payload loop on top of the unit tests.

Reproduction

export AIRFLOW_CONN_DATABRICKS_DEFAULT='{"conn_type":"databricks","host":"https://<workspace>","password":"<token>"}'
export PR66681_NOTEBOOK_PATH=/Users/<you>/airflow-pr66681-noop  # any existing notebook in the workspace
airflow dags test pr66681_realenv

Then verify the resulting Databricks job via:

curl -s -H "Authorization: Bearer $DATABRICKS_TOKEN" \
  "https://<workspace>/api/2.2/jobs/list?name=pr66681_realenv.wf&limit=1&expand_tasks=true" | jq '.jobs[0].settings.tasks'

dags/dag_pr66681_realenv.py

"""Real-environment validation DAG for PR #66681 (GH-47614)."""
from __future__ import annotations

import hashlib
import os
from datetime import datetime

from airflow.providers.databricks.hooks.databricks import DatabricksHook
from airflow.providers.databricks.operators.databricks import DatabricksNotebookOperator
from airflow.providers.databricks.operators.databricks_workflow import DatabricksWorkflowTaskGroup
from airflow.sdk import DAG, task

NOTEBOOK_PATH = os.environ["PR66681_NOTEBOOK_PATH"]
DAG_ID = "pr66681_realenv"
GROUP_ID = "wf"
CONN_ID = "databricks_default"


def _hook() -> DatabricksHook:
    return DatabricksHook(databricks_conn_id=CONN_ID)


def _expected_task_key(task_id_in_group: str) -> str:
    full_task_id = f"{GROUP_ID}.{task_id_in_group}"
    return hashlib.md5(f"{DAG_ID}__{full_task_id}".encode()).hexdigest()


with DAG(dag_id=DAG_ID, start_date=datetime(2026, 1, 1), schedule=None, catchup=False) as dag:
    with DatabricksWorkflowTaskGroup(
        group_id=GROUP_ID,
        databricks_conn_id=CONN_ID,
        job_clusters=[
            {
                "job_cluster_key": "shared",
                "new_cluster": {
                    "spark_version": "15.4.x-scala2.12",
                    "node_type_id": "i3.xlarge",
                    "num_workers": 0,
                    "spark_conf": {"spark.master": "local[*]"},
                    "custom_tags": {"ResourceClass": "SingleNode"},
                },
            }
        ],
    ) as wf:
        task_a = DatabricksNotebookOperator(
            task_id="task_a", notebook_path=NOTEBOOK_PATH, source="WORKSPACE", job_cluster_key="shared"
        )
        task_b = DatabricksNotebookOperator(
            task_id="task_b", notebook_path=NOTEBOOK_PATH, source="WORKSPACE", job_cluster_key="shared"
        )
        task_a >> task_b

    @task(task_id="verify_depends_on", trigger_rule="all_done")
    def verify_depends_on() -> None:
        jobs = _hook()._do_api_call(
            ("GET", "2.2/jobs/list"),
            {"name": f"{DAG_ID}.{GROUP_ID}", "limit": 1, "expand_tasks": True},
        )
        job = jobs["jobs"][0]
        tasks_by_key = {t["task_key"]: t for t in job["settings"]["tasks"]}
        a_key = _expected_task_key("task_a")
        b_key = _expected_task_key("task_b")
        assert tasks_by_key[a_key].get("depends_on", []) == []
        assert tasks_by_key[b_key]["depends_on"] == [{"task_key": a_key}]
        _hook()._do_api_call(("POST", "2.2/jobs/delete"), {"job_id": job["job_id"]})

    wf >> verify_depends_on()

Approving — note I'm not a committer, so this is only a non-binding LGTM; a committer still needs to sign off before merge.

[potiuk]

@Vamsi-klu A few things need addressing before review — see our Pull Request quality criteria.

• ❌ CI fails: MyPy providers checks (and possibly other checks — see the Checks tab for the full list).

No rush.

---

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

---

Drafted-by: Claude Code (Opus 4.7); reviewed by @potiuk before posting

[jscheffl]

Gentle bump on this PR. The requested validation has been added, CI is green, and @moomindani provided a non-binding approval with real-workspace validation. Could a committer please take another look when available?

Thanks.

Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting

There is one comment still open. Is this resolved? Then please resolve comment.

[jscheffl]

Code-wise looks good to me. If somebody with deep DBX know how around would be great.