feat: update to support airflow 3

.github/workflows/pr.yml

@@ -115,6 +115,9 @@ jobs:
           STAGING_SM2A_ADMIN_PASSWORD: ${{ secrets.STAGING_SM2A_ADMIN_PASSWORD }}
           STAGING_SM2A_API_URL: ${{ vars.STAGING_SM2A_API_URL }}
           DATASET_DAG_NAME: ${{ vars.DATASET_DAG_NAME }}
+          STAGING_AIRFLOW_API_VERSION: ${{ vars.STAGING_AIRFLOW_API_VERSION }}
+          AIRFLOW_JWT_SECRET: ${{ secrets.AIRFLOW_JWT_SECRET }}
+          AIRFLOW_JWT_SUB: ${{ secrets.AIRFLOW_JWT_SUB }}
         run: |
           pip install -r ./scripts/requirements.txt
 

.github/workflows/promote.yml

@@ -64,7 +64,10 @@ jobs:
           SM2A_ADMIN_USERNAME: ${{ secrets.SM2A_ADMIN_USERNAME }}
           SM2A_ADMIN_PASSWORD: ${{ secrets.SM2A_ADMIN_PASSWORD }}
           SM2A_API_URL: ${{ vars.SM2A_API_URL }}
-          PROMOTION_DAG: ${{ vars.PROMOTION_DAG_NAME }}
+          PROMOTION_DAG_NAME: ${{ vars.PROMOTION_DAG_NAME }}
+          PRODUCTION_AIRFLOW_API_VERSION: ${{ vars.PRODUCTION_AIRFLOW_API_VERSION }}
+          AIRFLOW_JWT_SECRET: ${{ secrets.AIRFLOW_JWT_SECRET }}
+          AIRFLOW_JWT_SUB: ${{ secrets.AIRFLOW_JWT_SUB }}
 
         run: |
           pip install -r ./scripts/requirements.txt
@@ -91,5 +94,7 @@ jobs:
                 echo "$collection_id ($config_type) successfully promoted to Production ✅"
               else
                 echo "$collection_id ($config_type) failed to promote to Production ❌"
+                failed=1
               fi
-          done
+          done
+          [ "${failed:-0}" -eq 0 ] || exit 1

docs/README.md

