FabledSteward/plugins/homeassistant/updates.py

# plugins/homeassistant/updates.py
"""Installing one update: take a backup, confirm it landed, install, and follow
it to its end. Its row holds where it has got to, so Steward restarting, or Home
Assistant restarting as an update does, only delays the next step."""
from __future__ import annotations

import logging
from datetime import datetime, timezone

from sqlalchemy import select
from sqlalchemy.exc import IntegrityError

from . import client
from .models import HomeAssistantConnection, HomeAssistantUpdate
from .service import BACKUP_EVENT, LAST_BACKUP, is_newer, parse_moment
from .timeline import update_entry

logger = logging.getLogger(__name__)

UNDER_WAY = ("backing_up", "installing")
BACKUP_DEADLINE_S = 30 * 60
# An OS update reboots the machine, and a Pi is slow to come back.
INSTALL_DEADLINE_S = {"addon": 20 * 60, "core": 30 * 60, "os": 45 * 60}


class CannotStart(Exception):
    """Why an update wasn't started, said to the person."""


async def _mark_refused(app, connection_id: str, said: str) -> None:
    async with app.db_sessionmaker() as session:
        async with session.begin():
            row = await session.get(HomeAssistantConnection, connection_id)
            if row is not None:
                row.refused, row.last_error = True, said


async def under_way(session) -> HomeAssistantUpdate | None:
    return (await session.execute(select(HomeAssistantUpdate).where(
        HomeAssistantUpdate.state.in_(UNDER_WAY)))).scalars().first()


async def start_update(app, entity_id: str, username: str, transport=None) -> HomeAssistantUpdate:
    """Start installing `entity_id`, an explicit action that may contact Home
    Assistant. Raises CannotStart, with nothing started, unless the update is
    still waiting and the signed-in user may install it."""
    async with app.db_sessionmaker() as session:
        connection = (await session.execute(select(HomeAssistantConnection))).scalars().first()
        if connection is None:
            raise CannotStart("Home Assistant isn't connected")
        if connection.refused:
            raise CannotStart("Home Assistant refused Steward's sign-in")
        if await under_way(session) is not None:
            raise CannotStart("another update is under way")
        if not connection.can_install:
            raise CannotStart("the Home Assistant user who signed in isn't an administrator")
    url, verify = connection.url, connection.verify_tls
    known = next((u for u in connection.updates or [] if u["entity_id"] == entity_id), None)
    if known is None or not known["kind"]:
        raise CannotStart("that isn't a Core, OS or add-on update")
    try:
        token = await client.access_token(connection, transport)
        current = await client.read_state(url, entity_id, token, verify, transport)
        backup = await client.read_state(url, LAST_BACKUP, token, verify, transport)
        event = await client.read_state(url, BACKUP_EVENT, token, verify, transport)
    except client.Refused as exc:
        await _mark_refused(app, connection.id, str(exc))
        raise CannotStart(str(exc).rstrip(".")) from None
    except client.Unreachable as exc:
        raise CannotStart(str(exc).rstrip(".")) from None
    if backup is None:
        raise CannotStart("Home Assistant has no automatic backups to confirm one with")
    if current is None or current.get("state") != "on":
        raise CannotStart("that update is no longer waiting")
    attrs = current.get("attributes") or {}
    run = HomeAssistantUpdate(
        entity_id=entity_id, title=known["title"], kind=known["kind"],
        from_version=attrs.get("installed_version"), to_version=attrs.get("latest_version"),
        started_by=username, backup_before=backup.get("state"),
        backup_event_before=(event or {}).get("state"))
    async with app.db_sessionmaker() as session:
        try:
            async with session.begin():
                session.add(run)
                await session.flush()
                run_id = run.id
        except IntegrityError:
            raise CannotStart("another update is under way") from None
    try:
        await client.call_action(url, "backup", "create_automatic", {}, token, verify, transport)
    except client.Refused as exc:
        await _mark_refused(app, connection.id, str(exc))
        await _finish(app, run_id, "failed", f"The backup couldn't be started: {exc} Nothing was installed.")
        raise CannotStart(str(exc).rstrip(".")) from None
    except (client.Unreachable, client.Rejected) as exc:
        await _finish(app, run_id, "failed", f"The backup couldn't be started: {exc} Nothing was installed.")
        raise CannotStart(str(exc).rstrip(".")) from None
    async with app.db_sessionmaker() as session:
        return await session.get(HomeAssistantUpdate, run_id)


