Monitor the workflow run breeze actually dispatched (#71305)

* Monitor the workflow run breeze actually dispatched

The docs-publish driver looked up the newest run of the workflow right
after dispatching it. When the dispatch had not registered yet, that is
the previous run - typically already finished - so the driver reported
success and moved straight on to the airflow-site refresh and the
S3-to-GitHub sync while the docs build it was supposed to gate on was
still running. A scheduled run, or a run someone else started, was picked
up the same way.

Terminal conclusions other than success and failure were also treated as
nothing to report, so a cancelled or timed-out run let the rest of the
chain proceed as if the docs had been published.

* Keep breeze workflow-run monitoring on the dispatched run only

Several of the workflows breeze dispatches - apache/airflow-site's build.yml
among them - also run on push and pull request. Such a run registering while
we poll would be taken for the one we dispatched, so only workflow_dispatch
runs are considered now.

A transient gh failure during that poll also aborted the whole publishing
chain, even though the caller is already retrying.

Author: potiuk

dev/breeze/src/airflow_breeze/utils/gh_workflow_utils.py

@@ -28,6 +28,9 @@
 from airflow_breeze.utils.github import run_gh_command
 from airflow_breeze.utils.shared_options import get_dry_run
 
+NEW_RUN_TIMEOUT_SECONDS = 180
+NEW_RUN_POLL_SECONDS = 5
+
 
 def tigger_workflow(workflow_name: str, repo: str, branch: str = "main", **kwargs):
     """
@@ -93,12 +96,20 @@ def make_sure_gh_is_installed():
         sys.exit(1)
 
 
-def get_workflow_run_id(workflow_name: str, repo: str) -> int:
+def get_latest_workflow_run_id(workflow_name: str, repo: str, *, exit_on_error: bool = True) -> int | None:
     """
-    Get the latest workflow run ID for a given workflow name and repository.
+    Get the latest dispatched workflow run ID for a given workflow name and repository.
+
+    Only ``workflow_dispatch`` runs are considered - several of the workflows we drive here also run
+    on push or pull request (apache/airflow-site's `build.yml`, for one), and such a run registering
+    between our two lookups would otherwise be mistaken for the one we dispatched.
 
     :param workflow_name: The name of the workflow to check.
     :param repo: The repository in the format 'owner/repo'.
+    :param exit_on_error: Whether a failing `gh` call should terminate breeze. Pass False when the
+        caller polls and can afford to retry a transient failure.
+    :return: The run id, or None when the workflow has never been dispatched (or when the lookup
+        failed and ``exit_on_error`` is False).
     """
     make_sure_gh_is_installed()
     command = [
@@ -109,6 +120,8 @@ def get_workflow_run_id(workflow_name: str, repo: str) -> int:
         workflow_name,
         "--repo",
         repo,
+        "--event",
+        "workflow_dispatch",
         "--limit",
         "1",
         "--json",
@@ -117,21 +130,49 @@ def get_workflow_run_id(workflow_name: str, repo: str) -> int:
 
     result = run_gh_command(command, capture_output=True)
     if result.returncode != 0:
+        if not exit_on_error:
+            console_print(f"[yellow]Error fetching workflow run ID: {result.stderr} - retrying.[/yellow]")
+            return None
         console_print(f"[red]Error fetching workflow run ID: {result.stderr}[/red]")
         sys.exit(1)
 
     runs_data = result.stdout.strip()
     if not runs_data:
-        console_print("[red]No workflow runs found.[/red]")
-        sys.exit(1)
+        return None
 
-    run_id = json.loads(runs_data)[0].get("databaseId")
+    runs = json.loads(runs_data)
+    return runs[0].get("databaseId") if runs else None
 
-    console_print(
-        f"[blue]Running workflow {workflow_name} at https://github.com/{repo}/actions/runs/{run_id}[/blue]",
-    )
 
-    return run_id
+def wait_for_new_workflow_run(workflow_name: str, repo: str, previous_run_id: int | None) -> int:
+    """
+    Wait until a run newer than ``previous_run_id`` shows up and return its id.
+
+    Run ids increase monotonically, so anything above the id observed just before the dispatch is
+    the run we started. Taking whatever run is newest would instead latch onto an unrelated one -
+    a scheduled run, or another maintainer's - whenever ours has not registered yet, and report
+    that run's result as ours.
+
+    :param workflow_name: The name of the workflow that was dispatched.
+    :param repo: The repository in the format 'owner/repo'.
+    :param previous_run_id: The newest run id seen before dispatching, or None if there was none.
+    """
+    deadline = time.monotonic() + NEW_RUN_TIMEOUT_SECONDS
+    while True:
+        run_id = get_latest_workflow_run_id(workflow_name, repo, exit_on_error=False)
+        if run_id is not None and (previous_run_id is None or run_id > previous_run_id):
+            console_print(
+                f"[blue]Running workflow {workflow_name} at "
+                f"https://github.com/{repo}/actions/runs/{run_id}[/blue]",
+            )
+            return run_id
+        if time.monotonic() >= deadline:
+            console_print(
+                f"[red]Timed out after {NEW_RUN_TIMEOUT_SECONDS}s waiting for the dispatched run of "
+                f"{workflow_name} in {repo} to appear.[/red]"
+            )
+            sys.exit(1)
+        time.sleep(NEW_RUN_POLL_SECONDS)
 
 
 def get_workflow_run_info(run_id: str, repo: str, fields: str) -> dict:
@@ -188,12 +229,14 @@ def monitor_workflow_run(run_id: str, repo: str):
         if status == "completed":
             if conclusion == "success":
                 console_print(f"[green]Workflow {name} run {run_id} completed successfully.[/green]")
-            elif conclusion == "failure":
-                console_print(
-                    f"[red]Workflow {name} run {run_id} failed, see for more info: https://github.com/{repo}/actions/runs/{run_id}[/red]"
-                )
-                sys.exit(1)
-            break
+                break
+            # Anything else - failure, cancelled, timed_out, action_required - means the run did not
+            # produce what the caller is about to chain further work onto.
+            console_print(
+                f"[red]Workflow {name} run {run_id} finished with conclusion '{conclusion}', "
+                f"see for more info: https://github.com/{repo}/actions/runs/{run_id}[/red]"
+            )
+            sys.exit(1)
 
         # Check status of jobs every 30 seconds
         time.sleep(30)
@@ -203,6 +246,7 @@ def trigger_workflow_and_monitor(
     workflow_name: str, repo: str, branch: str = "main", monitor=True, **workflow_fields
 ):
     make_sure_gh_is_installed()
+    previous_run_id = None if get_dry_run() else get_latest_workflow_run_id(workflow_name, repo)
     tigger_workflow(
         workflow_name=workflow_name,
         repo=repo,
@@ -213,9 +257,10 @@ def trigger_workflow_and_monitor(
     if get_dry_run():
         return
 
-    workflow_run_id = get_workflow_run_id(
+    workflow_run_id = wait_for_new_workflow_run(
         workflow_name=workflow_name,
         repo=repo,
+        previous_run_id=previous_run_id,
     )
 
     console_print(

dev/breeze/tests/test_gh_workflow_utils.py

@@ -16,9 +16,19 @@
 # under the License.
 from __future__ import annotations
 
+import contextlib
+import subprocess
 from unittest import mock
 
-from airflow_breeze.utils.gh_workflow_utils import trigger_workflow_and_monitor
+import pytest
+
+from airflow_breeze.utils.gh_workflow_utils import (
+    NEW_RUN_TIMEOUT_SECONDS,
+    get_latest_workflow_run_id,
+    monitor_workflow_run,
+    trigger_workflow_and_monitor,
+    wait_for_new_workflow_run,
+)
 from airflow_breeze.utils.shared_options import set_dry_run
 
 
@@ -38,3 +48,87 @@ def test_trigger_workflow_and_monitor_stops_after_the_dispatch_in_dry_run(_, moc
         set_dry_run(False)
 
     mock_monitor.assert_not_called()
+
+
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.run_gh_command")
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.make_sure_gh_is_installed")
+def test_get_latest_workflow_run_id_only_looks_at_dispatched_runs(_, mock_run_gh_command):
+    """Push and pull_request runs of the same workflow must not be mistaken for the dispatched one."""
+    mock_run_gh_command.return_value = subprocess.CompletedProcess(
+        args=[], returncode=0, stdout='[{"databaseId": 123}]', stderr=""
+    )
+
+    assert get_latest_workflow_run_id("build.yml", "apache/airflow-site") == 123
+
+    command = mock_run_gh_command.call_args.args[0]
+    assert command[command.index("--event") + 1] == "workflow_dispatch"
+
+
+@pytest.mark.parametrize(
+    ("exit_on_error", "expectation"),
+    [(True, pytest.raises(SystemExit)), (False, contextlib.nullcontext())],
+)
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.run_gh_command")
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.make_sure_gh_is_installed")
+def test_get_latest_workflow_run_id_lets_a_polling_caller_survive_a_failed_lookup(
+    _, mock_run_gh_command, exit_on_error, expectation
+):
+    mock_run_gh_command.return_value = subprocess.CompletedProcess(
+        args=[], returncode=1, stdout="", stderr="could not connect to api.github.com"
+    )
+
+    with expectation:
+        assert get_latest_workflow_run_id("build.yml", "apache/airflow", exit_on_error=exit_on_error) is None
+
+
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.time.sleep")
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.get_latest_workflow_run_id")
+def test_wait_for_new_workflow_run_retries_after_a_failed_lookup(mock_latest, _):
+    """A transient `gh` failure while polling must not abort the whole publishing chain."""
+    mock_latest.side_effect = [None, 222]
+
+    assert wait_for_new_workflow_run("build.yml", "apache/airflow-site", previous_run_id=111) == 222
+    assert all(call.kwargs["exit_on_error"] is False for call in mock_latest.call_args_list)
+
+
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.time.sleep")
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.get_latest_workflow_run_id")
+def test_wait_for_new_workflow_run_ignores_the_run_that_predates_the_dispatch(mock_latest, _):
+    mock_latest.side_effect = [111, 111, 222]
+
+    assert wait_for_new_workflow_run("publish-docs-to-s3.yml", "apache/airflow", previous_run_id=111) == 222
+
+
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.time.sleep")
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.get_latest_workflow_run_id")
+def test_wait_for_new_workflow_run_accepts_the_first_ever_run(mock_latest, _):
+    mock_latest.return_value = 42
+
+    assert wait_for_new_workflow_run("build.yml", "apache/airflow-site", previous_run_id=None) == 42
+
+
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.time.monotonic")
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.time.sleep")
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.get_latest_workflow_run_id")
+def test_wait_for_new_workflow_run_gives_up_when_no_new_run_appears(mock_latest, _, mock_monotonic):
+    mock_latest.return_value = 111
+    mock_monotonic.side_effect = [0, 0, NEW_RUN_TIMEOUT_SECONDS]
+
+    with pytest.raises(SystemExit) as exc_info:
+        wait_for_new_workflow_run("publish-docs-to-s3.yml", "apache/airflow", previous_run_id=111)
+
+    assert exc_info.value.code == 1
+
+
+@pytest.mark.parametrize("conclusion", ["failure", "cancelled", "timed_out", "action_required"])
+@mock.patch("airflow_breeze.utils.gh_workflow_utils.get_workflow_run_info")
+def test_monitor_workflow_run_fails_on_any_unsuccessful_conclusion(mock_info, conclusion):
+    mock_info.side_effect = [
+        {"jobs": []},
+        {"status": "completed", "conclusion": conclusion, "name": "Publish Docs to S3"},
+    ]
+
+    with pytest.raises(SystemExit) as exc_info:
+        monitor_workflow_run(run_id="123", repo="apache/airflow")
+
+    assert exc_info.value.code == 1