reconcile.py

"""App Automator — reconcile Databricks Apps to their scheduled on/off windows.

This script is run by the serverless job defined in resources/reconcile_job.yml.
It is intentionally simple and idempotent: every run it figures out, for each
app in config/apps.yml, whether the app *should* be running right now, compares
that to the app's actual compute state, and starts or stops it if they differ.

Because it's idempotent, a missed run (or running it twice) is harmless — the
next run just brings everything back in line.

You can also run it locally to test:
    python src/reconcile.py --config config/apps.yml --dry-run true -p <profile>
"""

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass, field
from datetime import datetime, time
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

import yaml
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.apps import ComputeState

# States that mean "compute is up or on its way up" -> treat app as ON.
_RUNNING_STATES = {ComputeState.ACTIVE, ComputeState.STARTING, ComputeState.UPDATING}
# States that mean "compute is down or on its way down" -> treat app as OFF.
_STOPPED_STATES = {ComputeState.STOPPED, ComputeState.STOPPING, ComputeState.ERROR}

_DAY_ALIASES = {
    "mon": 0, "monday": 0,
    "tue": 1, "tues": 1, "tuesday": 1,
    "wed": 2, "weds": 2, "wednesday": 2,
    "thu": 3, "thur": 3, "thurs": 3, "thursday": 3,
    "fri": 4, "friday": 4,
    "sat": 5, "saturday": 5,
    "sun": 6, "sunday": 6,
}
_DEFAULT_DAYS = [0, 1, 2, 3, 4]  # Mon–Fri


@dataclass
class AppSchedule:
    """One app's parsed schedule."""

    name: str
    on_time: time
    off_time: time
    tz: ZoneInfo
    days: list[int]
    enabled: bool = True
    # Populated during reconcile for the summary line.
    notes: list[str] = field(default_factory=list)

    def should_be_on(self, now: datetime) -> bool:
        """Return True if this app should be running at `now` (tz-aware)."""
        local = now.astimezone(self.tz)
        if self.on_time == self.off_time:
            # Degenerate window -> never on.
            return False

        overnight = self.off_time < self.on_time
        if overnight:
            # Window wraps past midnight, e.g. 22:00 -> 06:00.
            # It's "on" if we're after on-time today (evening portion) or
            # before off-time today (morning portion). The day-of-week gate
            # applies to the day the window STARTED (the evening).
            in_evening = local.time() >= self.on_time
            in_morning = local.time() < self.off_time
            if in_evening:
                return local.weekday() in self.days
            if in_morning:
                # Belongs to yesterday's window; check yesterday's weekday.
                yesterday = (local.weekday() - 1) % 7
                return yesterday in self.days
            return False

        # Normal same-day window, e.g. 08:00 -> 19:00.
        if local.weekday() not in self.days:
            return False
        return self.on_time <= local.time() < self.off_time


def _parse_time(value: str, app_name: str) -> time:
    try:
        hh, mm = str(value).strip().split(":")
        return time(int(hh), int(mm))
    except Exception as exc:  # noqa: BLE001 - surface a clear config error
        raise ValueError(
            f"App '{app_name}': invalid time {value!r}, expected 'HH:MM' (24h)."
        ) from exc


def _parse_days(value, app_name: str) -> list[int]:
    if value is None:
        return list(_DEFAULT_DAYS)
    if not isinstance(value, list):
        raise ValueError(f"App '{app_name}': 'days' must be a list, e.g. [mon, tue].")
    out: list[int] = []
    for d in value:
        key = str(d).strip().lower()
        if key not in _DAY_ALIASES:
            raise ValueError(
                f"App '{app_name}': unknown day {d!r}. Use mon/tue/wed/thu/fri/sat/sun."
            )
        out.append(_DAY_ALIASES[key])
    return sorted(set(out))


def _resolve_tz(name: str, app_name: str) -> ZoneInfo:
    try:
        return ZoneInfo(name)
    except ZoneInfoNotFoundError as exc:
        raise ValueError(
            f"App '{app_name}': unknown timezone {name!r}. Use an IANA name "
            f"like 'America/New_York'."
        ) from exc