@@ -0,0 +1,42 @@
+# README
+
+The workflows read config from two GitHub Environments (under Settings -> Environments): `staging` and `production`. A job only sees the secrets and vars of the environment it declares. The same secret name, therefore, can hold a different value in each stage.
+
+## Which workflow uses which environment
+
+- `staging` is used by [pr.yml](../.github/workflows/pr.yml). It runs when a PR is opened or updated and publishes to staging Airflow.
+- `production` is used by [promote.yml](../.github/workflows/promote.yml) which [promotion-checker.yml](../.github/workflows/promotion-checker.yml) triggers after the PR is approved. It publishes to production Airflow.
+
+For the GitHub `staging` environment, the following config should be set:
+
+- Secrets, for a staging instance using Airflow 2, `STAGING_SM2A_ADMIN_USERNAME` and `STAGING_SM2A_ADMIN_PASSWORD`.
+- Secrets, for a staging instance using Airflow 3, `AIRFLOW_JWT_SECRET` and `AIRFLOW_JWT_SUB`.
+- Secret for the MDX PR step: `APP_PEM`.
+- Vars: `STAGING_SM2A_API_URL`, `STAGING_AIRFLOW_API_VERSION`, `DATASET_DAG_NAME`, and `ARTIFACT_RETENTION_DAYS`.
+- Vars for the MDX step: `SKIP_MDX_PR_STEP`, `VEDA_CONFIG_REPO_ORG`, `VEDA_CONFIG_REPO_NAME`, `VEDA_CONFIG_APP_ID`, and `GH_ACTOR_EMAIL`.
+
+For the GitHub `production` environment, the following config should be set:
+
+- Secrets, for a production instance using Airflow 2, `SM2A_ADMIN_USERNAME` and `SM2A_ADMIN_PASSWORD`.
+- Secrets, for a production instance using Airflow 3, `AIRFLOW_JWT_SECRET` and `AIRFLOW_JWT_SUB` *which should be different values from what is set in staging*.
+- Vars: `SM2A_API_URL`, `PRODUCTION_AIRFLOW_API_VERSION` and `PROMOTION_DAG_NAME`.
+
+> Note: `*_AIRFLOW_API_VERSION` is the Airflow major version (can be either 2 or 3) which defaults to 2 if unset. The JWT secrets **are only needed for Airflow 3**.
+
+### Repo-level secrets
+
+Set repo level secrets under `Settings` -> `Secrets and variables` -> `Actions`.
+
+- `WORKFLOW_TRIGGER_TOKEN` which is a repo-level secret that isn't environment scoped. It is required for `promote.yml`.
+
+## Environment Variables Reference for Collection and Dataset Promotion
+
+> See the [terminology note](../README.md#step-4-promote-to-production) in the root README for what "promotion" implies.
+>
+
+| Variable | Set in | Used by | Default | Notes |
+| -- | -- | -- | -- | -- |
+| STAGING_AIRFLOW_API_VERSION | Staging GitHub Env | Staging Publish/Promotion | 2 | The Airflow Version for the Staging Environment.  For Airflow 2 = API v1 + Basic auth, Airflow 3 = API v2 + JWT |
+| PRODUCTION_AIRFLOW_API_VERSION | Production GitHub Env | Production Promotion | 2 | Same as above |
+| AIRFLOW_JWT_SECRET | Both (optional) | Any Publish/Promotion step that uses Airflow 3 | none | Conditionally required when the resolved API version is 3. Must match the secret that the Airflow 3 deployment signs with. Can be found in Secrets Manager |
+| AIRFLOW_JWT_SUB | Both (optional) | Airflow 3 only | none | Conditionally required when the resolved API version is 3. The subject claim for the token (the Integer ID from Airflow's user table) |

scripts/airflow_api.py

@@ -0,0 +1,123 @@
+"""
+Supports Airflow 2 and 3 (API v1 / v2)
+"""
+
+import os
+import uuid
+import json
+import http.client
+
+from typing import Dict, Any
+from base64 import b64encode
+
+
+class AirflowAPIError(Exception):
+    def __init__(self, msg, status_code=None):
+        super().__init__(msg)
+        self.status_code = status_code
+
+    pass
+
+
+def _build_request_body(
+    conf: Dict[str, Any], dag_id: str, note: str = "", api_version: str = "2"
+) -> Dict[str, Any]:
+    """Build the request body for the DagRun"""
+    run_id = os.getenv("GITHUB_RUN_ID")
+    if run_id:
+        run_url = (
+            f"{os.getenv('GITHUB_SERVER_URL')}/{os.getenv('GITHUB_REPOSITORY')}"
+            f"/actions/runs/{run_id}"
+        )
+        default_note = f"Run from GitHub Actions: {run_url}"
+    else:
+        default_note = "Run from GitHub Actions veda-data workflow"
+
+    body = {
+        "conf": conf,
+        "dag_run_id": f"{dag_id}-{uuid.uuid4()}",
+        "note": note or default_note,
+    }
+    if api_version != "2":
+        body["logical_date"] = None
+    return body
+
+
+def _generate_jwt_token(secret: str, sub: str, expiration_time: int = 3600) -> str:
+    """Generate a HS512 JWT Token for Airflow API authentication"""
+    import jwt
+    import time
+
+    payload = {
+        "iss": "airflow",
+        "sub": sub,
+        "aud": "apache-airflow",
+        "nbf": int(time.time()),
+        "iat": int(time.time()),
+        "exp": int(time.time()) + expiration_time,
+    }
+    return jwt.encode(payload, secret, algorithm="HS512")
+
+
+def trigger_dag_run(
+    base_api_url: str,
+    dag_id: str,
+    conf: Dict[str, Any],
+    username: str,
+    password: str,
+    api_version: str = None,
+) -> Dict[str, Any]:
+    """
+    Trigger a DAG run with version-aware API handling
+
+    Raises: AirflowAPIError if it fails
+    """
+
+    if api_version is None:
+        api_version = os.getenv("AIRFLOW_API_VERSION") or "2"
+
+    api_path = f"/api/v{'1' if api_version == '2' else '2'}/dags/{dag_id}/dagRuns"
+    request_body = _build_request_body(conf, dag_id, api_version=api_version)
+
+    if api_version == "2":
+        api_token = b64encode(f"{username}:{password}".encode()).decode()
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": "Basic " + api_token,
+        }
+    else:
+        jwt_secret = os.getenv("AIRFLOW_JWT_SECRET")
+        if not jwt_secret:
+            raise AirflowAPIError("AIRFLOW_JWT_SECRET environment variable not set")
+        jwt_sub = os.getenv("AIRFLOW_JWT_SUB")
+        if not jwt_sub:
+            raise AirflowAPIError("AIRFLOW_JWT_SUB environment variable not set")
+        access_token = _generate_jwt_token(jwt_secret, jwt_sub)
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {access_token}",
+        }
+
+    try:
+        http_conn = http.client.HTTPSConnection(base_api_url)
+        http_conn.request("POST", api_path, json.dumps(request_body), headers)
+        response = http_conn.getresponse()
+        response_data = response.read()
+        http_conn.close()
+
+        if response.status >= 400:
+            raise AirflowAPIError(
+                f"Airflow API v{api_version} returns "
+                f"{response.status}: {response_data.decode()}",
+                status_code=response.status,
+            )
+
+        return {
+            "statusCode": response.status,
+            "body": response_data.decode(),
+        }
+
+    except AirflowAPIError:
+        raise
+    except Exception as e:
+        raise AirflowAPIError(f"Failed to trigger the DAG run: {str(e)}")

