appfactory-worker/app/main.py

import json
import os
import signal
import socket
import sqlite3
import subprocess
import threading
import time
from pathlib import Path

from fastapi import FastAPI

DB_FILE = Path("/opt/appfactory/data/appfactory/appfactory.db")
TOOLS_DIR = Path("/tools")
MAINTENANCE_DIR = Path("/opt/appfactory/workspace/appfactory-tools/maintenance")
ALERTS_DIR = Path("/opt/appfactory/workspace/appfactory-tools/alerts")

WORKER_ID = os.getenv("APPFACTORY_WORKER_ID", socket.gethostname())
POLL_SECONDS = int(os.getenv("APPFACTORY_WORKER_POLL_SECONDS", "3"))
HEARTBEAT_SECONDS = int(os.getenv("APPFACTORY_WORKER_HEARTBEAT_SECONDS", "10"))
CANCEL_CHECK_SECONDS = int(os.getenv("APPFACTORY_CANCEL_CHECK_SECONDS", "2"))

app = FastAPI(title="AppFactory Worker")

_worker_thread = None
_heartbeat_thread = None
_worker_started = False


def get_connection():
    con = sqlite3.connect(DB_FILE, timeout=30)
    con.row_factory = sqlite3.Row
    return con


def append_job_log(job_id: int, stream: str, message: str):
    con = get_connection()
    con.execute(
        """
        INSERT INTO job_logs (job_id, stream, message, created_at)
        VALUES (?, ?, ?, CURRENT_TIMESTAMP)
        """,
        (job_id, stream, message),
    )
    con.commit()
    con.close()


def get_job_status(job_id: int) -> str | None:
    con = get_connection()
    row = con.execute(
        "SELECT status FROM jobs WHERE id = ?",
        (job_id,),
    ).fetchone()
    con.close()

    if not row:
        return None

    return row["status"]


def update_worker_status(current_job_id=None):
    con = get_connection()

    con.execute(
        """
        INSERT INTO workers (
            id,
            status,
            started_at,
            last_seen_at,
            current_job_id,
            metadata_json
        )
        VALUES (
            ?,
            'online',
            CURRENT_TIMESTAMP,
            CURRENT_TIMESTAMP,
            ?,
            ?
        )
        ON CONFLICT(id)
        DO UPDATE SET
            status = 'online',
            last_seen_at = CURRENT_TIMESTAMP,
            current_job_id = excluded.current_job_id,
            metadata_json = excluded.metadata_json
        """,
        (
            WORKER_ID,
            current_job_id,
            json.dumps(
                {
                    "poll_seconds": POLL_SECONDS,
                    "heartbeat_seconds": HEARTBEAT_SECONDS,
                    "cancel_check_seconds": CANCEL_CHECK_SECONDS,
                },
                ensure_ascii=False,
            ),
        ),
    )

    con.commit()
    con.close()


def heartbeat_loop():
    while True:
        try:
            update_worker_status()
        except Exception as exc:
            print(f"Heartbeat error: {exc}", flush=True)

        time.sleep(HEARTBEAT_SECONDS)


def cleanup_stale_locks():
    con = get_connection()

    con.execute(
        """
        DELETE FROM deployment_locks
        WHERE expires_at IS NOT NULL
          AND expires_at < CURRENT_TIMESTAMP
        """
    )

    con.commit()
    con.close()


def acquire_lock(job: dict) -> bool:
    con = get_connection()

    try:
        con.execute(
            """
            INSERT INTO deployment_locks (
                target_type,
                target_id,
                job_id,
                worker_id,
                acquired_at,
                expires_at
            )
            VALUES (
                ?,
                ?,
                ?,
                ?,
                CURRENT_TIMESTAMP,
                datetime('now', '+30 minutes')
            )
            """,
            (
                job.get("target_type") or job["type"],
                job["target_id"],
                job["id"],
                WORKER_ID,
            ),
        )

        con.commit()
        return True

    except sqlite3.IntegrityError:
        return False

    finally:
        con.close()