async def _finish(app, run_id: str, state: str, message: str | None) -> None:
    async with app.db_sessionmaker() as session:
        async with session.begin():
            run = await session.get(HomeAssistantUpdate, run_id)
            run.state, run.message, run.finished_at = state, message, datetime.now(timezone.utc)
            # The page should show the new versions now, not at the next poll.
            connection = (await session.execute(select(HomeAssistantConnection))).scalars().first()
            if connection is not None:
                connection.last_polled_at = None
            try:
                async with session.begin_nested():
                    await update_entry(session, run)
            except Exception:
                logger.exception("Home Assistant update %s: its timeline entry could not be written", run_id)


async def advance(app, transport=None) -> None:
    """Move the update under way one step on, if there is one."""
    async with app.db_sessionmaker() as session:
        run = await under_way(session)
        connection = (await session.execute(select(HomeAssistantConnection))).scalars().first()
    if run is None:
        return
    if connection is None:
        await _finish(app, run.id, "failed", "The Home Assistant connection was removed.")
        return
    try:
        if run.state == "backing_up":
            await _advance_backup(app, run, connection, transport)
        else:
            await _advance_install(app, run, connection, transport)
    except client.Refused as exc:
        await _mark_refused(app, connection.id, str(exc))
        await _finish(app, run.id, "failed", str(exc))
    except client.Rejected as exc:
        await _finish(app, run.id, "failed", f"Home Assistant wouldn't install it: {exc}")
    except client.Unreachable as exc:
        # Expected while Home Assistant restarts: the next tick asks again.
        logger.info("Home Assistant update %s: %s", run.id, exc)
        await _past_deadline(app, run)


def _elapsed(since: datetime) -> float:
    return (datetime.now(timezone.utc) - since).total_seconds()


async def _past_deadline(app, run) -> bool:
    if run.state == "backing_up" and _elapsed(run.started_at) > BACKUP_DEADLINE_S:
        await _finish(app, run.id, "failed",
                      f"No new backup within {BACKUP_DEADLINE_S // 60} minutes. Nothing was installed.")
        return True
    if run.state == "installing":
        limit = INSTALL_DEADLINE_S.get(run.kind, INSTALL_DEADLINE_S["os"])
        if _elapsed(run.install_sent_at or run.backup_at or run.started_at) > limit:
            await _finish(app, run.id, "failed",
                          f"Still not on {run.to_version} after {limit // 60} minutes. Look in Home Assistant.")
            return True
    return False


async def _advance_backup(app, run, connection, transport) -> None:
    token = await client.access_token(connection, transport)
    url, verify = connection.url, connection.verify_tls
    backup = await client.read_state(url, LAST_BACKUP, token, verify, transport)
    value = (backup or {}).get("state")
    if is_newer(value, run.backup_before):
        async with app.db_sessionmaker() as session:
            async with session.begin():
                row = await session.get(HomeAssistantUpdate, run.id)
                row.state, row.backup_at = "installing", parse_moment(value)
        return
    event = await client.read_state(url, BACKUP_EVENT, token, verify, transport)
    if event and event.get("state") != run.backup_event_before \
            and (event.get("attributes") or {}).get("event_type") == "failed":
        reason = (event.get("attributes") or {}).get("failed_reason")
        await _finish(app, run.id, "failed",
                      f"The backup failed{f' ({reason})' if reason else ''}. Nothing was installed.")
        return
    await _past_deadline(app, run)


async def _mark_sent(app, run_id: str, at: datetime | None) -> None:
    async with app.db_sessionmaker() as session:
        async with session.begin():
            row = await session.get(HomeAssistantUpdate, run_id)
            row.install_sent_at = at


async def _advance_install(app, run, connection, transport) -> None:
    token = await client.access_token(connection, transport)
    url, verify = connection.url, connection.verify_tls
    if run.install_sent_at is None:
        # Marked first: Steward stopping here leaves an install that runs out its
        # deadline, where asking twice would fail one that is going fine.
        await _mark_sent(app, run.id, datetime.now(timezone.utc))
        try:
            await client.call_action(url, "update", "install", {"entity_id": run.entity_id}, token, verify, transport)
        except client.Unreachable:
            await _mark_sent(app, run.id, None)
            raise
        return
    current = await client.read_state(url, run.entity_id, token, verify, transport)
    attrs = (current or {}).get("attributes") or {}
    installed = attrs.get("installed_version")
    if installed and (installed == run.to_version
                      or (current.get("state") == "off" and installed != run.from_version)):
        await _finish(app, run.id, "done", None)
        return
    await _past_deadline(app, run)