[03] Pipeline state as a snapshot blob; DB becomes a downstream consumer

backend/src/data/main.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 from copy import deepcopy
+import traceback
 from typing import Dict, List, Optional, Tuple
 
 from src.constants import CURR_YEAR, DISABLE_GCS
@@ -38,6 +39,7 @@
     update_team_years as update_team_years_db,
     update_teams as update_teams_db,
 )
+from src.google.snapshot import read_snapshot, write_snapshot
 from src.google.storage import write_objs as write_objs_storage
 
 
@@ -51,13 +53,16 @@ def process_year(
     all_team_years: Optional[Dict[int, Dict[int, TeamYear]]],
 ) -> List[Team]:
     timer = Timer()
+    curr_year_gcs = year_num == CURR_YEAR and not DISABLE_GCS
     orig_objs = deepcopy(objs)
     if all_team_years is None:
         all_team_years = defaultdict(dict)
-        for year in range(max(2002, year_num - 4), year_num):
-            team_years = get_team_years_db(year=year)
-            for ty in team_years:
-                all_team_years[ty.year][ty.team] = ty
+        try:
+            for year in range(max(2002, year_num - 4), year_num):
+                for ty in get_team_years_db(year=year):
+                    all_team_years[ty.year][ty.team] = ty
+        except Exception:
+            traceback.print_exc()
 
     new_teams, objs = process_year_tba(year_num, teams, objs, tba_partial, cache)
     teams += new_teams