def release_lock(job: dict):
    con = get_connection()

    con.execute(
        """
        DELETE FROM deployment_locks
        WHERE target_type = ?
          AND target_id = ?
          AND job_id = ?
        """,
        (
            job.get("target_type") or job["type"],
            job["target_id"],
            job["id"],
        ),
    )

    con.commit()
    con.close()


def extend_lock(job: dict):
    con = get_connection()

    con.execute(
        """
        UPDATE deployment_locks
        SET expires_at = datetime('now', '+30 minutes')
        WHERE target_type = ?
          AND target_id = ?
          AND job_id = ?
        """,
        (
            job.get("target_type") or job["type"],
            job["target_id"],
            job["id"],
        ),
    )

    con.commit()
    con.close()


def recover_stale_jobs():
    con = get_connection()

    rows = con.execute(
        """
        SELECT id
        FROM jobs
        WHERE status IN ('running', 'cancelled_requested')
          AND started_at < datetime('now', '-15 minutes')
        """
    ).fetchall()

    for row in rows:
        job_id = row["id"]

        con.execute(
            """
            UPDATE jobs
            SET status = 'failed',
                finished_at = CURRENT_TIMESTAMP,
                error_text = 'Worker recovery: stale running job'
            WHERE id = ?
            """,
            (job_id,),
        )

        con.execute(
            """
            INSERT INTO job_logs (job_id, stream, message, created_at)
            VALUES (?, 'system', 'Worker recovery marked stale running job as failed', CURRENT_TIMESTAMP)
            """,
            (job_id,),
        )

        con.execute(
            """
            DELETE FROM deployment_locks
            WHERE job_id = ?
            """,
            (job_id,),
        )

    con.commit()
    con.close()


def update_job_status(job_id: int, status: str, error_text: str | None = None, result: dict | None = None):
    con = get_connection()

    if status == "running":
        con.execute(
            """
            UPDATE jobs
            SET status = 'running',
                started_at = COALESCE(started_at, CURRENT_TIMESTAMP),
                worker_id = ?
            WHERE id = ?
            """,
            (WORKER_ID, job_id),
        )
    elif status in ("success", "failed", "cancelled"):
        con.execute(
            """
            UPDATE jobs
            SET status = ?,
                finished_at = CURRENT_TIMESTAMP,
                error_text = ?,
                result_json = ?
            WHERE id = ?
            """,
            (
                status,
                error_text,
                json.dumps(result or {}, ensure_ascii=False),
                job_id,
            ),
        )

    con.commit()
    con.close()


def get_next_job():
    con = get_connection()

    con.execute("BEGIN IMMEDIATE")

    row = con.execute(
        """
        SELECT *
        FROM jobs
        WHERE status = 'queued'
        ORDER BY id ASC
        LIMIT 1
        """
    ).fetchone()

    if not row:
        con.commit()
        con.close()
        return None

    job_id = row["id"]

    con.execute(
        """
        UPDATE jobs
        SET status = 'running',
            started_at = CURRENT_TIMESTAMP,
            worker_id = ?
        WHERE id = ?
          AND status = 'queued'
        """,
        (WORKER_ID, job_id),
    )

    con.commit()
    con.close()

    return dict(row)


def build_env(job: dict):
    payload = json.loads(job.get("payload_json") or "{}")

    env = os.environ.copy()
    env["APPFACTORY_TRIGGER_SOURCE"] = job.get("source") or "worker"
    env["APPFACTORY_TRIGGERED_BY_USER_ID"] = str(job.get("created_by_user_id") or "")
    env["APPFACTORY_TRIGGERED_BY_USERNAME"] = str(job.get("created_by_username") or "")
    env["APPFACTORY_TRIGGERED_BY_DISPLAY_NAME"] = str(job.get("created_by_display_name") or "")
    env["APPFACTORY_COMMIT_AUTHOR"] = str(payload.get("commit_author") or "")
    env["APPFACTORY_PUSHER"] = str(payload.get("pusher") or "")
    env["APPFACTORY_COMMIT_SHA"] = str(payload.get("commit_sha") or "")

    if job.get("type") == "run_alert_script":
        env["APPFACTORY_ALERT_EVENT_ID"] = str(payload.get("alert_event_id") or "")
        env["APPFACTORY_ALERT_RULE_ID"] = str(payload.get("rule_id") or "")
        env["APPFACTORY_ALERT_EVENT_TYPE"] = str(payload.get("event_type") or "")
        env["APPFACTORY_ALERT_SERVICE_ID"] = str(payload.get("service_id") or "")
        env["APPFACTORY_ALERT_INCIDENT_ID"] = str(payload.get("incident_id") or "")
        env["APPFACTORY_ALERT_PAYLOAD_JSON"] = json.dumps(payload.get("alert_payload") or {}, ensure_ascii=False)

    return env


