diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 0ca67fc..0676a24 100644 --- a/.github/workflows/pr.yml +++ b/.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 diff --git a/.github/workflows/promote.yml b/.github/workflows/promote.yml index 3e8c5b0..f94b678 100644 --- a/.github/workflows/promote.yml +++ b/.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 \ No newline at end of file + done + [ "${failed:-0}" -eq 0 ] || exit 1 \ No newline at end of file diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..41dd9df --- /dev/null +++ b/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) | diff --git a/scripts/airflow_api.py b/scripts/airflow_api.py new file mode 100644 index 0000000..583f840 --- /dev/null +++ b/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)}") diff --git a/scripts/promote_collection.py b/scripts/promote_collection.py index 1ad6004..80d6dd2 100644 --- a/scripts/promote_collection.py +++ b/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 ") + 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) diff --git a/scripts/promote_dataset.py b/scripts/promote_dataset.py index 7f99474..1cd3e20 100644 --- a/scripts/promote_dataset.py +++ b/scripts/promote_dataset.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 class MissingFieldError(Exception): @@ -13,22 +12,12 @@ class MissingFieldError(Exception): def validate_discovery_item_config(item: Dict[str, Any]) -> Dict[str, Any]: - if "bucket" not in item: - raise MissingFieldError( - f"Missing required field 'bucket' in discovery item: {item}" - ) - if "discovery" not in item: - raise MissingFieldError( - f"Missing required field 'discovery' in discovery item: {item}" - ) - if "filename_regex" not in item: - raise MissingFieldError( - f"Missing required field 'filename_regex' in discovery item: {item}" - ) - if "prefix" not in item: - raise MissingFieldError( - f"Missing required field 'prefix' in discovery item: {item}" - ) + required_fields = ["bucket", "discovery", "filename_regex", "prefix"] + for field in required_fields: + if field not in item: + raise MissingFieldError( + f"Missing required field '{field}' in discovery item: {item}" + ) return item @@ -37,40 +26,27 @@ def publish_to_staging(payload): dataset_pipeline_dag = os.getenv("DATASET_DAG_NAME", "veda_dataset_pipeline") username = os.getenv("STAGING_SM2A_ADMIN_USERNAME") password = os.getenv("STAGING_SM2A_ADMIN_PASSWORD") + api_version_env = os.getenv("STAGING_AIRFLOW_API_VERSION") or "2" - api_token = b64encode(f"{username}:{password}".encode()).decode() - - if not base_api_url or not api_token: + if not all([base_api_url, username, password]): raise ValueError( - "STAGING_SM2A_API_URL or STAGING_SM2A_ADMIN_USERNAME" - + " or STAGING_SM2A_ADMIN_PASSWORD is not" - + " set in the environment variables." + "Missing required env vars STAGING_SM2A_API_URL, " + "STAGING_SM2A_ADMIN_USERNAME, STAGING_SM2A_ADMIN_PASSWORD" ) - headers = { - "Content-Type": "application/json", - "Authorization": "Basic " + api_token, - } - - body = { - **payload, - "dag_run_id": f"{dataset_pipeline_dag}-{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/{dataset_pipeline_dag}/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=dataset_pipeline_dag, + conf=payload.get("conf", {}), + username=username, + password=password, + api_version=api_version_env, + ) + print(json.dumps({"statusCode": result["statusCode"]})) + return result + except AirflowAPIError as e: + print(json.dumps({"statusCode": e.status_code or 500, "error": str(e)})) def promote_to_production(payload): @@ -78,37 +54,29 @@ def promote_to_production(payload): promotion_dag = os.getenv("PROMOTION_DAG_NAME", "veda_promotion_pipeline") username = os.getenv("SM2A_ADMIN_USERNAME") password = os.getenv("SM2A_ADMIN_PASSWORD") + api_version_env = os.getenv("PRODUCTION_AIRFLOW_API_VERSION") or "2" - api_token = b64encode(f"{username}:{password}".encode()).decode() - - if not base_api_url or not api_token: + if not all([base_api_url, username, password]): raise ValueError( - "SM2A_API_URL or SM2A_ADMIN_USERNAME or SM2A_ADMIN_PASSWORD is not" - + " set in the environment variables." + "Missing required env vars SM2A_API_URL, " + "SM2A_ADMIN_USERNAME, SM2A_ADMIN_PASSWORD" ) - headers = { - "Content-Type": "application/json", - "Authorization": "Basic " + api_token, - } - payload["conf"].setdefault("transfer", False) - body = { - **payload, - "dag_run_id": f"{promotion_dag}-{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/{promotion_dag}/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=promotion_dag, + conf=payload.get("conf", {}), + username=username, + password=password, + api_version=api_version_env, + ) + print(json.dumps({"statusCode": result["statusCode"]})) + return result + except AirflowAPIError as e: + print(json.dumps({"statusCode": e.status_code or 500, "error": str(e)})) if __name__ == "__main__": @@ -132,8 +100,17 @@ def promote_to_production(payload): publish_to_staging(dag_payload) except IndexError: - print("Usage: promote_collection.py ") + print("Usage: promote_dataset.py ") + 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) + except MissingFieldError as e: + print(f"Error: {e}") + sys.exit(1) diff --git a/scripts/requirements.txt b/scripts/requirements.txt index 4818cc5..8eb9318 100644 --- a/scripts/requirements.txt +++ b/scripts/requirements.txt @@ -1 +1,2 @@ -pyyaml \ No newline at end of file +pyyaml +PyJWT