def load_config(path: str) -> list[AppSchedule]:
    """Parse config/apps.yml into a list of AppSchedule objects."""
    with open(path, "r", encoding="utf-8") as fh:
        raw = yaml.safe_load(fh) or {}

    settings = raw.get("settings") or {}
    default_tz_name = settings.get("default_timezone", "UTC")

    apps_raw = raw.get("apps") or []
    if not isinstance(apps_raw, list):
        raise ValueError("Config 'apps' must be a list. See config/apps.yml comments.")

    schedules: list[AppSchedule] = []
    for entry in apps_raw:
        if not isinstance(entry, dict) or "name" not in entry:
            raise ValueError(f"Each app entry needs a 'name'. Got: {entry!r}")
        name = str(entry["name"]).strip()
        for required in ("start", "stop"):
            if required not in entry:
                raise ValueError(f"App '{name}': missing required field '{required}'.")
        tz_name = entry.get("timezone", default_tz_name)
        schedules.append(
            AppSchedule(
                name=name,
                on_time=_parse_time(entry["start"], name),
                off_time=_parse_time(entry["stop"], name),
                tz=_resolve_tz(tz_name, name),
                days=_parse_days(entry.get("days"), name),
                enabled=bool(entry.get("enabled", True)),
            )
        )
    return schedules


def _current_state(w: WorkspaceClient, app_name: str) -> ComputeState | None:
    """Return the app's current compute state, or None if it doesn't exist."""
    from databricks.sdk.errors import NotFound

    try:
        app = w.apps.get(name=app_name)
    except NotFound:
        return None
    if app.compute_status is None:
        return None
    return app.compute_status.state


def reconcile(config_path: str, dry_run: bool, profile: str | None = None) -> int:
    """Bring every managed app in line with its schedule. Returns exit code."""
    schedules = load_config(config_path)
    # In the job, auth comes from the job's identity; locally, from the profile.
    w = WorkspaceClient(profile=profile) if profile else WorkspaceClient()

    now = datetime.now(tz=ZoneInfo("UTC"))
    print(f"App Automator reconcile @ {now.isoformat()}  (dry_run={dry_run})")
    print(f"Managing {len(schedules)} app(s) from {config_path}\n")

    changed = 0
    errors = 0
    warnings = 0
    for sched in schedules:
        if not sched.enabled:
            print(f"  - {sched.name}: skipped (enabled: false)")
            continue

        want_on = sched.should_be_on(now)
        state = _current_state(w, sched.name)
        if state is None:
            # Missing app = likely a typo or an app not created yet. Warn and
            # keep going — don't fail the whole scheduled job over one entry.
            print(f"  ! {sched.name}: NOT FOUND in workspace — skipping "
                  f"(check the name in `databricks apps list`).")
            warnings += 1
            continue

        is_on = state in _RUNNING_STATES
        want_word = "ON" if want_on else "OFF"
        local = now.astimezone(sched.tz).strftime("%a %H:%M %Z")

        if want_on == is_on:
            print(f"  = {sched.name}: already {want_word} (state={state.value}, {local})")
            continue

        # A change is needed.
        action = "START" if want_on else "STOP"
        print(f"  > {sched.name}: {action} (state={state.value} -> want {want_word}, {local})")
        changed += 1
        if dry_run:
            continue
        try:
            if want_on:
                w.apps.start(name=sched.name)
            else:
                w.apps.stop(name=sched.name)
        except Exception as exc:  # noqa: BLE001 - keep going for other apps
            print(f"    ERROR {action} {sched.name}: {exc}")
            errors += 1

    verb = "would change" if dry_run else "changed"
    print(f"\nDone. {verb} {changed} app(s), {warnings} warning(s), "
          f"{errors} error(s).")
    # Only a genuine start/stop API failure fails the job. Missing apps are
    # warnings so a stray config entry doesn't turn every run red.
    return 1 if errors else 0


def _str2bool(value: str) -> bool:
    return str(value).strip().lower() in {"1", "true", "yes", "y", "on"}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Reconcile Databricks Apps to schedule.")
    parser.add_argument("--config", required=True, help="Path to apps.yml")
    parser.add_argument(
        "--dry-run", default="false", help="true = log only, don't start/stop"
    )
    parser.add_argument(
        "-p", "--profile", default=None, help="CLI profile (local runs only)"
    )
    args = parser.parse_args(argv)
    return reconcile(args.config, _str2bool(args.dry_run), args.profile)


if __name__ == "__main__":
    # Only exit non-zero on real errors. A bare sys.exit(0) is reported as a
    # task failure by the Databricks serverless Python task runner, so on
    # success we simply fall through and let the task terminate normally.
    _code = main()
    if _code:
        sys.exit(_code)