devplacepy/devplacepy/services/devii/workspace/controller.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

from typing import Any

from devplacepy.database import get_table, resolve_by_slug
from devplacepy.services.containers import store
from devplacepy.services.containers.workspace import (
    flags,
    provision,
    quota,
    tunnels,
)
from devplacepy.services.containers.workspace.provision import WorkspaceError
from devplacepy.utils import generate_uid


class WorkspaceController:
    def __init__(self, owner_kind: str, owner_id: str, admin: bool = False) -> None:
        self.owner_kind = owner_kind
        self.owner_id = owner_id
        self.admin = admin

    def _user(self) -> dict | None:
        if self.owner_kind != "user" or not self.owner_id:
            return None
        return get_table("users").find_one(uid=self.owner_id)

    def _user_by_name(self, username: str) -> dict | None:
        if not username:
            return None
        return get_table("users").find_one(username=username)

    def _project(self, slug: str) -> dict | None:
        return resolve_by_slug(get_table("projects"), slug)

    def _resolve(self, slug: str, owner_uid: str = "") -> dict:
        project = self._project(slug)
        if not project:
            raise WorkspaceError(f"project not found: {slug}")
        target = owner_uid or self.owner_id
        instance = provision.find_for_project(project["uid"], target)
        if not instance:
            raise WorkspaceError("no workspace exists for this project")
        return instance

    async def dispatch(self, name: str, args: dict[str, Any]) -> Any:
        handler = getattr(self, f"_{name}", None)
        if handler is None:
            return {"error": f"unknown workspace action: {name}"}
        try:
            result = handler(args)
            if hasattr(result, "__await__"):
                return await result
            return result
        except WorkspaceError as error:
            return {"error": str(error)}

    async def _workspace_open(self, args: dict) -> Any:
        user = self._user()
        if not user:
            raise WorkspaceError("sign in to use workspaces")
        project = self._project(args.get("project_slug", ""))
        if not project:
            raise WorkspaceError("project not found")
        instance = await provision.ensure(project, user)
        instance = provision.resume(instance)
        provision.write_manifest(instance)
        view = provision.view(instance)
        view["editor_url"] = (
            f"/projects/{project.get('slug') or project['uid']}/workspace"
        )
        return view

    def _workspace_status(self, args: dict) -> Any:
        return provision.view(self._resolve(args.get("project_slug", "")))

    def _workspace_list(self, args: dict) -> Any:
        table = get_table("instances")
        filters: dict[str, Any] = {"is_workspace": 1, "deleted_at": None}
        if not self.admin:
            filters["workspace_owner_uid"] = self.owner_id
        return {
            "workspaces": [provision.view(row) for row in table.find(**filters)]
        }

    def _workspace_stop(self, args: dict) -> Any:
        instance = self._resolve(args.get("project_slug", ""))
        provision.stop(instance)
        return {"ok": True, "status": "stopping", "uid": instance["uid"]}

    def _workspace_delete(self, args: dict) -> Any:
        instance = self._resolve(args.get("project_slug", ""))
        for row in tunnels.list_for_instance(instance["uid"]):
            tunnels.soft_delete(row["uid"], self.owner_id)
        store.delete_instance(instance["uid"], self.owner_id)
        return {"ok": True, "deleted": instance["uid"]}

    def _tunnel_list(self, args: dict) -> Any:
        instance = self._resolve(args.get("project_slug", ""))
        return {"tunnels": tunnels.list_for_instance(instance["uid"])}

    def _tunnel_create(self, args: dict) -> Any:
        instance = self._resolve(args.get("project_slug", ""))
        port = int(args.get("container_port") or 0)
        if port <= 0:
            raise WorkspaceError("container_port is required")
        limits = quota.resolve(instance.get("workspace_owner_uid", ""), instance)
        if limits.max_tunnels and tunnels.count_for_instance(
            instance["uid"]
        ) >= limits.max_tunnels:
            raise WorkspaceError(f"tunnel limit reached ({limits.max_tunnels})")
        row = tunnels.create(
            instance, args.get("label", ""), port, self.owner_id
        )
        if not row:
            raise WorkspaceError("could not create tunnel")
        provision.write_manifest(instance)
        return {
            "ok": True,
            "tunnel": row,
            "url": f"https://{row['hostname']}",
            "public": True,
        }

    def _tunnel_delete(self, args: dict) -> Any:
        instance = self._resolve(args.get("project_slug", ""))
        uid = args.get("tunnel_uid", "")
        row = tunnels.get(uid)
        if not row or row.get("instance_uid") != instance["uid"]:
            raise WorkspaceError("tunnel not found on this workspace")
        tunnels.soft_delete(uid, self.owner_id)
        provision.write_manifest(instance)
        return {"ok": True, "deleted": uid}

    def _workspace_quota_get(self, args: dict) -> Any:
        target = self.owner_id
        username = args.get("username", "")
        if username:
            if not self.admin:
                raise WorkspaceError("only administrators may read another user's quota")
            other = self._user_by_name(username)
            if not other:
                raise WorkspaceError(f"user not found: {username}")
            target = other["uid"]
        limits = quota.resolve(target)
        return {
            "max_workspaces": limits.max_workspaces,
            "max_tunnels": limits.max_tunnels,
            "disk_quota_mb": limits.disk_quota_mb,
            "egress_quota_mb": limits.egress_quota_mb,
            "idle_stop_minutes": limits.idle_stop_minutes,
            "retention_days": limits.retention_days,
            "used_workspaces": provision.count_for_owner(target),
        }

    def _workspace_quota_set(self, args: dict) -> Any:
        if not self.admin:
            raise WorkspaceError("administrators only")
        other = self._user_by_name(args.get("username", ""))
        if not other:
            raise WorkspaceError("user not found")
        table = get_table(quota.RULES_TABLE)
        existing = table.find_one(
            owner_kind="user", owner_id=other["uid"], deleted_at=None
        )
        payload = {
            key: int(args.get(key) or 0)
            for key in quota.RULE_COLUMNS
            if args.get(key) is not None
        }
        if existing:
            table.update({"uid": existing["uid"], **payload}, ["uid"])
            uid = existing["uid"]
        else:
            uid = generate_uid()
            table.insert(
                {
                    "uid": uid,
                    "owner_kind": "user",
                    "owner_id": other["uid"],
                    "label": args.get("label", "") or other.get("username", ""),
                    "created_at": "",
                    "updated_at": "",
                    "deleted_at": None,
                    "deleted_by": None,
                    **payload,
                }
            )
        return {"ok": True, "rule_uid": uid, "applied": payload}

    def _workspace_flag_list(self, args: dict) -> Any:
        status = args.get("status", "open")
        if self.admin:
            return {"flags": flags.list_flags(status=status)}
        return {"flags": flags.list_flags(user_uid=self.owner_id, status=status)}

    def _workspace_flag_raise(self, args: dict) -> Any:
        if not self.admin:
            raise WorkspaceError("administrators only")
        owner = self._user_by_name(args.get("username", ""))
        if not owner:
            raise WorkspaceError("user not found")
        instance = self._resolve(args.get("project_slug", ""), owner["uid"])
        row = flags.raise_flag(
            instance,
            flags.KIND_MANUAL,
            args.get("severity", "warn"),
            args.get("detail", ""),
        )
        return {"ok": True, "flag": row}

    def _workspace_flag_resolve(self, args: dict) -> Any:
        if not self.admin:
            raise WorkspaceError("administrators only")
        status = args.get("status", "resolved")
        if not flags.set_status(args.get("flag_uid", ""), status, self.owner_id):
            raise WorkspaceError("flag not found or invalid status")
        return {"ok": True, "status": status}

    def _workspace_suspend(self, args: dict) -> Any:
        if not self.admin:
            raise WorkspaceError("administrators only")
        owner = self._user_by_name(args.get("username", ""))
        if not owner:
            raise WorkspaceError("user not found")
        instance = self._resolve(args.get("project_slug", ""), owner["uid"])
        reason = args.get("reason", "")
        if not reason:
            raise WorkspaceError("a reason is required and is shown to the owner")
        provision.suspend(instance, self.owner_id, reason)
        return {"ok": True, "suspended": instance["uid"], "reason": reason}

    def _workspace_unsuspend(self, args: dict) -> Any:
        if not self.admin:
            raise WorkspaceError("administrators only")
        owner = self._user_by_name(args.get("username", ""))
        if not owner:
            raise WorkspaceError("user not found")
        instance = self._resolve(args.get("project_slug", ""), owner["uid"])
        provision.unsuspend(instance)
        return {"ok": True, "resumed": instance["uid"]}