scripts/promote_collection.py

@@ -1,11 +1,10 @@
 from typing import Dict, Any
 
-import http.client
 import json
 import sys
 import os
-import uuid
-from base64 import b64encode
+
+from airflow_api import AirflowAPIError, trigger_dag_run
 
 
 def trigger_collection_dag(payload: Dict[str, Any], stage: str):
@@ -18,10 +17,12 @@ def trigger_collection_dag(payload: Dict[str, Any], stage: str):
         api_url_env = "STAGING_SM2A_API_URL"
         username_env = "STAGING_SM2A_ADMIN_USERNAME"
         password_env = "STAGING_SM2A_ADMIN_PASSWORD"
+        api_version_env = "STAGING_AIRFLOW_API_VERSION"
     elif stage == "production":
         api_url_env = "SM2A_API_URL"
         username_env = "SM2A_ADMIN_USERNAME"
         password_env = "SM2A_ADMIN_PASSWORD"
+        api_version_env = "PRODUCTION_AIRFLOW_API_VERSION"
     else:
         raise ValueError(
             f"Invalid stage provided: {stage}. Must be 'staging' or 'production'."
@@ -30,38 +31,24 @@ def trigger_collection_dag(payload: Dict[str, Any], stage: str):
     base_api_url = os.getenv(api_url_env)
     username = os.getenv(username_env)
     password = os.getenv(password_env)
+    api_version = os.getenv(api_version_env) or "2"
 
     if not all([base_api_url, username, password]):
-        raise ValueError(
-            f"Missing one or more environment variables: "
-            f"stage is None={stage is None}, "
-            f"username is None={username_env is None}, "
-            f"password is None={password_env is None}"
-        )
-
-    api_token = b64encode(f"{username}:{password}".encode()).decode()
+        raise ValueError(f"Missing required environment variables for stage '{stage}' ")
 
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": "Basic " + api_token,
-    }
-
-    body = {
-        **payload,
-        "dag_run_id": f"{dag_name}-{uuid.uuid4()}",
-        "note": "Run from GitHub Actions eic-data",
-    }
-    http_conn = http.client.HTTPSConnection(base_api_url)
-    http_conn.request(
-        "POST", f"/api/v1/dags/{dag_name}/dagRuns", json.dumps(body), headers
-    )
-    response = http_conn.getresponse()
-    response_data = response.read()
-    http_conn.close()
-
-    print(json.dumps({"statusCode": response.status}))
-    print(response_data.decode())
-    return {"statusCode": response.status, "body": response_data.decode()}
+    try:
+        result = trigger_dag_run(
+            base_api_url=base_api_url,
+            dag_id=dag_name,
+            conf=payload.get("conf", {}),
+            username=username,
+            password=password,
+            api_version=api_version,
+        )
+        print(json.dumps({"statusCode": result["statusCode"], "body": result["body"]}))
+        return result
+    except AirflowAPIError as e:
+        print(json.dumps({"statusCode": e.status_code or 500, "error": str(e)}))
 
 
 if __name__ == "__main__":
@@ -74,7 +61,13 @@ def trigger_collection_dag(payload: Dict[str, Any], stage: str):
 
     except IndexError:
         print("Usage: promote_collection.py <file_name> <stage>")
+        sys.exit(1)
     except FileNotFoundError:
         print(f"Error: File '{sys.argv[1]}' not found.")
+        sys.exit(1)
     except json.JSONDecodeError:
-        raise ValueError(f"Invalid JSON content in file {sys.argv[1]}")
+        print(f"Invalid JSON content in file {sys.argv[1]}")
+        sys.exit(1)
+    except ValueError as e:
+        print(f"Error: {e}")
+        sys.exit(1)