def validate_script_name(script_name: str):
    if not script_name:
        raise ValueError("Missing script_name")

    if "/" in script_name or "\\" in script_name or ".." in script_name:
        raise ValueError("Invalid script name")

    if not script_name.endswith(".sh"):
        raise ValueError("Script must end with .sh")


def command_for_job(job: dict):
    job_type = job["type"]
    target_id = job["target_id"]

    if job_type == "deploy_app":
        return [str(TOOLS_DIR / "deploy-app.sh"), target_id]

    if job_type == "deploy_core_service":
        return [str(TOOLS_DIR / "deploy-core-service.sh"), target_id]

    if job_type == "run_script":
        payload = json.loads(job.get("payload_json") or "{}")
        script_name = payload.get("script_name") or target_id

        validate_script_name(script_name)

        script_path = MAINTENANCE_DIR / script_name

        if not script_path.exists():
            raise ValueError(f"Script not found: {script_name}")

        if not script_path.is_file():
            raise ValueError(f"Script is not a file: {script_name}")

        return [str(script_path)]

    if job_type == "run_alert_script":
        payload = json.loads(job.get("payload_json") or "{}")
        script_name = payload.get("script_name") or target_id

        validate_script_name(script_name)

        script_path = ALERTS_DIR / script_name

        if not script_path.exists():
            raise ValueError(f"Alert script not found: {script_name}")

        if not script_path.is_file():
            raise ValueError(f"Alert script is not a file: {script_name}")

        return [str(script_path)]

    raise ValueError(f"Unsupported job type: {job_type}")


def update_alert_event_from_job(job: dict, status: str, error_text: str | None):
    if job.get("type") != "run_alert_script":
        return

    payload = json.loads(job.get("payload_json") or "{}")
    alert_event_id = payload.get("alert_event_id")

    if not alert_event_id:
        return

    con = get_connection()

    if status == "success":
        alert_status = "success"
    elif status == "cancelled":
        alert_status = "cancelled"
    else:
        alert_status = "failed"

    con.execute(
        """
        UPDATE alert_events
        SET status = ?,
            error_text = ?,
            processed_at = CURRENT_TIMESTAMP
        WHERE id = ?
        """,
        (
            alert_status,
            error_text,
            alert_event_id,
        ),
    )

    con.commit()
    con.close()


def terminate_process_group(process: subprocess.Popen, job_id: int):
    append_job_log(job_id, "system", "Cancel requested. Terminating process group...")

    try:
        os.killpg(os.getpgid(process.pid), signal.SIGTERM)
        process.wait(timeout=10)
        append_job_log(job_id, "system", "Process group terminated.")
    except subprocess.TimeoutExpired:
        append_job_log(job_id, "system", "Process group did not terminate in time. Killing process group...")
        try:
            os.killpg(os.getpgid(process.pid), signal.SIGKILL)
        except Exception:
            process.kill()
        process.wait(timeout=10)
        append_job_log(job_id, "system", "Process group killed.")
    except Exception as exc:
        append_job_log(job_id, "stderr", f"Error while terminating process group: {exc}")
        try:
            process.kill()
            process.wait(timeout=10)
        except Exception as kill_exc:
            append_job_log(job_id, "stderr", f"Error while killing process: {kill_exc}")


