refactor(github): address PR review comments for external_trigger and generic dispatch args

Author: RexBearIU

xlml/utils/github.py

@@ -14,6 +14,7 @@
 
 """Utilities for GitHub API integration."""
 
+from typing import Any
 import requests
 from airflow.decorators import task
 from airflow.exceptions import AirflowFailException
@@ -22,60 +23,80 @@
 
 @task
 def validate_git_trigger() -> None:
-  """Validates that the DAG run has required GitHub parameters.
-
-  Prevents manual runs from the Airflow UI by ensuring github_run_id,
-  github_repo, and github_token are present.
-  """
+  """Validates that the DAG run has required GitHub parameters and was externally triggered."""
   context = get_current_context()
+  dag_run = context.get("dag_run")
+
+  if not dag_run or not dag_run.external_trigger:
+    raise AirflowFailException(
+        "This DAG should not be run manually from the Airflow UI."
+    )
+
   params = context.get("params", {})
   run_id = params.get("github_run_id")
   repo = params.get("github_repo")
   token = params.get("github_token") or params.get("github_callback_token")
 
   if not run_id or not repo or not token:
     raise AirflowFailException(
-        "Missing required GitHub parameters (run_id, repo, token). "
-        "This DAG should not be run manually from the Airflow UI."
+        "Missing required GitHub parameters (run_id, repo, token)."
     )
 
 
 @task
-def fire_github_callback(test_type: str | None = None) -> None:
+def fire_github_callback(
+    test_type: str | None = None,
+    repo: str | None = None,
+    token: str | None = None,
+    event_type: str = "airflow-dag-complete",
+    client_payload: dict[str, Any] | None = None,
+) -> None:
   """Fires a GitHub repository_dispatch callback with the DAG run result."""
   context = get_current_context()
   params = context.get("params", {})
   dag_run = context.get("dag_run")
 
-  repo = params.get("github_repo")
-  token = params.get("github_token") or params.get("github_callback_token")
+  target_repo = repo or params.get("github_repo")
+  target_token = (
+      token or params.get("github_token") or params.get("github_callback_token")
+  )
 
-  if not repo or not token:
+  if not target_repo or not target_token:
     raise AirflowFailException(
         "Missing required GitHub parameters (repo, token) to fire callback."
     )
 
-  client_payload = {
-      "state": "success",
-      "dag_id": dag_run.dag_id if dag_run else None,
-      "dag_run_id": dag_run.run_id if dag_run else None,
-      "sha": params.get("commit_sha") or params.get("sha"),
-      "github_run_id": params.get("github_run_id"),
-  }
-  if test_type:
-    client_payload["test_type"] = test_type
+  if client_payload is None:
+    payload: dict[str, Any] = {
+        "state": "success",
+        "dag_id": dag_run.dag_id if dag_run else None,
+        "dag_run_id": dag_run.run_id if dag_run else None,
+        "sha": params.get("commit_sha")
+        or params.get("maxtext_sha")
+        or params.get("sha"),
+        "github_run_id": params.get("github_run_id"),
+    }
+    if test_type:
+      payload["test_type"] = test_type
+  else:
+    payload = dict(client_payload)
+    if test_type and "test_type" not in payload:
+      payload["test_type"] = test_type
 
   response = requests.post(
-      f"https://api.github.com/repos/{repo}/dispatches",
+      f"https://api.github.com/repos/{target_repo}/dispatches",
       headers={
-          "Authorization": f"Bearer {token}",
+          "Authorization": f"Bearer {target_token}",
           "Accept": "application/vnd.github+json",
           "X-GitHub-Api-Version": "2022-11-28",
       },
       json={
-          "event_type": "airflow-dag-complete",
-          "client_payload": client_payload,
+          "event_type": event_type,
+          "client_payload": payload,
       },
       timeout=30,
   )
   response.raise_for_status()
+
+
+trigger_github_repository_dispatch = fire_github_callback

xlml/utils/github_test.py

@@ -24,23 +24,44 @@ class GithubTest(unittest.TestCase):
 
   @mock.patch("xlml.utils.github.get_current_context")
   def test_validate_git_trigger_success(self, mock_context):
+    mock_dag_run = mock.MagicMock()
+    mock_dag_run.external_trigger = True
     mock_context.return_value = {
+        "dag_run": mock_dag_run,
         "params": {
             "github_run_id": "12345",
             "github_repo": "owner/repo",
             "github_token": "secret_token",
-        }
+        },
     }
     # Should not raise exception
     github.validate_git_trigger.function()
 
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_validate_git_trigger_manual_run(self, mock_context):
+    mock_dag_run = mock.MagicMock()
+    mock_dag_run.external_trigger = False
+    mock_context.return_value = {
+        "dag_run": mock_dag_run,
+        "params": {
+            "github_run_id": "12345",
+            "github_repo": "owner/repo",
+            "github_token": "secret_token",
+        },
+    }
+    with self.assertRaises(AirflowFailException):
+      github.validate_git_trigger.function()
+
   @mock.patch("xlml.utils.github.get_current_context")
   def test_validate_git_trigger_missing_params(self, mock_context):
+    mock_dag_run = mock.MagicMock()
+    mock_dag_run.external_trigger = True
     mock_context.return_value = {
+        "dag_run": mock_dag_run,
         "params": {
             "github_run_id": "12345",
             # missing repo and token
-        }
+        },
     }
     with self.assertRaises(AirflowFailException):
       github.validate_git_trigger.function()
@@ -89,6 +110,37 @@ def test_fire_github_callback_success(self, mock_context, mock_post):
     )
     mock_response.raise_for_status.assert_called_once()
 
+  @mock.patch("requests.post")
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_trigger_github_repository_dispatch_explicit_args(
+      self, mock_context, mock_post
+  ):
+    mock_context.return_value = {}
+    mock_response = mock.MagicMock()
+    mock_post.return_value = mock_response
+
+    github.trigger_github_repository_dispatch.function(
+        repo="custom/repo",
+        token="custom_token",
+        event_type="custom-event",
+        client_payload={"foo": "bar"},
+    )
+
+    mock_post.assert_called_once_with(
+        "https://api.github.com/repos/custom/repo/dispatches",
+        headers={
+            "Authorization": "Bearer custom_token",
+            "Accept": "application/vnd.github+json",
+            "X-GitHub-Api-Version": "2022-11-28",
+        },
+        json={
+            "event_type": "custom-event",
+            "client_payload": {"foo": "bar"},
+        },
+        timeout=30,
+    )
+    mock_response.raise_for_status.assert_called_once()
+
   @mock.patch("xlml.utils.github.get_current_context")
   def test_fire_github_callback_missing_token(self, mock_context):
     mock_context.return_value = {