devplacepy/devplacepy/services/devii/tasks/guards.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

from datetime import datetime, timedelta
from typing import Any, Callable, NamedTuple, Optional

from .context import in_task_run
from .limits import create_quota, owner_is_admin, run_quota
from .schedule import (
    DEFAULT_MAX_RUNS,
    MAX_LIFETIME_DAYS,
    RECURRING_KINDS,
    Schedule,
    from_iso,
    to_iso,
)

REASON_NOT_A_USER = "tasks belong to a signed-in account"
REASON_EXPIRED = "task lifetime expired"
REASON_MAX_RUNS = "run limit reached"
REASON_FAILURES = "too many consecutive failures"
REASON_OWNER_IDLE = "owner has been inactive"
REASON_NESTED = "a task run may not schedule more work"
REASON_RUN_QUOTA = "daily run limit reached"
REASON_BUDGET = "automation budget reached"
REASON_CREATE_QUOTA = "daily task limit reached"

BudgetProbe = Callable[[str, str], bool]

FIELD_MAX_PER_OWNER = "devii_task_max_per_owner"
DEFAULT_MAX_PER_OWNER = 10
BUDGET_RETRY_HOURS = 1


class Deferral(NamedTuple):
    reason: str
    retry_at: datetime


class AutomationDenied(Exception):
    def __init__(self, reason: str, retry_at: Optional[datetime] = None) -> None:
        super().__init__(reason)
        self.reason = reason
        self.retry_at = retry_at


def automation_allowed(owner_kind: str, owner_id: str) -> bool:
    return owner_kind == "user" and bool(owner_id)


def nesting_allowed(owner_id: str) -> bool:
    return not in_task_run() or owner_is_admin(owner_id)


def max_active_per_owner() -> int:
    from devplacepy.database import get_int_setting

    return get_int_setting(FIELD_MAX_PER_OWNER, DEFAULT_MAX_PER_OWNER)


def task_columns(schedule: Schedule, reference: datetime) -> dict[str, Any]:
    columns = dict(schedule.columns())
    if schedule.kind in RECURRING_KINDS and columns.get("max_runs") is None:
        columns["max_runs"] = DEFAULT_MAX_RUNS
    columns["expires_at"] = to_iso(schedule.expiry(reference))
    return columns


def expiry_of(row: dict[str, Any]) -> Optional[datetime]:
    for column, offset in (("expires_at", None), ("created_at", MAX_LIFETIME_DAYS)):
        stamp = row.get(column)
        if not stamp:
            continue
        try:
            moment = from_iso(str(stamp)[:19])
        except ValueError:
            continue
        return moment if offset is None else moment + timedelta(days=offset)
    return None


def retire_reason(
    row: dict[str, Any], reference: datetime, max_failures: int = 0
) -> Optional[str]:
    owner_kind = str(row.get("owner_kind") or "")
    owner_id = str(row.get("owner_id") or "")
    if not automation_allowed(owner_kind, owner_id):
        return REASON_NOT_A_USER

    expiry = expiry_of(row)
    if expiry is not None and reference >= expiry:
        return REASON_EXPIRED

    max_runs = row.get("max_runs")
    if max_runs is not None and int(row.get("run_count") or 0) >= int(max_runs):
        return REASON_MAX_RUNS

    if max_failures > 0 and int(row.get("failure_count") or 0) >= max_failures:
        return REASON_FAILURES

    return None


def budget_deferral(
    row: dict[str, Any],
    reference: datetime,
    budget_exceeded: Optional[BudgetProbe] = None,
) -> Optional[Deferral]:
    if budget_exceeded is None:
        return None
    owner_kind = str(row.get("owner_kind") or "")
    owner_id = str(row.get("owner_id") or "")
    if not budget_exceeded(owner_kind, owner_id):
        return None
    return Deferral(REASON_BUDGET, reference + timedelta(hours=BUDGET_RETRY_HOURS))


def run_deferral(db: Any, row: dict[str, Any], reference: datetime) -> Deferral:
    owner_kind = str(row.get("owner_kind") or "")
    owner_id = str(row.get("owner_id") or "")
    quota = run_quota(db, owner_kind, owner_id, reference)
    retry_at = quota.free_at or reference + timedelta(hours=BUDGET_RETRY_HOURS)
    return Deferral(REASON_RUN_QUOTA, retry_at)


def creation_denial(
    db: Any, owner_kind: str, owner_id: str, reference: datetime
) -> Optional[AutomationDenied]:
    if not automation_allowed(owner_kind, owner_id):
        return AutomationDenied(REASON_NOT_A_USER)
    if not nesting_allowed(owner_id):
        return AutomationDenied(REASON_NESTED)
    quota = create_quota(db, owner_kind, owner_id, reference)
    if quota.exceeded:
        return AutomationDenied(REASON_CREATE_QUOTA, quota.free_at)
    return None