feat: update promotion scripts to accommodate Airflow v1 and v2 (for Airflow 3 upgrade)

.github/workflows/promote.yml

@@ -64,7 +64,7 @@ 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 }}
 
         run: |
           pip install -r ./scripts/requirements.txt
@@ -91,5 +91,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
+          [ "${failed:-0}" -eq 0 ] || exit 1

.github/workflows/test-promotion-script.yml

@@ -0,0 +1,69 @@
+name: Test promotion script
+
+on:
+  push:
+    branches: ['feat/airflow3-flag']
+    paths:
+      # Run the workflow only if files inside this path are updated
+      - 'scripts/**'
+      - 'tests/**'
+      - '.github/workflows/test-promotion-script.yml'
+
+jobs:
+  test-promotion-script:
+    runs-on: ubuntu-latest
+    environment: staging
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+
+      - name: Set up Python
+        uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
+        with:
+          python-version: '3.9'
+
+      - name: Install dependencies
+        run: pip install -r ./scripts/requirements.txt
+
+      - name: Test promotion script (Airflow 3)
+        env:
+            STAGING_SM2A_API_URL: ${{ vars.TESTING_AIRFLOW3_SM2A_API_URL }}
+            STAGING_SM2A_ADMIN_USERNAME: ${{ secrets.TESTING_AIRFLOW3_SM2A_ADMIN_USERNAME}}
+            STAGING_SM2A_ADMIN_PASSWORD: ${{ secrets.TESTING_AIRFLOW3_SM2A_ADMIN_PASSWORD}}
+            STAGING_AIRFLOW_API_VERSION: "3"
+            AIRFLOW_JWT_SECRET: ${{ secrets.AIRFLOW_JWT_SECRET }}
+            AIRFLOW_JWT_SUB: ${{ secrets.AIRFLOW_JWT_SUB }}
+            PYTHONPATH: .
+        run: |
+          response=$(python3 scripts/promote_collection.py tests/test_collection.json staging)
+          status_code=$(echo "$response" | jq -r '.statusCode' | head -n1)
+              echo "Status Code: $status_code"
+
+              # Update status message based on response code
+              if [[ $status_code -eq 200 ]] || [[ $status_code -eq 201 ]]; then
+                echo "Airflow 3 test passed"
+              else
+                error_msg=$(echo "$response" | jq -r '.error')
+                echo "$collection_id failed to publish ❌"
+                echo "Error: $error_msg"
+                exit 1
+              fi
+
+      - name: Test promotion script (Airflow 2)
+        env:
+            STAGING_SM2A_API_URL: ${{ vars.TESTING_AIRFLOW2_SM2A_API_URL }}
+            STAGING_SM2A_ADMIN_USERNAME: ${{ secrets.TESTING_AIRFLOW2_SM2A_ADMIN_USERNAME}}
+            STAGING_SM2A_ADMIN_PASSWORD: ${{ secrets.TESTING_AIRFLOW2_SM2A_ADMIN_PASSWORD}}
+            STAGING_AIRFLOW_API_VERSION: "2"
+            PYTHONPATH: .
+        run: |
+          response=$(python3 scripts/promote_collection.py tests/test_collection_tenant.json staging)
+          status_code=$(echo "$response" | jq -r '.statusCode' | head -n1)
+              echo "Status Code: $status_code"
+
+              # Update status message based on response code
+              if [[ $status_code -eq 200 ]] || [[ $status_code -eq 201 ]]; then
+                echo "Airflow 2 test passed"
+              else
+                echo "$collection_id failed to publish ❌"
+                exit 1
+              fi

docs/README.md

@@ -0,0 +1,11 @@
+
+# README
+
+## Environment Variables Reference for Collection and Dataset Promotion
+
+| Variable | Used by | Default | Notes |
+| -- | -- | -- | -- |
+| STAGING_AIRFLOW_API_VERSION | Staging 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 Promotion | 2 | Same as above |
+| AIRFLOW_JWT_SECRET | Any 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 | Airflow 3 only | none | Conditionally required when the resolved API verison 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 veda-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)