refresh_service.py

"""
DAG Refresh Service — lightweight Flask sidecar for git-syncing DAGs.

Endpoints:
    GET/POST /refresh  — pull latest DAGs from the configured git repo
    GET      /health   — liveness check (also pings Airflow webserver)
    GET      /config   — show current (non-sensitive) configuration
    GET      /status   — git SHA, last sync time, repo info
"""

import fcntl
import logging
import os
import subprocess
import time
from datetime import datetime, timezone

from flask import Flask, jsonify, request

# ---- Logging ----
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("refresh-service")

app = Flask(__name__)

# ---- Configuration ----
AIRFLOW_HOME = os.environ.get("AIRFLOW_HOME", "/opt/airflow")
DAGS_DIR = os.path.join(AIRFLOW_HOME, "dags")
DAG_REPO_URL = os.environ.get("DAG_REPO_URL", "https://github.com/subhamgiri460/myworkflows.git")
DAG_REPO_BRANCH = os.environ.get("DAG_REPO_BRANCH", "main")
DAG_REPO_TOKEN = os.environ.get("DAG_REPO_TOKEN", "")
AIRFLOW_WEB_PORT = os.environ.get("AIRFLOW_WEB_PORT", "8080")
REFRESH_SERVICE_PORT = os.environ.get("REFRESH_SERVICE_PORT", "5000")

LOCK_FILE = "/tmp/dag_sync.lock"
_last_sync: dict = {}


def _get_authenticated_url() -> str:
    """Inject PAT token into the repo URL for private repos."""
    if DAG_REPO_TOKEN:
        return DAG_REPO_URL.replace("https://", f"https://{DAG_REPO_TOKEN}@")
    return DAG_REPO_URL


def _run_git(args: list[str], cwd: str | None = None, timeout: int = 60) -> subprocess.CompletedProcess:
    """Run a git command and return the result."""
    return subprocess.run(
        ["git"] + args,
        cwd=cwd,
        capture_output=True,
        text=True,
        timeout=timeout,
    )


def _get_head_sha() -> str | None:
    """Return the current HEAD SHA in the DAGs directory."""
    try:
        result = _run_git(["rev-parse", "--short", "HEAD"], cwd=DAGS_DIR, timeout=10)
        return result.stdout.strip() if result.returncode == 0 else None
    except Exception:
        return None


def _sync_repo() -> tuple[dict, int]:
    """Clone or update the DAG repository. Returns (response_body, status_code)."""
    global _last_sync

    os.makedirs(DAGS_DIR, exist_ok=True)

    auth_url = _get_authenticated_url()
    is_update = os.path.isdir(os.path.join(DAGS_DIR, ".git"))

    if is_update:
        # Update remote URL in case PAT token changed
        _run_git(["remote", "set-url", "origin", auth_url], cwd=DAGS_DIR)

        fetch = _run_git(["fetch", "origin", DAG_REPO_BRANCH], cwd=DAGS_DIR)
        if fetch.returncode != 0:
            logger.error("git fetch failed: %s", fetch.stderr)
            return {"status": "error", "message": f"git fetch failed: {fetch.stderr}"}, 500

        reset = _run_git(["reset", "--hard", f"origin/{DAG_REPO_BRANCH}"], cwd=DAGS_DIR)
        if reset.returncode != 0:
            logger.error("git reset failed: %s", reset.stderr)
            return {"status": "error", "message": f"git reset failed: {reset.stderr}"}, 500
        action = "updated"
    else:
        clone = _run_git(
            ["clone", "--depth", "1", "--branch", DAG_REPO_BRANCH, auth_url, DAGS_DIR],
            timeout=120,
        )
        if clone.returncode != 0:
            logger.error("git clone failed: %s", clone.stderr)
            return {"status": "error", "message": f"git clone failed: {clone.stderr}"}, 500
        action = "cloned"

    sha = _get_head_sha()
    _last_sync = {
        "action": action,
        "sha": sha,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

    logger.info("DAGs %s — commit %s", action, sha)
    return {
        "status": "success",
        "message": f"DAGs {action} successfully",
        "repo": DAG_REPO_URL,
        "branch": DAG_REPO_BRANCH,
        "sha": sha,
    }, 200


# ---- Request logging ----
@app.before_request
def log_request():
    logger.info("%s %s", request.method, request.path)


# ---- Endpoints ----
@app.route("/refresh", methods=["GET", "POST"])
def refresh_dags():
    """Pull latest DAGs from git. Uses a file lock to prevent concurrent syncs."""
    try:
        lock_fd = open(LOCK_FILE, "w")
        acquired = False
        try:
            fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            acquired = True
        except BlockingIOError:
            return jsonify({"status": "busy", "message": "A sync is already in progress"}), 429

        if acquired:
            body, status = _sync_repo()
            fcntl.flock(lock_fd, fcntl.LOCK_UN)
            lock_fd.close()
            return jsonify(body), status

    except subprocess.TimeoutExpired:
        return jsonify({"status": "error", "message": "Git operation timed out"}), 504
    except Exception as e:
        logger.exception("Unexpected error during refresh")
        return jsonify({"status": "error", "message": str(e)}), 500


@app.route("/health")
def health():
    """Liveness probe — also checks if Airflow webserver is reachable."""
    airflow_ok = False
    try:
        import urllib.request
        resp = urllib.request.urlopen(f"http://127.0.0.1:{AIRFLOW_WEB_PORT}/health", timeout=5)
        airflow_ok = resp.status == 200
    except Exception:
        pass

    return jsonify({
        "status": "healthy",
        "airflow_webserver": "up" if airflow_ok else "starting",
    }), 200


@app.route("/config")
def show_config():
    """Return non-sensitive configuration values."""
    return jsonify({
        "dags_dir": DAGS_DIR,
        "repo_url": DAG_REPO_URL,
        "branch": DAG_REPO_BRANCH,
        "airflow_web_port": AIRFLOW_WEB_PORT,
        "refresh_service_port": REFRESH_SERVICE_PORT,
    }), 200


@app.route("/status")
def status():
    """Return the last sync status, current git SHA, and uptime."""
    return jsonify({
        "last_sync": _last_sync or "no sync yet",
        "current_sha": _get_head_sha(),
        "repo": DAG_REPO_URL,
        "branch": DAG_REPO_BRANCH,
    }), 200


if __name__ == "__main__":
    port = int(REFRESH_SERVICE_PORT)
    app.run(host="0.0.0.0", port=port)