clickstream-ch-kafka-superset-demo/generator/src/clickstream_generator/launch.py

"""Глаголы и профили запуска генератора."""

from __future__ import annotations

import argparse
import os
import shlex
import sys
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone


LAUNCH_ENV_KEYS = (
    "GEN_SEED",
    "GEN_MODEL_T0",
    "GEN_MODEL_TIMEZONE",
    "GEN_MODEL_TIME_SPEED",
    "GEN_TICK_SECONDS",
    "GEN_LAMBDA_BASE_PER_MIN",
    "GEN_JITTER_PCT",
    "GEN_MIN_EVENTS_PER_TICK",
    "GEN_MAX_EVENTS_PER_TICK",
)


@dataclass(frozen=True)
class LaunchProfile:
    """Именованный набор настроек запуска."""

    duration: str
    env: dict[str, str]


PROFILES = {
    "ci": LaunchProfile(
        duration="6h",
        env={
            "GEN_SEED": "4242",
            "GEN_MODEL_T0": "2026-01-01T00:00:00+00:00",
            "GEN_MODEL_TIMEZONE": "UTC",
            "GEN_MODEL_TIME_SPEED": "1",
            "GEN_TICK_SECONDS": "60",
            "GEN_LAMBDA_BASE_PER_MIN": "60",
            "GEN_JITTER_PCT": "0",
            "GEN_MIN_EVENTS_PER_TICK": "1",
            "GEN_MAX_EVENTS_PER_TICK": "1000",
        },
    ),
    "daily-wave": LaunchProfile(
        duration="3d",
        env={
            "GEN_SEED": "4242",
            "GEN_MODEL_T0": "2026-01-01T00:00:00+00:00",
            "GEN_MODEL_TIMEZONE": "UTC",
            "GEN_MODEL_TIME_SPEED": "60",
            "GEN_TICK_SECONDS": "1",
            "GEN_LAMBDA_BASE_PER_MIN": "60",
            "GEN_JITTER_PCT": "0",
            "GEN_MIN_EVENTS_PER_TICK": "1",
            "GEN_MAX_EVENTS_PER_TICK": "1000",
        },
    ),
}


def parse_duration(value: str) -> timedelta:
    """Разбирает короткую длительность: 30m, 6h, 2d."""
    if len(value) < 2:
        raise ValueError("Duration must look like 6h or 2d")
    amount_text = value[:-1]
    unit = value[-1]
    if not amount_text.isdigit():
        raise ValueError("Duration amount must be a positive integer")
    amount = int(amount_text)
    if amount <= 0:
        raise ValueError("Duration amount must be > 0")
    if unit == "s":
        return timedelta(seconds=amount)
    if unit == "m":
        return timedelta(minutes=amount)
    if unit == "h":
        return timedelta(hours=amount)
    if unit == "d":
        return timedelta(days=amount)
    raise ValueError("Duration unit must be one of: s, m, h, d")


def build_launch_env(
    verb: str,
    *,
    profile_name: str = "daily-wave",
    duration: str | None = None,
    overrides: dict[str, str] | None = None,
) -> dict[str, str]:
    """Возвращает старые env-переменные для нового глагола запуска."""
    if verb not in {"backfill", "continue", "next-day", "reset"}:
        raise ValueError("глагол должен быть backfill, continue, next-day или reset")
    if profile_name not in PROFILES:
        known = ", ".join(sorted(PROFILES))
        raise ValueError(f"Unknown launch profile {profile_name!r}; known: {known}")

    profile = PROFILES[profile_name]
    overrides = overrides or {}
    env = dict(profile.env)
    for key in LAUNCH_ENV_KEYS:
        if key in overrides and overrides[key] != "":
            env[key] = overrides[key]

    selected_duration = (
        duration
        or overrides.get("GEN_HISTORY_DURATION")
        or profile.duration
    )
    env["GEN_LAUNCH_VERB"] = verb
    env["GEN_LAUNCH_PROFILE"] = profile_name

    if verb == "backfill":
        env["GEN_RUN_MODE"] = "backfill"
        env["GEN_STATE_RESET"] = "true"
        env["GEN_HISTORY_DURATION"] = selected_duration
        env["GEN_MODEL_T_END"] = _model_t_end(env["GEN_MODEL_T0"], selected_duration)
    elif verb == "continue":
        env["GEN_RUN_MODE"] = "live"
        env["GEN_STATE_RESET"] = "false"
        env["GEN_HISTORY_DURATION"] = selected_duration
        env["GEN_MODEL_T_END"] = _model_t_end(env["GEN_MODEL_T0"], selected_duration)
    elif verb == "next-day":
        current_t_end = overrides.get("GEN_MODEL_T_END")
        if not current_t_end:
            raise ValueError("для next-day требуется GEN_MODEL_T_END")
        env["GEN_RUN_MODE"] = "next-day"
        env["GEN_STATE_RESET"] = "false"
        env["GEN_MODEL_T_END"] = current_t_end
    else:
        env["GEN_RUN_MODE"] = "live"
        env["GEN_STATE_RESET"] = "true"

    return env


def shell_assignments(env: dict[str, str]) -> str:
    """Печатает env в виде, пригодном для eval в Bash."""
    return "\n".join(
        f"export {key}={shlex.quote(value)}"
        for key, value in sorted(env.items())
    )


def _model_t_end(model_t0: str, duration: str) -> str:
    t0 = datetime.fromisoformat(model_t0.replace("Z", "+00:00"))
    if t0.tzinfo is None:
        raise ValueError("GEN_MODEL_T0 must include timezone")
    return (t0.astimezone(timezone.utc) + parse_duration(duration)).isoformat()


def _parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Глаголы запуска генератора")
    parser.add_argument(
        "verb",
        choices=("backfill", "continue", "next-day", "reset"),
    )
    parser.add_argument(
        "--profile",
        default=(
            os.getenv("PROFILE")
            or os.getenv("GEN_LAUNCH_PROFILE")
            or "daily-wave"
        ),
        choices=sorted(PROFILES),
    )
    parser.add_argument("--duration", default=os.getenv("GEN_HISTORY_DURATION"))
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    """CLI для Bash-скриптов запуска."""
    args = _parse_args(sys.argv[1:] if argv is None else argv)
    env = build_launch_env(
        args.verb,
        profile_name=args.profile,
        duration=args.duration,
        overrides=os.environ,
    )
    print(shell_assignments(env))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())