docker-images/python-tools/app/web.py

"""Flask application."""

from __future__ import annotations

import ipaddress
import json
from datetime import datetime
from pathlib import Path

from flask import Flask, abort, redirect, render_template, request, send_file, url_for

from .config import Settings
from .discovery import TaskCatalog
from .inputs import validate_inputs, validate_uploads
from .runner import RunManager
from .scheduler import BackgroundScheduler
from .store import RunStore


def create_app(settings: Settings | None = None) -> Flask:
    settings = settings or Settings.from_environment()
    catalog = TaskCatalog(settings.task_root)
    store = RunStore(settings.database_path)
    store.initialize()
    store.mark_interrupted()
    manager = RunManager(settings, catalog, store)
    manager.cleanup_transient_files()
    scheduler = BackgroundScheduler(catalog, manager)
    scheduler.start()

    project_root = Path(__file__).resolve().parent.parent
    app = Flask(
        __name__,
        template_folder=str(project_root / "templates"),
        static_folder=str(project_root / "static"),
    )
    app.config["settings"] = settings
    app.extensions["catalog"] = catalog
    app.extensions["run_manager"] = manager
    app.extensions["run_store"] = store
    app.extensions["background_scheduler"] = scheduler

    @app.before_request
    def allow_only_proxy() -> None:
        if not settings.allowed_proxy_ips:
            return
        remote_address = request.remote_addr
        if remote_address is None:
            abort(403)
        try:
            normalized = str(ipaddress.ip_address(remote_address))
        except ValueError:
            abort(403)
        if normalized not in settings.allowed_proxy_ips:
            abort(403)

    @app.template_filter("timestamp")
    def timestamp(value: str | None) -> str:
        if not value:
            return "—"
        return datetime.fromisoformat(value).astimezone().strftime("%Y-%m-%d %H:%M:%S %Z")

    def render_index(
        *,
        form_values: dict[str, dict[str, list[str]]] | None = None,
        input_errors: dict[str, dict[str, str]] | None = None,
        status: int = 200,
    ):
        groups = catalog.groups()
        latest = {
            ("group", group.id): store.last_run("group", group.id) for group in groups
        }
        for group in groups:
            for task in group.tasks:
                latest[("task", task.id)] = store.last_run("task", task.id)
        return (
            render_template(
                "index.html",
                groups=groups,
                latest=latest,
                form_values=form_values or {},
                input_errors=input_errors or {},
            ),
            status,
        )

    @app.get("/")
    def index():
        return render_index()

    @app.post("/tasks/<path:task_id>/run")
    def run_task(task_id: str):
        task = catalog.task(task_id)
        if task is None:
            abort(404)
        submitted_values = {
            field.name: request.form.getlist(field.name) for field in task.inputs
        }
        input_data, errors = validate_inputs(task.inputs, submitted_values)
        uploads, upload_errors = validate_uploads(task.inputs, request.files)
        errors.update(upload_errors)
        if errors:
            if request.accept_mimetypes.best == "application/json":
                return {"errors": errors}, 400
            return render_index(
                form_values={
                    task.id: {
                        field.name: [] if field.sensitive else submitted_values[field.name]
                        for field in task.inputs
                    }
                },
                input_errors={task.id: errors},
                status=400,
            )
        run, _ = manager.run_task(task, input_data, uploads)
        if request.accept_mimetypes.best == "application/json":
            return {
                "artifacts_url": url_for("run_artifacts", run_id=run.id),
                "detail_url": url_for("run_detail", run_id=run.id),
                "output_url": url_for("run_output", run_id=run.id),
                "run_id": run.id,
                "state_url": url_for("run_state", run_id=run.id),
            }, 202
        return redirect(url_for("run_detail", run_id=run.id))

    @app.post("/groups/<path:group_id>/run")
    def run_group(group_id: str):
        group = catalog.group(group_id)
        if group is None or not group.can_run_as_group:
            abort(404)
        run, _ = manager.run_group(group)
        return redirect(url_for("run_detail", run_id=run.id))

    @app.get("/runs/<int:run_id>")
    def run_detail(run_id: int):
        run = store.get(run_id)
        if run is None:
            abort(404)
        execution_task = execution_replay_task(run)
        return render_template(
            "run.html",
            run=run,
            children=store.children(run.id),
            output=read_run_output(run),
            execution_task=execution_task,
        )

    def execution_replay_task(run):
        if run.status != "succeeded" or run.target_kind != "task":
            return None
        task = catalog.task(run.target_path)
        if task is None or task.execution is None:
            return None
        if any(field.type == "file" or field.sensitive for field in task.inputs):
            return None
        try:
            input_data = json.loads(run.input_json)
        except json.JSONDecodeError:
            return None
        if not isinstance(input_data, dict):
            return None
        if input_data.get(task.execution.field.name) != task.execution.dry_run_value:
            return None
        return task

    @app.post("/runs/<int:run_id>/execute")
    def execute_dry_run(run_id: int):
        run = store.get(run_id)
        if run is None:
            abort(404)
        task = execution_replay_task(run)
        if task is None:
            abort(404)
        input_data = json.loads(run.input_json)
        input_data[task.execution.field.name] = task.execution.execute_value
        raw_values = {
            name: [str(value) for value in value] if isinstance(value, list) else [str(value)]
            for name, value in input_data.items()
        }
        validated, errors = validate_inputs(task.inputs, raw_values)
        if errors:
            abort(409)
        executed, _ = manager.run_task(task, validated)
        return redirect(url_for("run_detail", run_id=executed.id))

    @app.post("/runs/<int:run_id>/cancel")
    def cancel_run(run_id: int):
        run = manager.request_cancel(run_id)
        if run is None:
            abort(404)
        return redirect(url_for("run_detail", run_id=run.id))

    @app.get("/runs/<int:run_id>/output")
    def run_output(run_id: int):
        run = store.get(run_id)
        if run is None:
            abort(404)
        return {"output": read_run_output(run)}

    def read_run_output(run) -> str:
        if not run.log_path:
            return ""
        try:
            return Path(run.log_path).read_text(encoding="utf-8", errors="replace")
        except FileNotFoundError:
            return ""

    @app.get("/runs/<int:run_id>/artifacts")
    def run_artifacts(run_id: int):
        run = store.get(run_id)
        if run is None:
            abort(404)
        artifacts = []
        for name in manager.artifacts(run.id):
            artifact = {
                "name": name,
                "url": url_for("download_artifact", run_id=run.id, filename=name),
            }
            if name.casefold().endswith(".json"):
                artifact["view_url"] = url_for("view_artifact", run_id=run.id, filename=name)
            artifacts.append(artifact)
        return {"artifacts": artifacts}

    @app.get("/runs/<int:run_id>/artifacts/<filename>/view")
    def view_artifact(run_id: int, filename: str):
        run = store.get(run_id)
        if run is None or run.status != "succeeded" or not filename.casefold().endswith(".json"):
            abort(404)
        artifact = manager.artifact_path(run.id, filename)
        if artifact is None:
            abort(404)
        return send_file(artifact, mimetype="application/json", as_attachment=False, conditional=False)

    @app.get("/runs/<int:run_id>/artifacts/<filename>")
    def download_artifact(run_id: int, filename: str):
        run = store.get(run_id)
        if run is None or run.status != "succeeded":
            abort(404)
        artifact = manager.artifact_path(run.id, filename)
        if artifact is None:
            abort(404)
        response = send_file(
            artifact,
            as_attachment=True,
            download_name=artifact.name,
            conditional=False,
        )
        original_response = response.response

        def delete_after_transfer():
            try:
                yield from original_response
            finally:
                close = getattr(original_response, "close", None)
                if close is not None:
                    close()
                manager.delete_artifact(run.id, artifact.name)

        response.response = delete_after_transfer()
        response.direct_passthrough = False
        return response

    @app.get("/runs/<int:run_id>/state")
    def run_state(run_id: int):
        run = store.get(run_id)
        if run is None:
            abort(404)
        return {
            "status": run.status,
            "finished_at": run.finished_at,
            "exit_code": run.exit_code,
            "cancel_requested_at": run.cancel_requested_at,
        }

    @app.get("/history/<target_kind>/<path:target_path>")
    def history(target_kind: str, target_path: str):
        if target_kind not in {"task", "group"}:
            abort(404)
        return render_template(
            "history.html",
            target_kind=target_kind,
            target_path=target_path,
            runs=store.history(target_kind, target_path),
        )

    return app