ha-babytracker/custom_components/babytracker/babytracker_api/seed.py

"""Read the seed database that ``GET /account/seed`` returns.

The seed is a zip holding the app's SQLite database. It is the full account
state at some point in time; transaction logs are replayed on top of it.
"""

from __future__ import annotations

import shutil
import sqlite3
import zipfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterator, Optional, Union

from .models import TYPE_FIELD, format_date, measure

SEED_DB_NAME = "EasyLog.db"

PathLike = Union[str, Path]


def extract_seed_db(zip_path: PathLike, dest_path: PathLike) -> Path:
    """Extract the database from a seed zip to ``dest_path``.

    Copies only the one expected member, so paths inside the archive never
    decide where anything is written.
    """
    dest = Path(dest_path)
    with zipfile.ZipFile(zip_path) as archive:
        matches = [name for name in archive.namelist() if Path(name).name == SEED_DB_NAME]
        if not matches:
            raise ValueError(f"{SEED_DB_NAME} not found in seed zip")
        with archive.open(matches[0]) as source, dest.open("wb") as target:
            shutil.copyfileobj(source, target)
    return dest


def open_seed_db(db_path: PathLike) -> sqlite3.Connection:
    """Open an extracted seed database read-only, with rows addressable by name."""
    connection = sqlite3.connect(f"{Path(db_path).resolve().as_uri()}?mode=ro", uri=True)
    connection.row_factory = sqlite3.Row
    return connection


def _date_field(seconds: Optional[float]) -> Optional[str]:
    # Database times are Unix seconds. The app treats values <= 0 as unset.
    if not seconds or seconds <= 0:
        return None
    return format_date(datetime.fromtimestamp(seconds, timezone.utc))


def read_babies(connection: sqlite3.Connection) -> list[dict[str, Any]]:
    """Return babies as records usable in ``baby=`` of the record builders."""
    babies = []
    for row in connection.execute("SELECT ID, Timestamp, Name, DOB, DueDay, Gender, Picture FROM Baby"):
        baby: dict[str, Any] = {
            "gender": bool(row["Gender"]),  # True is boy.
            "name": row["Name"],
            "newFlage": False,
            "objectID": row["ID"],
        }
        if row["Picture"]:
            baby["pictureName"] = row["Picture"]
        for field, value in (("dob", row["DOB"]), ("dueDay", row["DueDay"]), ("timestamp", row["Timestamp"])):
            formatted = _date_field(value)
            if formatted:
                baby[field] = formatted
        babies.append(baby)
    return babies


def read_cursors(connection: sqlite3.Connection) -> dict[str, int]:
    """Return per-device SyncIDs already merged into the seed.

    Passing these to :meth:`BabyTrackerClient.pull` replays only the
    transactions that are newer than the seed.
    """
    rows = connection.execute("SELECT DeviceID, SyncID FROM MergedTransaction")
    return {row["DeviceID"]: row["SyncID"] for row in rows}


# Database tables holding activities, and the record type each corresponds to.
_ACTIVITY_TABLES = {
    "Allergen": "Allergen",
    "Bath": "Bath",
    "Diaper": "Diaper",
    "DoctorVisit": "DoctorVisit",
    "Formula": "Formula",
    "Growth": "Growth",
    "HealthQuestion": "HealthQuestion",
    "Journal": "Journal",
    "Joy": "Joy",
    "Medicine": "Medication",
    "Milestone": "Milestone",
    "Nursing": "Nursing",
    "OtherActivity": "OtherActivity",
    "OtherFeed": "OtherFeed",
    "Pump": "Pump",
    "Pumped": "Pumped",
    "Sick": "Sick",
    "Sleep": "Sleep",
    "Temperature": "Temperature",
    "Vaccine": "Vaccine",
}


def _names(connection: sqlite3.Connection, table: str) -> dict[str, str]:
    try:
        return {row["ID"]: row["Name"] for row in connection.execute(f'SELECT ID, Name FROM "{table}"')}
    except sqlite3.OperationalError:
        return {}


def _optional_measure(value: Optional[float], english: Any) -> Optional[dict[str, Any]]:
    # Growth columns default to 0, and the app reads 0 as "not measured".
    return measure(value, bool(english)) if value else None


def _activity_fields(table: str, row: dict[str, Any], names: dict[str, dict[str, str]]) -> dict[str, Any]:
    """Convert a table's own columns to the fields its record type uses in the log."""
    if table == "Diaper":
        return {"status": row.get("Status"), "amount": row.get("Amount")}
    if table == "Sleep":
        return {"duration": row.get("Duration") or 0}
    if table in ("Formula", "Pumped", "Pump"):
        return {"amount": measure(row.get("Amount") or 0, bool(row.get("IsEnglishScale")))}
    if table == "Nursing":
        return {
            "finishSide": row.get("FinishSide"),
            "leftDuration": row.get("LeftDuration") or 0,
            "rightDuration": row.get("RightDuration") or 0,
            "bothDuration": row.get("BothDuration") or 0,
        }
    if table == "OtherFeed":
        return {
            "feedType": {"name": names["OtherFeedSelection"].get(row.get("TypeID"))},
            "amount": measure(row.get("Amount") or 0, bool(row.get("IsEnglishScale"))),
            "unit": row.get("Unit"),
        }
    if table == "Temperature":
        return {"temperature": measure(row.get("Temp") or 0, bool(row.get("IsEnglishScale")))}
    if table == "Growth":
        # Length and head share one scale column.
        return {
            "weight": _optional_measure(row.get("Weight"), row.get("IsEnglishWeightScale")),
            "length": _optional_measure(row.get("Length"), row.get("IsEnglishLengthScale")),
            "head": _optional_measure(row.get("Head"), row.get("IsEnglishLengthScale")),
        }
    if table == "Medicine":
        return {
            "medicationSelection": {"name": names["MedicineSelection"].get(row.get("MedID"))},
            "amount": row.get("Amount"),
        }
    if table == "Vaccine":
        return {"vaccineType": {"name": names["VaccineSelection"].get(row.get("VaccID"))}}
    if table == "Milestone":
        return {"milestoneType": {"name": names["MilestoneSelection"].get(row.get("MilestoneSelectionID"))}}
    if table == "OtherActivity":
        return {
            "desc": {"name": names["OtherActivityDesc"].get(row.get("DescID"))},
            "duration": row.get("Duration") or 0,
        }
    return {}


def read_records(connection: sqlite3.Connection) -> Iterator[dict[str, Any]]:
    """Yield the seed's activities, shaped like the records in transaction logs.

    Tables missing from an older or newer database are skipped.
    """
    names = {
        table: _names(connection, table)
        for table in ("MedicineSelection", "MilestoneSelection", "OtherActivityDesc", "OtherFeedSelection", "VaccineSelection")
    }
    for table, record_type in _ACTIVITY_TABLES.items():
        try:
            rows = connection.execute(f'SELECT * FROM "{table}"').fetchall()
        except sqlite3.OperationalError:
            continue
        for sqlite_row in rows:
            row = dict(sqlite_row)
            record: dict[str, Any] = {
                TYPE_FIELD: record_type,
                "objectID": row["ID"],
                "time": datetime.fromtimestamp(row["Time"], timezone.utc),
                "timestamp": _date_field(row.get("Timestamp")),
                "note": row.get("Note") or "",
            }
            if row.get("BabyID"):
                record["baby"] = {"objectID": row["BabyID"]}
            record.update(_activity_fields(table, row, names))
            yield record