def execute_job(job: dict):
    job_id = int(job["id"])
    target = job["target_id"]

    update_worker_status(current_job_id=job_id)
    append_job_log(job_id, "system", f"Worker {WORKER_ID} started job {job_id} for {target}")
    update_job_status(job_id, "running")

    try:
        command = command_for_job(job)
        env = build_env(job)

        append_job_log(job_id, "system", "Command: " + " ".join(command))

        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            env=env,
            start_new_session=True,
        )

        while process.poll() is None:
            extend_lock(job)

            status = get_job_status(job_id)

            if status == "cancelled_requested":
                terminate_process_group(process, job_id)

                update_job_status(
                    job_id,
                    "cancelled",
                    error_text="Job cancelled by user request",
                    result={"cancelled": True},
                )
                update_alert_event_from_job(job, "cancelled", "Job cancelled by user request")
                append_job_log(job_id, "system", "Job cancelled.")
                return

            time.sleep(CANCEL_CHECK_SECONDS)

        stdout, stderr = process.communicate()
        returncode = process.returncode

        if stdout:
            append_job_log(job_id, "stdout", stdout)

        if stderr:
            append_job_log(job_id, "stderr", stderr)

        if get_job_status(job_id) == "cancelled_requested":
            update_job_status(
                job_id,
                "cancelled",
                error_text="Job cancelled by user request",
                result={"returncode": returncode, "cancelled": True},
            )
            update_alert_event_from_job(job, "cancelled", "Job cancelled by user request")
            append_job_log(job_id, "system", "Job cancelled after process finished.")
            return

        if returncode == 0:
            update_job_status(
                job_id,
                "success",
                result={
                    "returncode": returncode,
                    "stdout_length": len(stdout or ""),
                    "stderr_length": len(stderr or ""),
                },
            )
            update_alert_event_from_job(job, "success", None)
            append_job_log(job_id, "system", "Job finished successfully")
        else:
            error_text = f"Command failed with return code {returncode}"
            update_job_status(
                job_id,
                "failed",
                error_text=error_text,
                result={"returncode": returncode},
            )
            update_alert_event_from_job(job, "failed", error_text)
            append_job_log(job_id, "system", f"Job failed with return code {returncode}")

    except Exception as exc:
        error_text = str(exc)
        update_job_status(job_id, "failed", error_text=error_text)
        update_alert_event_from_job(job, "failed", error_text)
        append_job_log(job_id, "stderr", error_text)

    finally:
        update_worker_status(current_job_id=None)


def worker_loop():
    while True:
        try:
            cleanup_stale_locks()
            recover_stale_jobs()

            job = get_next_job()

            if job:
                if not acquire_lock(job):
                    append_job_log(
                        int(job["id"]),
                        "system",
                        "Deployment lock is already held for this target. Re-queueing job.",
                    )

                    con = get_connection()
                    con.execute(
                        """
                        UPDATE jobs
                        SET status = 'queued',
                            started_at = NULL,
                            worker_id = NULL
                        WHERE id = ?
                        """,
                        (job["id"],),
                    )
                    con.commit()
                    con.close()

                    time.sleep(POLL_SECONDS)
                    continue

                try:
                    execute_job(job)
                finally:
                    release_lock(job)
            else:
                time.sleep(POLL_SECONDS)

        except Exception as exc:
            print(f"Worker loop error: {exc}", flush=True)
            time.sleep(POLL_SECONDS)


@app.on_event("startup")
def start_worker():
    global _worker_thread
    global _heartbeat_thread
    global _worker_started

    if _worker_started:
        return

    cleanup_stale_locks()
    recover_stale_jobs()
    update_worker_status(current_job_id=None)

    _worker_started = True

    _heartbeat_thread = threading.Thread(target=heartbeat_loop, daemon=True)
    _heartbeat_thread.start()

    _worker_thread = threading.Thread(target=worker_loop, daemon=True)
    _worker_thread.start()


@app.get("/health")
def health():
    return {
        "status": "ok",
        "worker_id": WORKER_ID,
        "poll_seconds": POLL_SECONDS,
        "heartbeat_seconds": HEARTBEAT_SECONDS,
        "cancel_check_seconds": CANCEL_CHECK_SECONDS,
    }