altool/airflow

1"""2DAG Refresh Service — lightweight Flask sidecar for git-syncing DAGs.3 4Endpoints:5    GET/POST /refresh  — pull latest DAGs from the configured git repo6    GET      /health   — liveness check (also pings Airflow webserver)7    GET      /config   — show current (non-sensitive) configuration8    GET      /status   — git SHA, last sync time, repo info9"""10 11import fcntl12import logging13import os14import subprocess15import time16from datetime import datetime, timezone17 18from flask import Flask, jsonify, request19 20# ---- Logging ----21logging.basicConfig(22    level=logging.INFO,23    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",24    datefmt="%Y-%m-%d %H:%M:%S",25)26logger = logging.getLogger("refresh-service")27 28app = Flask(__name__)29 30# ---- Configuration ----31AIRFLOW_HOME = os.environ.get("AIRFLOW_HOME", "/opt/airflow")32DAGS_DIR = os.path.join(AIRFLOW_HOME, "dags")33DAG_REPO_URL = os.environ.get("DAG_REPO_URL", "https://github.com/subhamgiri460/myworkflows.git")34DAG_REPO_BRANCH = os.environ.get("DAG_REPO_BRANCH", "main")35DAG_REPO_TOKEN = os.environ.get("DAG_REPO_TOKEN", "")36AIRFLOW_WEB_PORT = os.environ.get("AIRFLOW_WEB_PORT", "8080")37REFRESH_SERVICE_PORT = os.environ.get("REFRESH_SERVICE_PORT", "5000")38 39LOCK_FILE = "/tmp/dag_sync.lock"40_last_sync: dict = {}41 42 43def _get_authenticated_url() -> str:44    """Inject PAT token into the repo URL for private repos."""45    if DAG_REPO_TOKEN:46        return DAG_REPO_URL.replace("https://", f"https://{DAG_REPO_TOKEN}@")47    return DAG_REPO_URL48 49 50def _run_git(args: list[str], cwd: str | None = None, timeout: int = 60) -> subprocess.CompletedProcess:51    """Run a git command and return the result."""52    return subprocess.run(53        ["git"] + args,54        cwd=cwd,55        capture_output=True,56        text=True,57        timeout=timeout,58    )59 60 61def _get_head_sha() -> str | None:62    """Return the current HEAD SHA in the DAGs directory."""63    try:64        result = _run_git(["rev-parse", "--short", "HEAD"], cwd=DAGS_DIR, timeout=10)65        return result.stdout.strip() if result.returncode == 0 else None66    except Exception:67        return None68 69 70def _sync_repo() -> tuple[dict, int]:71    """Clone or update the DAG repository. Returns (response_body, status_code)."""72    global _last_sync73 74    os.makedirs(DAGS_DIR, exist_ok=True)75 76    auth_url = _get_authenticated_url()77    is_update = os.path.isdir(os.path.join(DAGS_DIR, ".git"))78 79    if is_update:80        # Update remote URL in case PAT token changed81        _run_git(["remote", "set-url", "origin", auth_url], cwd=DAGS_DIR)82 83        fetch = _run_git(["fetch", "origin", DAG_REPO_BRANCH], cwd=DAGS_DIR)84        if fetch.returncode != 0:85            logger.error("git fetch failed: %s", fetch.stderr)86            return {"status": "error", "message": f"git fetch failed: {fetch.stderr}"}, 50087 88        reset = _run_git(["reset", "--hard", f"origin/{DAG_REPO_BRANCH}"], cwd=DAGS_DIR)89        if reset.returncode != 0:90            logger.error("git reset failed: %s", reset.stderr)91            return {"status": "error", "message": f"git reset failed: {reset.stderr}"}, 50092        action = "updated"93    else:94        clone = _run_git(95            ["clone", "--depth", "1", "--branch", DAG_REPO_BRANCH, auth_url, DAGS_DIR],96            timeout=120,97        )98        if clone.returncode != 0:99            logger.error("git clone failed: %s", clone.stderr)100            return {"status": "error", "message": f"git clone failed: {clone.stderr}"}, 500101        action = "cloned"102 103    sha = _get_head_sha()104    _last_sync = {105        "action": action,106        "sha": sha,107        "timestamp": datetime.now(timezone.utc).isoformat(),108    }109 110    logger.info("DAGs %s — commit %s", action, sha)111    return {112        "status": "success",113        "message": f"DAGs {action} successfully",114        "repo": DAG_REPO_URL,115        "branch": DAG_REPO_BRANCH,116        "sha": sha,117    }, 200118 119 120# ---- Request logging ----121@app.before_request122def log_request():123    logger.info("%s %s", request.method, request.path)124 125 126# ---- Endpoints ----127@app.route("/refresh", methods=["GET", "POST"])128def refresh_dags():129    """Pull latest DAGs from git. Uses a file lock to prevent concurrent syncs."""130    try:131        lock_fd = open(LOCK_FILE, "w")132        acquired = False133        try:134            fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)135            acquired = True136        except BlockingIOError:137            return jsonify({"status": "busy", "message": "A sync is already in progress"}), 429138 139        if acquired:140            body, status = _sync_repo()141            fcntl.flock(lock_fd, fcntl.LOCK_UN)142            lock_fd.close()143            return jsonify(body), status144 145    except subprocess.TimeoutExpired:146        return jsonify({"status": "error", "message": "Git operation timed out"}), 504147    except Exception as e:148        logger.exception("Unexpected error during refresh")149        return jsonify({"status": "error", "message": str(e)}), 500150 151 152@app.route("/health")153def health():154    """Liveness probe — also checks if Airflow webserver is reachable."""155    airflow_ok = False156    try:157        import urllib.request158        resp = urllib.request.urlopen(f"http://127.0.0.1:{AIRFLOW_WEB_PORT}/health", timeout=5)159        airflow_ok = resp.status == 200160    except Exception:161        pass162 163    return jsonify({164        "status": "healthy",165        "airflow_webserver": "up" if airflow_ok else "starting",166    }), 200167 168 169@app.route("/config")170def show_config():171    """Return non-sensitive configuration values."""172    return jsonify({173        "dags_dir": DAGS_DIR,174        "repo_url": DAG_REPO_URL,175        "branch": DAG_REPO_BRANCH,176        "airflow_web_port": AIRFLOW_WEB_PORT,177        "refresh_service_port": REFRESH_SERVICE_PORT,178    }), 200179 180 181@app.route("/status")182def status():183    """Return the last sync status, current git SHA, and uptime."""184    return jsonify({185        "last_sync": _last_sync or "no sync yet",186        "current_sha": _get_head_sha(),187        "repo": DAG_REPO_URL,188        "branch": DAG_REPO_BRANCH,189    }), 200190 191 192if __name__ == "__main__":193    port = int(REFRESH_SERVICE_PORT)194    app.run(host="0.0.0.0", port=port)