cross-py-builder/ctrl/app.py

import datetime
import json
import logging
import os
import queue

from flask import Flask, jsonify, request, send_from_directory, Response
from werkzeug.exceptions import HTTPException

import ctrl.db as db
import ctrl.settings as settings
import ctrl.workers as wrk
from ctrl.scheduler import make_scheduler

log = logging.getLogger("ctrl")

app = Flask(__name__, static_folder="static", static_url_path="")
scheduler = make_scheduler()

VIEW_STATIC = os.path.join(os.path.dirname(__file__), "static")

_initialized = False


def _ensure_ready():
    global _initialized
    if _initialized:
        return
    settings.ensure_dirs()
    db.init_db()
    _initialized = True


@app.before_request
def _before_request():
    _ensure_ready()


@app.route("/")
def index():
    return send_from_directory(VIEW_STATIC, "index.html")


@app.get("/api/workers")
def api_workers():
    return jsonify(wrk.probe_all())


@app.post("/api/workers")
def api_add_worker():
    data = request.get_json(silent=True) or {}
    spec = str(data.get("spec") or "").strip()
    host = str(data.get("host") or "").strip()
    port = data.get("port")
    try:
        if spec:
            w = wrk.parse_worker(spec)
        elif host and port:
            w = wrk.parse_worker(f"{host}:{int(port)}")
        else:
            return jsonify({"error": "Provide spec ('host:port') or host+port"}), 400
    except (ValueError, TypeError) as e:
        return jsonify({"error": str(e)}), 400

    db.add_worker(w["host"], w["port"])
    db.set_queued_jobs_for_worker(w["url"])
    log.info("Added worker %s:%s", w["host"], w["port"])
    status = wrk.probe_worker(w["url"])
    return jsonify({"host": w["host"], "port": w["port"], "online": status is not None, "status": status}), 201


@app.delete("/api/workers/<int:port>/<path:host>")
def api_remove_worker(host, port):
    db.remove_worker(host, port)
    log.info("Removed worker %s:%s", host, port)
    return jsonify({"removed": True, "host": host, "port": port}), 200


@app.get("/api/capabilities")
def api_capabilities():
    combos = {
        (w.get("os"), w.get("cpu"))
        for w in wrk.probe_all()
        if w.get("os") and w.get("cpu")
    }
    return jsonify([{"os": o, "cpu": c} for o, c in sorted(combos)])


@app.post("/api/jobs")
def api_create_job():
    os_req = (request.form.get("os") or "").strip() or None
    cpu_req = (request.form.get("cpu") or "").strip() or None

    if request.files and "file" in request.files:
        upload = request.files["file"]
        if not upload.filename:
            return jsonify({"error": "No file selected"}), 400
        source_dir = os.path.join(settings.DATA_DIR, "sources")
        os.makedirs(source_dir, exist_ok=True)
        fname = os.path.basename(upload.filename)
        dest = os.path.join(source_dir, fname)
        upload.save(dest)
        job_id = db.create_job("upload", fname, os_req, cpu_req)
    else:
        data = request.get_json(silent=True) or {}
        repo_url = (data.get("repo_url") or request.form.get("repo_url") or "").strip()
        if not repo_url:
            return jsonify({"error": "Provide a file upload or a repo_url"}), 400
        job_id = db.create_job("git", repo_url, os_req, cpu_req)

    scheduler.start()
    log.info("Job %s created: type=%s os=%s cpu=%s", job_id, "upload" if "file" in request.files else "git", os_req, cpu_req)
    return jsonify({"id": job_id}), 201


@app.get("/api/jobs")
def api_list_jobs():
    return jsonify(db.list_jobs())


@app.get("/api/jobs/<job_id>")
def api_get_job(job_id):
    job = db.get_job(job_id)
    if not job:
        return jsonify({"error": "not found"}), 404
    return jsonify(job)


@app.get("/api/jobs/<job_id>/artifacts/<name>")
def api_download(job_id, name):
    job = db.get_job(job_id)
    if not job:
        return jsonify({"error": "not found"}), 404
    safe = os.path.basename(name)
    out_dir = os.path.join(settings.BUILDS_DIR, job_id)
    return send_from_directory(out_dir, safe, as_attachment=True)


@app.post("/api/jobs/<job_id>/cancel")
def api_cancel(job_id):
    job = db.get_job(job_id)
    if not job:
        return jsonify({"error": "not found"}), 404
    if job["status"] in ("queued",):
        db.update_job(job_id, status="cancelled", finished_at=datetime.datetime.now().isoformat())
        return jsonify({"id": job_id, "status": "cancelled"})
    return jsonify({"error": f"Cannot cancel job in state {job['status']}"}), 400


@app.get("/api/jobs/<job_id>/stream")
def api_stream(job_id):
    job = db.get_job(job_id)
    if not job:
        return jsonify({"error": "not found"}), 404

    q = scheduler.subscribe(job_id)

    def gen():
        # initial snapshot
        snap = db.get_job(job_id)
        if snap:
            yield f"event: job\ndata: {json.dumps({'status': snap['status']})}\n\n"
            for line in (snap.get("log") or "").splitlines():
                yield f"event: log\ndata: {json.dumps(line)}\n\n"
        try:
            while True:
                try:
                    frame = q.get(timeout=15)
                    yield frame
                except queue.Empty:
                    # heartbeat to keep connection alive
                    yield ": keepalive\n\n"
                job_now = db.get_job(job_id)
                if job_now and job_now["status"] in ("done", "failed", "cancelled"):
                    break
        finally:
            scheduler.unsubscribe(job_id, q)

    return Response(gen(), mimetype="text/event-stream")


@app.errorhandler(Exception)
def _handle_error(err):
    if isinstance(err, HTTPException):
        return err
    log.exception("Unhandled error serving %s %s", request.method, request.path)
    return jsonify({"error": f"{type(err).__name__}: {err}"}), 500


def main():
    _ensure_ready()
    scheduler.start()
    _configure_logging()
    log.info("Controller listening on 0.0.0.0:%s (%d workers configured)", settings.PORT, len(db.list_workers()))
    if settings.DEBUG:
        app.run(host="0.0.0.0", port=settings.PORT, threaded=True)
        return
    from waitress import serve
    serve(app, host="0.0.0.0", port=settings.PORT, threads=8)


def _configure_logging():
    level = logging.DEBUG if settings.DEBUG else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    # waitress's own request access log to stderr by default; keep ours separate.
    logging.getLogger("waitress").setLevel(logging.INFO)


if __name__ == "__main__":
    main()