@@ -74,13 +79,23 @@ def process_year(
     objs = process_year_epa(objs, all_team_years)
     timer.print(str(year_num) + " EPA")
 
-    write_objs_db(year_num, objs, orig_objs if partial else None, not partial)
-    timer.print(str(year_num) + " Write DB")
+    if curr_year_gcs:
+        write_snapshot(year_num, objs, teams)
+        timer.print(str(year_num) + " Write Snapshot")
 
-    if year_num == CURR_YEAR and not DISABLE_GCS:
-        write_objs_storage(objs, orig_objs if partial else None)
+        write_objs_storage(objs, orig_objs if partial else None, teams)
         timer.print(str(year_num) + " Write Storage")
 
+        try:
+            db_orig = read_objs_db(year_num) if partial else None
+            write_objs_db(year_num, objs, db_orig, not partial)
+        except Exception:
+            traceback.print_exc()
+        timer.print(str(year_num) + " Write DB")
+    else:
+        write_objs_db(year_num, objs, orig_objs if partial else None, not partial)
+        timer.print(str(year_num) + " Write DB")
+
     return teams
 
 
@@ -138,15 +153,23 @@ def update_curr_year(partial: bool, tba_partial: bool):
     year = CURR_YEAR
     timer = Timer()
 
-    teams = get_teams_db()
-    timer.print("Load Teams")
-
-    if partial:
-        objs: objs_type = read_objs_db(year)
-        timer.print("Read Objs")
-    else:
-        objs = create_objs(year)
-        timer.print("Create Objs")
+    objs: Optional[objs_type] = None
+    teams: Optional[List[Team]] = None
+    if partial and not DISABLE_GCS:
+        loaded = read_snapshot(year)
+        if loaded is not None:
+            objs, teams = loaded
+            timer.print("Read Snapshot")
+
+    if objs is None or teams is None:
+        teams = get_teams_db()
+        timer.print("Load Teams")
+        if partial:
+            objs = read_objs_db(year)
+            timer.print("Read Objs")
+        else:
+            objs = create_objs(year)
+            timer.print("Create Objs")
 
     teams = process_year(
         year, partial, tba_partial, year < CURR_YEAR, teams, objs, None

backend/src/data/utils.py

@@ -53,13 +53,16 @@ def read_objs(year: int) -> objs_type:
     if year_obj is None:
         raise Exception("Year not found")
 
+    def by_pk(objs: list) -> dict:
+        return {o.pk(): o for o in sorted(objs, key=lambda o: o.pk())}
+
     return (
         year_obj,
-        {t.pk(): t for t in get_team_years_db(year=year)},
-        {e.pk(): e for e in get_events_db(year=year)},
-        {te.pk(): te for te in get_team_events_db(year=year)},
-        {m.pk(): m for m in get_matches_db(year=year)},
-        {e.pk(): e for e in get_etags_db(year=year)},
+        by_pk(get_team_years_db(year=year)),
+        by_pk(get_events_db(year=year)),
+        by_pk(get_team_events_db(year=year)),
+        by_pk(get_matches_db(year=year)),
+        by_pk(get_etags_db(year=year)),
     )
 
 

backend/src/google/snapshot.py

@@ -0,0 +1,120 @@
+from enum import Enum
+import json
+import os
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Type
+from uuid import uuid4
+import zlib
+
+import attr
+from sqlalchemy import inspect
+from sqlalchemy.sql.sqltypes import Enum as SQLEnum
+
+from src.data.utils import objs_type
+from src.db.models import ETag, Event, Match, Team, TeamEvent, TeamYear, Year
+from src.db.models.etag import ETagORM
+from src.db.models.event import EventORM
+from src.db.models.main import Model, ModelORM
+from src.db.models.match import MatchORM
+from src.db.models.team import TeamORM
+from src.db.models.team_event import TeamEventORM
+from src.db.models.team_year import TeamYearORM
+from src.db.models.year import YearORM
+from src.google.storage import _bucket
+
+SNAPSHOT_SCHEMA = 1
+SNAPSHOT_PREFIX = "state"
+
+
+def snapshot_key(year: int) -> str:
+    return f"{SNAPSHOT_PREFIX}/snapshot.{year}"
+
+
+def _enum_fields(orm_type: Type[ModelORM]) -> Dict[str, Type[Enum]]:
+    fields: Dict[str, Type[Enum]] = {}
+    for column in inspect(orm_type).columns:
+        if isinstance(column.type, SQLEnum) and column.type.enum_class is not None:
+            fields[column.name] = column.type.enum_class
+    return fields
+
+
+def _load(
+    model_cls: Type[Model], orm_type: Type[ModelORM], data: Dict[str, Any]
+) -> Any:
+    obj = model_cls.from_dict(data)
+    for name, enum_cls in _enum_fields(orm_type).items():
+        value = getattr(obj, name)
+        if value is not None and not isinstance(value, enum_cls):
+            setattr(obj, name, enum_cls(value))
+    return obj
+
+
+def _dump_values(objs: Mapping[str, Model]) -> List[Dict[str, Any]]:
+    return [attr.asdict(o) for o in sorted(objs.values(), key=lambda o: o.pk())]
+
+
+def _load_values(
+    model_cls: Type[Model], orm_type: Type[ModelORM], rows: List[Dict[str, Any]]
+) -> Dict[str, Any]:
+    loaded = [_load(model_cls, orm_type, row) for row in rows]
+    return {o.pk(): o for o in loaded}
+
+
+def serialize(objs: objs_type, teams: List[Team]) -> bytes:
+    year_obj = objs[0]
+    payload = {
+        "schema": SNAPSHOT_SCHEMA,
+        "year": year_obj.year,
+        "teams": [attr.asdict(t) for t in sorted(teams, key=lambda t: t.team)],
+        "objs": {
+            "year": attr.asdict(year_obj),
+            "team_years": _dump_values(objs[1]),
+            "events": _dump_values(objs[2]),
+            "team_events": _dump_values(objs[3]),
+            "matches": _dump_values(objs[4]),
+            "etags": _dump_values(objs[5]),
+        },
+    }
+    return zlib.compress(json.dumps(payload).encode("utf-8"))
+
+
+def deserialize(raw: bytes) -> Tuple[objs_type, List[Team]]:
+    payload = json.loads(zlib.decompress(raw).decode("utf-8"))
+    schema = payload.get("schema")
+    if schema != SNAPSHOT_SCHEMA:
+        raise ValueError(
+            f"snapshot schema {schema} != expected {SNAPSHOT_SCHEMA}"
+        )
+    data = payload["objs"]
+    objs: objs_type = (
+        _load(Year, YearORM, data["year"]),
+        _load_values(TeamYear, TeamYearORM, data["team_years"]),
+        _load_values(Event, EventORM, data["events"]),
+        _load_values(TeamEvent, TeamEventORM, data["team_events"]),
+        _load_values(Match, MatchORM, data["matches"]),
+        _load_values(ETag, ETagORM, data["etags"]),
+    )
+    teams = [_load(Team, TeamORM, row) for row in payload["teams"]]
+    return objs, teams
+
+
+def write_snapshot(year: int, objs: objs_type, teams: List[Team]) -> None:
+    bucket = _bucket()
+    key = snapshot_key(year)
+    tmp_key = f"{key}.{os.getpid()}.{uuid4().hex}.tmp"
+    bucket.blob(tmp_key).upload_from_string(
+        serialize(objs, teams), "application/octet-stream"
+    )
+    bucket.copy_blob(bucket.blob(tmp_key), bucket, key)
+    bucket.blob(tmp_key).delete()
+
+
+def read_snapshot(year: int) -> Optional[Tuple[objs_type, List[Team]]]:
+    try:
+        raw = _bucket().blob(snapshot_key(year)).download_as_bytes()
+    except Exception:
+        return None
+    try:
+        return deserialize(raw)
+    except Exception as e:
+        print(f"snapshot unreadable, falling back to DB path: {e}")
+        return None

backend/src/google/storage.py

@@ -3,14 +3,15 @@
 from datetime import datetime, timedelta, timezone
 import gzip
 import json
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, TypeVar
 import zlib
 
 from google.cloud import storage
 
 from src.constants import CURR_YEAR, HIST_EPOCH, PROD
 from src.data.utils import nan_safe_eq, objs_type
 from src.db.functions import get_noteworthy_matches, get_upcoming_matches
+from src.db.models.team import Team
 from src.db.read.event import get_events as get_events_db
 from src.db.read.team import get_teams as get_teams_db
 from src.db.read.team_year import get_team_years as get_team_years_db
@@ -63,6 +64,17 @@ def upload_file_to_gcs(data: Any, object_name: str) -> None:
     _upload_bytes(_bucket(), object_name, compress(data), None)
 
 
+T = TypeVar("T")
+
+
+def _best_effort(label: str, fn: Callable[[], T], default: T) -> T:
+    try:
+        return fn()
+    except Exception as e:
+        print(f"skipped {label}: {e}")
+        return default
+
+
 def read_manifest() -> Optional[Manifest]:
     try:
         raw = _bucket().blob(MANIFEST_OBJECT).download_as_bytes()
@@ -104,6 +116,7 @@ def _publish(plan: UploadPlan) -> None:
 def write_objs(
     objs: objs_type,
     orig_objs: Optional[objs_type] = None,
+    teams: Optional[List[Team]] = None,
 ) -> None:
     year = CURR_YEAR
     year_obj = objs[0]
@@ -114,7 +127,8 @@ def add(object_name: str, data: Any) -> None:
         rendered[object_name] = compress(data)
 
     # teams/all
-    teams = get_teams_db()
+    if teams is None:
+        teams = _best_effort("teams", get_teams_db, [])
     add("teams/all", _read_all_teams(teams))
 
     # team_years/{CURR_YEAR}
@@ -129,7 +143,9 @@ def add(object_name: str, data: Any) -> None:
     )
 
     # events/all
-    add("events/all", _read_all_events(get_events_db()))
+    all_events = _best_effort("events/all", get_events_db, None)
+    if all_events is not None:
+        add("events/all", _read_all_events(all_events))
 
     # events/{CURR_YEAR}
     events = list(objs[2].values())
@@ -181,38 +197,52 @@ def add(object_name: str, data: Any) -> None:
     add("team_to_events", team_to_events)
 
     # team/{team.team}
-    all_team_years = get_team_years_db()
-    team_years_by_team: Dict[int, List[Any]] = defaultdict(list)
-    for ty in all_team_years:
-        team_years_by_team[ty.team].append(ty)
-    teams_by_num = {t.team: t for t in teams}
-    for num in {ty.team for ty in team_years}:
-        team_obj = teams_by_num.get(num)
-        if team_obj is None:
-            continue
-        add(f"team/{num}", _read_team(team_obj, team_years_by_team.get(num, [])))
+    all_team_years = _best_effort("team pages", get_team_years_db, None)
+    if all_team_years is not None:
+        team_years_by_team: Dict[int, List[Any]] = defaultdict(list)
+        for ty in all_team_years:
+            if ty.year != year:
+                team_years_by_team[ty.team].append(ty)
+        for ty in team_years:
+            team_years_by_team[ty.team].append(ty)
+        teams_by_num = {t.team: t for t in teams}
+        for num in {ty.team for ty in team_years}:
+            team_obj = teams_by_num.get(num)
+            if team_obj is None:
+                continue
+            add(f"team/{num}", _read_team(team_obj, team_years_by_team.get(num, [])))
 
     # noteworthy_matches/{year}
-    noteworthy_matches = get_noteworthy_matches(
-        year=year, country=None, state=None, district=None, elim=None, week=None
+    noteworthy_matches = _best_effort(
+        "noteworthy_matches",
+        lambda: get_noteworthy_matches(
+            year=year, country=None, state=None, district=None, elim=None, week=None
+        ),
+        None,
     )
-    add(f"noteworthy_matches/{year}", _read_noteworthy_matches(noteworthy_matches))
+    if noteworthy_matches is not None:
+        add(f"noteworthy_matches/{year}", _read_noteworthy_matches(noteworthy_matches))
 
     # upcoming_matches?limit=20&metric={predicted_time | max_epa | sum_epa | diff_epa}
     for metric in ["predicted_time", "max_epa", "sum_epa", "diff_epa"]:
-        upcoming_matches = get_upcoming_matches(
-            country=None,
-            state=None,
-            district=None,
-            elim=None,
-            minutes=-1,
-            limit=20,
-            metric=metric,
-        )
-        add(
-            f"upcoming_matches.limit=20.metric={metric}",
-            _read_upcoming_matches(upcoming_matches),
+        upcoming_matches = _best_effort(
+            f"upcoming_matches.{metric}",
+            lambda metric=metric: get_upcoming_matches(
+                country=None,
+                state=None,
+                district=None,
+                elim=None,
+                minutes=-1,
+                limit=20,
+                metric=metric,
+            ),
+            None,
         )
+        if upcoming_matches is not None:
+            add(
+                f"upcoming_matches.limit=20.metric={metric}",
+                _read_upcoming_matches(upcoming_matches),
+            )
 
     cycle = datetime.now(timezone.utc).isoformat()
     plan = plan_uploads(rendered, prev, cycle, hist_epoch=HIST_EPOCH)