devplacepy/devplacepy/push/store.py

# retoor <retoor@molodetz.nl>

import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

from devplacepy.database import db, get_table
from devplacepy.push.providers import DEFAULT_PROVIDER
from devplacepy.utils import generate_uid

logger = logging.getLogger(__name__)

TABLE = "push_registration"


@dataclass(frozen=True)
class RegistrationWrite:
    record: dict[str, Any]
    created: bool
    revived: bool

    @property
    def probe(self) -> bool:
        return self.created or self.revived


def table():
    return get_table(TABLE)


def provider_of(registration: dict[str, Any]) -> str:
    return registration.get("provider") or DEFAULT_PROVIDER


def active_for_user(user_uid: str) -> list[dict[str, Any]]:
    return list(table().find(user_uid=user_uid, deleted_at=None))


def _filled(fields: dict[str, Any]) -> dict[str, Any]:
    filled: dict[str, Any] = {}
    for key, value in fields.items():
        if value is None:
            continue
        if isinstance(value, str):
            value = value.strip()
            if not value:
                continue
        filled[key] = value
    return filled


def _prefer_live(rows: list[dict[str, Any]]) -> dict[str, Any] | None:
    if not rows:
        return None
    live = [row for row in rows if not row.get("deleted_at")]
    return (live or rows)[-1]


def _lookup(user_uid: str, provider: str, **identity: Any) -> dict[str, Any] | None:
    return _prefer_live(
        list(table().find(user_uid=user_uid, provider=provider, **identity))
    )


def _with_id(record: dict[str, Any]) -> dict[str, Any]:
    if record.get("id"):
        return record
    found = table().find_one(uid=record.get("uid"))
    return found or record


def _retire_duplicates(
    user_uid: str,
    provider: str,
    keep_id: int,
    fields: dict[str, Any],
    previous_token: str | None = None,
) -> None:
    token = fields.get("token")
    client_id = fields.get("client_id")
    for row in table().find(user_uid=user_uid, provider=provider, deleted_at=None):
        if row["id"] == keep_id:
            continue
        if client_id and row.get("client_id") == client_id:
            mark_dead(row["id"])
            continue
        if token and row.get("token") == token:
            mark_dead(row["id"])
            continue
        if previous_token and row.get("token") == previous_token:
            mark_dead(row["id"])


def _merge(
    existing: dict[str, Any], fields: dict[str, Any]
) -> RegistrationWrite:
    revived = bool(existing.get("deleted_at"))
    patch: dict[str, Any] = {"id": existing["id"]}
    if revived:
        patch["deleted_at"] = None
    changed = revived
    for key, value in fields.items():
        if existing.get(key) != value:
            patch[key] = value
            changed = True
    if not changed:
        return RegistrationWrite(_with_id(existing), False, False)
    patch["registered_at"] = datetime.now(timezone.utc).isoformat()
    previous_token = existing.get("token")
    table().update(patch, ["id"])
    merged = {**existing, **patch}
    if revived:
        merged["deleted_at"] = None
    merged = _with_id(merged)
    _retire_duplicates(
        existing["user_uid"],
        existing["provider"],
        merged["id"],
        fields,
        previous_token=previous_token if previous_token != fields.get("token") else None,
    )
    logger.info(
        "Updated %s push subscription for user %s%s",
        existing.get("provider"),
        existing.get("user_uid"),
        " (revived)" if revived else "",
    )
    return RegistrationWrite(merged, False, revived)


def register(
    user_uid: str, provider: str, fields: dict[str, Any]
) -> RegistrationWrite:
    fields = _filled(fields)
    registrations = table()
    client_id = fields.get("client_id")
    token = fields.get("token")
    endpoint = fields.get("endpoint")

    if client_id:
        existing = _lookup(user_uid, provider, client_id=client_id)
        if existing:
            return _merge(existing, fields)

    if token:
        existing = _lookup(user_uid, provider, token=token)
        if existing:
            return _merge(existing, fields)

    if endpoint:
        existing = _lookup(user_uid, provider, endpoint=endpoint)
        if existing:
            return _merge(existing, fields)

    live = registrations.find_one(
        user_uid=user_uid, provider=provider, deleted_at=None, **fields
    )
    if live:
        return RegistrationWrite(_with_id(live), False, False)

    now = datetime.now(timezone.utc).isoformat()
    record = {
        "uid": generate_uid(),
        "user_uid": user_uid,
        "provider": provider,
        "created_at": now,
        "registered_at": now,
        "deleted_at": None,
        **fields,
    }
    inserted = registrations.insert(record)
    if isinstance(inserted, int):
        record["id"] = inserted
    record = _with_id(record)
    if record.get("id"):
        _retire_duplicates(user_uid, provider, record["id"], fields)
    logger.info("Registered %s push subscription for user %s", provider, user_uid)
    return RegistrationWrite(record, True, False)


def unregister(user_uid: str, provider: str, fields: dict[str, Any]) -> bool:
    fields = _filled(fields)
    client_id = fields.get("client_id")
    token = fields.get("token")
    endpoint = fields.get("endpoint")

    row = None
    if client_id:
        row = _lookup(user_uid, provider, client_id=client_id)
    if row is None and token:
        row = _lookup(user_uid, provider, token=token)
    if row is None and endpoint:
        row = _lookup(user_uid, provider, endpoint=endpoint)
    if row is None or row.get("deleted_at"):
        return False

    table().update(
        {"id": row["id"], "deleted_at": datetime.now(timezone.utc).isoformat()},
        ["id"],
    )
    logger.info("Unregistered %s push subscription for user %s", provider, user_uid)
    return True


def mark_dead(registration_id: int, dead_before: str | None = None) -> None:
    if dead_before is not None:
        row = table().find_one(id=registration_id)
        registered_at = row.get("registered_at") if row else None
        if registered_at and registered_at > dead_before:
            logger.info(
                "Skipped marking push subscription id=%s dead: registered again at %s "
                "after the provider confirmed it dead at %s",
                registration_id,
                registered_at,
                dead_before,
            )
            return
    table().update(
        {"id": registration_id, "deleted_at": datetime.now(timezone.utc).isoformat()},
        ["id"],
    )
    logger.info("Soft-deleted dead push subscription id=%s", registration_id)


def prune(cutoff: str) -> int:
    if TABLE not in db.tables:
        return 0
    rows = list(table().find(deleted_at={"<": cutoff}))
    if not rows:
        return 0
    table().delete(deleted_at={"<": cutoff})
    return len(rows)


def counts() -> dict[str, int]:
    if TABLE not in db.tables:
        return {}
    totals: dict[str, int] = {"dead": 0}
    for row in db.query(
        f"SELECT provider AS provider, deleted_at IS NULL AS live, COUNT(*) AS total "
        f"FROM {TABLE} GROUP BY provider, deleted_at IS NULL"
    ):
        provider = row["provider"] or DEFAULT_PROVIDER
        if row["live"]:
            totals[provider] = totals.get(provider, 0) + int(row["total"])
        else:
            totals["dead"] += int(row["total"])
    return totals