devplacepy/devplacepy/services/devii/tasks/limits.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

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

from .schedule import from_iso, to_iso

logger = logging.getLogger("devii.tasks.limits")

RUNS_TABLE = "devii_task_runs"
TASKS_TABLE = "devii_tasks"
WINDOW_HOURS = 24

FIELD_MEMBER_CREATE = "devii_task_member_create_24h"
FIELD_MEMBER_RUNS = "devii_task_member_runs_24h"
FIELD_ADMIN_CREATE = "devii_task_admin_create_24h"
FIELD_ADMIN_RUNS = "devii_task_admin_runs_24h"

DEFAULT_MEMBER_CREATE = 5
DEFAULT_MEMBER_RUNS = 10
DEFAULT_ADMIN_CREATE = 5
DEFAULT_ADMIN_RUNS = 100

RUNS_SQL = (
    f"SELECT created_at FROM {RUNS_TABLE} WHERE owner_kind = :kind AND owner_id = :owner "
    "AND created_at >= :cutoff ORDER BY created_at LIMIT :cap"
)
CREATIONS_SQL = (
    f"SELECT created_at FROM {TASKS_TABLE} WHERE owner_kind = :kind AND owner_id = :owner "
    "AND created_at >= :cutoff ORDER BY created_at LIMIT :cap"
)
RESERVE_SQL = (
    f"INSERT INTO {RUNS_TABLE} (uid, owner_kind, owner_id, task_uid, created_at) "
    "SELECT :uid, :kind, :owner, :task, :now WHERE ("
    f"SELECT COUNT(*) FROM {RUNS_TABLE} WHERE owner_kind = :kind AND owner_id = :owner "
    "AND created_at >= :cutoff) < :limit"
)
SAMPLE_CAP = 1000


class Quota(NamedTuple):
    used: int
    limit: int
    free_at: Optional[datetime]

    @property
    def exceeded(self) -> bool:
        return self.limit > 0 and self.used >= self.limit

    @property
    def remaining(self) -> int:
        if self.limit <= 0:
            return -1
        return max(0, self.limit - self.used)


def owner_is_admin(owner_id: str) -> bool:
    from devplacepy.database import get_admin_uids

    return bool(owner_id) and owner_id in get_admin_uids()


def _setting(name: str, default: int) -> int:
    from devplacepy.database import get_int_setting

    return get_int_setting(name, default)


def create_limit(is_admin: bool) -> int:
    if is_admin:
        return _setting(FIELD_ADMIN_CREATE, DEFAULT_ADMIN_CREATE)
    return _setting(FIELD_MEMBER_CREATE, DEFAULT_MEMBER_CREATE)


def run_limit(is_admin: bool) -> int:
    if is_admin:
        return _setting(FIELD_ADMIN_RUNS, DEFAULT_ADMIN_RUNS)
    return _setting(FIELD_MEMBER_RUNS, DEFAULT_MEMBER_RUNS)


def window_start(reference: datetime) -> datetime:
    return reference - timedelta(hours=WINDOW_HOURS)


def _stamps(
    db: Any, sql: str, table: str, owner_kind: str, owner_id: str, cutoff: str
) -> list[str]:
    if table not in db.tables:
        return []
    rows = db.query(sql, kind=owner_kind, owner=owner_id, cutoff=cutoff, cap=SAMPLE_CAP)
    return [str(row["created_at"]) for row in rows if row.get("created_at")]


def _quota(stamps: list[str], limit: int) -> Quota:
    used = len(stamps)
    if limit <= 0 or used < limit:
        return Quota(used, limit, None)
    oldest_kept = stamps[used - limit]
    try:
        free_at = from_iso(oldest_kept[:19]) + timedelta(hours=WINDOW_HOURS)
    except ValueError:
        free_at = None
    return Quota(used, limit, free_at)


def run_quota(db: Any, owner_kind: str, owner_id: str, reference: datetime) -> Quota:
    limit = run_limit(owner_is_admin(owner_id))
    stamps = _stamps(
        db, RUNS_SQL, RUNS_TABLE, owner_kind, owner_id, to_iso(window_start(reference))
    )
    return _quota(stamps, limit)


def create_quota(db: Any, owner_kind: str, owner_id: str, reference: datetime) -> Quota:
    limit = create_limit(owner_is_admin(owner_id))
    stamps = _stamps(
        db,
        CREATIONS_SQL,
        TASKS_TABLE,
        owner_kind,
        owner_id,
        to_iso(window_start(reference)),
    )
    return _quota(stamps, limit)


def record_run(
    db: Any, owner_kind: str, owner_id: str, task_uid: str, reference: datetime
) -> None:
    from devplacepy.utils import generate_uid

    db[RUNS_TABLE].insert(
        {
            "uid": generate_uid(),
            "owner_kind": owner_kind,
            "owner_id": owner_id,
            "task_uid": task_uid,
            "created_at": to_iso(reference),
        }
    )


def reserve_run(
    db: Any, owner_kind: str, owner_id: str, task_uid: str, reference: datetime
) -> bool:
    import sqlalchemy

    from devplacepy.utils import generate_uid

    limit = run_limit(owner_is_admin(owner_id))
    if limit <= 0 or RUNS_TABLE not in db.tables:
        record_run(db, owner_kind, owner_id, task_uid, reference)
        return True
    params = {
        "uid": generate_uid(),
        "kind": owner_kind,
        "owner": owner_id,
        "task": task_uid,
        "now": to_iso(reference),
        "cutoff": to_iso(window_start(reference)),
        "limit": limit,
    }
    with db:
        result = db.executable.execute(sqlalchemy.text(RESERVE_SQL), params)
    return result.rowcount == 1


def insert_task_within_quota(
    db: Any,
    record: dict[str, Any],
    owner_kind: str,
    owner_id: str,
    reference: datetime,
) -> bool:
    import sqlalchemy

    limit = create_limit(owner_is_admin(owner_id))
    if limit <= 0 or TASKS_TABLE not in db.tables:
        db[TASKS_TABLE].insert(record)
        return True

    columns = list(record)
    statement = sqlalchemy.text(
        f"INSERT INTO {TASKS_TABLE} ({', '.join(columns)}) "
        f"SELECT {', '.join(':' + name for name in columns)} WHERE ("
        f"SELECT COUNT(*) FROM {TASKS_TABLE} WHERE owner_kind = :q_kind "
        "AND owner_id = :q_owner AND created_at >= :q_cutoff) < :q_limit"
    )
    params = dict(record)
    params.update(
        {
            "q_kind": owner_kind,
            "q_owner": owner_id,
            "q_cutoff": to_iso(window_start(reference)),
            "q_limit": limit,
        }
    )
    with db:
        result = db.executable.execute(statement, params)
    return result.rowcount == 1


def prune_runs(db: Any, reference: datetime, keep_hours: int = WINDOW_HOURS * 2) -> int:
    if RUNS_TABLE not in db.tables:
        return 0
    cutoff = to_iso(reference - timedelta(hours=keep_hours))
    table = db[RUNS_TABLE]
    stale = table.count(created_at={"<": cutoff})
    if stale:
        table.delete(created_at={"<": cutoff})
    return stale