Jonathan
6f325d30b9
Two row-leak bugs found in whole-branch review: - discovery.run_discovery now prunes DiscoverySeedContribution rows whose seedMbid is no longer a WatchedArtist (unfollow/removal has no FK cascade) and recomputes the touched candidates' scores this sweep, instead of letting a dead seed's contribution inflate the aggregate forever. - maybe_run_scan's exception handler around scan_chunk now calls clear_worklist and clears scan.id, matching the drain path. Previously an exception mid-chunk left the already-done rows plus the remaining worklist for that scan_id orphaned forever, since the next scan mints a fresh uuid. Also clear scan.id in the build_worklist failure handler for consistency. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
251 lines
10 KiB
Python
251 lines
10 KiB
Python
from dataclasses import dataclass
|
||
|
||
import psycopg
|
||
|
||
from lyra_worker.browser import MbBrowser
|
||
from lyra_worker.release_filter import is_core_release
|
||
from lyra_worker.similarity.base import SimilaritySource
|
||
|
||
_TRUE = {"1", "true", "yes", "on"}
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class DiscoveryConfig:
|
||
enabled: bool = False
|
||
interval_hours: int = 168
|
||
chunk_size: int = 5
|
||
similar_per_seed: int = 20
|
||
albums_per_artist: int = 1
|
||
min_score: float = 0.0
|
||
|
||
@classmethod
|
||
def from_config(cls, config: dict) -> "DiscoveryConfig":
|
||
def _int(key: str, default: int) -> int:
|
||
try:
|
||
return int(config[key])
|
||
except (KeyError, TypeError, ValueError):
|
||
return default
|
||
|
||
def _float(key: str, default: float) -> float:
|
||
try:
|
||
return float(config[key])
|
||
except (KeyError, TypeError, ValueError):
|
||
return default
|
||
|
||
return cls(
|
||
enabled=str(config.get("discover.enabled", "")).strip().lower() in _TRUE,
|
||
interval_hours=_int("discover.intervalHours", 168),
|
||
chunk_size=_int("discover.chunkSize", 5),
|
||
similar_per_seed=_int("discover.similarPerSeed", 20),
|
||
albums_per_artist=_int("discover.albumsPerArtist", 1),
|
||
min_score=_float("discover.minScore", 0.0),
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class DiscoveryResult:
|
||
artists: int
|
||
albums: int
|
||
seeds: int = 0
|
||
|
||
|
||
def _healthy_sources(sources):
|
||
"""Return sources whose health() is truthy; a source whose health() raises
|
||
is logged and skipped, never aborting the sweep."""
|
||
live = []
|
||
for s in sources:
|
||
try:
|
||
if s.health():
|
||
live.append(s)
|
||
except Exception as e: # a bad health check must not abort the sweep
|
||
print(f"worker: discovery source {getattr(s, 'name', '?')} health check failed: {e}", flush=True)
|
||
return live
|
||
|
||
|
||
def _seed_artists(conn: psycopg.Connection, cfg: DiscoveryConfig):
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
'SELECT mbid, name FROM "WatchedArtist" '
|
||
'WHERE "lastDiscoveredAt" IS NULL '
|
||
' OR "lastDiscoveredAt" < now() - make_interval(hours => %s) '
|
||
'ORDER BY "lastDiscoveredAt" ASC NULLS FIRST LIMIT %s',
|
||
(cfg.interval_hours, cfg.chunk_size),
|
||
)
|
||
return cur.fetchall()
|
||
|
||
|
||
def _followed_mbids(conn: psycopg.Connection) -> set[str]:
|
||
with conn.cursor() as cur:
|
||
cur.execute('SELECT mbid FROM "WatchedArtist"')
|
||
return {r[0] for r in cur.fetchall()}
|
||
|
||
|
||
def _upsert_artist(conn: psycopg.Connection, cand: dict) -> int:
|
||
dedupe = f"artist:{cand['mbid']}:"
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
'INSERT INTO "DiscoverySuggestion" (id, kind, "artistMbid", "artistName", '
|
||
'score, "seedCount", sources, status, "dedupeKey", "createdAt", "updatedAt") '
|
||
"VALUES (gen_random_uuid()::text, 'artist', %s, %s, %s, %s, %s, 'pending', %s, now(), now()) "
|
||
'ON CONFLICT ("dedupeKey") DO UPDATE SET '
|
||
'score = EXCLUDED.score, "seedCount" = EXCLUDED."seedCount", '
|
||
'sources = EXCLUDED.sources, "updatedAt" = now() '
|
||
"WHERE \"DiscoverySuggestion\".status = 'pending'",
|
||
(cand["mbid"], cand["name"], cand["score"], cand["seed_count"],
|
||
sorted(cand["sources"]), dedupe),
|
||
)
|
||
return cur.rowcount
|
||
|
||
|
||
def _existing_rg_mbids(conn: psycopg.Connection) -> set[str]:
|
||
with conn.cursor() as cur:
|
||
cur.execute('SELECT "rgMbid" FROM "MonitoredRelease"')
|
||
return {r[0] for r in cur.fetchall()}
|
||
|
||
|
||
def _upsert_album(conn: psycopg.Connection, artist: dict, rg) -> int:
|
||
dedupe = f"album:{artist['mbid']}:{rg.rg_mbid}"
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
'INSERT INTO "DiscoverySuggestion" (id, kind, "artistMbid", "artistName", '
|
||
'"rgMbid", album, "primaryType", "secondaryTypes", "firstReleaseDate", '
|
||
'score, "seedCount", sources, status, "dedupeKey", "createdAt", "updatedAt") '
|
||
"VALUES (gen_random_uuid()::text, 'album', %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, "
|
||
"'pending', %s, now(), now()) "
|
||
'ON CONFLICT ("dedupeKey") DO UPDATE SET score = EXCLUDED.score, '
|
||
'"updatedAt" = now() '
|
||
"WHERE \"DiscoverySuggestion\".status = 'pending'",
|
||
(artist["mbid"], artist["name"], rg.rg_mbid, rg.title,
|
||
rg.primary_type or None, list(rg.secondary_types), rg.first_release_date or None,
|
||
artist["score"], artist["seed_count"], sorted(artist["sources"]), dedupe),
|
||
)
|
||
return cur.rowcount
|
||
|
||
|
||
def _derive_albums(conn: psycopg.Connection, browser: MbBrowser, surfaced: list[dict],
|
||
cfg: DiscoveryConfig) -> int:
|
||
if cfg.albums_per_artist <= 0:
|
||
return 0
|
||
have = _existing_rg_mbids(conn)
|
||
albums = 0
|
||
for artist in surfaced:
|
||
try:
|
||
groups = browser.browse_release_groups(artist["mbid"])
|
||
except Exception as e: # one artist's browse failure must not abort the sweep
|
||
print(f"worker: discovery album browse failed for {artist['mbid']}: {e}", flush=True)
|
||
continue
|
||
core = [g for g in groups
|
||
if is_core_release(g.primary_type, g.secondary_types) and g.rg_mbid not in have]
|
||
core.sort(key=lambda g: g.first_release_date or "", reverse=True)
|
||
for rg in core[: cfg.albums_per_artist]:
|
||
albums += _upsert_album(conn, artist, rg)
|
||
have.add(rg.rg_mbid)
|
||
return albums
|
||
|
||
|
||
def _record_contributions(conn: psycopg.Connection, seed_mbid: str, live, followed: set[str],
|
||
cfg: DiscoveryConfig) -> set[str]:
|
||
"""Replace seed_mbid's contribution rows with its current similar-artist scores.
|
||
Returns the set of candidate mbids whose aggregate may have changed (old contributors
|
||
of this seed ∪ newly inserted), so callers recompute drop-offs too."""
|
||
with conn.cursor() as cur:
|
||
cur.execute('SELECT "candidateMbid" FROM "DiscoverySeedContribution" WHERE "seedMbid" = %s',
|
||
(seed_mbid,))
|
||
affected = {r[0] for r in cur.fetchall()}
|
||
cur.execute('DELETE FROM "DiscoverySeedContribution" WHERE "seedMbid" = %s', (seed_mbid,))
|
||
|
||
per_cand: dict[str, dict] = {}
|
||
for src in live:
|
||
try:
|
||
similar = src.similar_artists(seed_mbid)
|
||
except Exception as e: # a bad source/seed must not abort the sweep
|
||
print(f"worker: discovery source {src.name} failed for {seed_mbid}: {e}", flush=True)
|
||
continue
|
||
for sa in sorted(similar, key=lambda a: a.score, reverse=True)[: cfg.similar_per_seed]:
|
||
if sa.mbid in followed:
|
||
continue
|
||
c = per_cand.setdefault(sa.mbid, {"name": sa.name, "score": 0.0, "sources": set()})
|
||
c["score"] += sa.score
|
||
c["sources"].add(src.name)
|
||
if sa.name and not c["name"]:
|
||
c["name"] = sa.name
|
||
|
||
with conn.cursor() as cur:
|
||
for mbid, c in per_cand.items():
|
||
cur.execute(
|
||
'INSERT INTO "DiscoverySeedContribution" (id, "candidateMbid", "candidateName", '
|
||
'"seedMbid", score, sources, "updatedAt") '
|
||
'VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, now())',
|
||
(mbid, c["name"], seed_mbid, c["score"], sorted(c["sources"])),
|
||
)
|
||
affected.add(mbid)
|
||
return affected
|
||
|
||
|
||
def _recompute_suggestions(conn: psycopg.Connection, affected: set[str],
|
||
cfg: DiscoveryConfig) -> list[dict]:
|
||
"""Recompute each affected candidate's DiscoverySuggestion from the sum of its
|
||
contributions. Returns the candidates that produced a live pending suggestion."""
|
||
if not affected:
|
||
return []
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
'SELECT "candidateMbid", "candidateName", score, "seedMbid", sources '
|
||
'FROM "DiscoverySeedContribution" WHERE "candidateMbid" = ANY(%s)',
|
||
(list(affected),))
|
||
rows = cur.fetchall()
|
||
|
||
agg: dict[str, dict] = {}
|
||
for mbid, name, score, seed, sources in rows:
|
||
a = agg.setdefault(mbid, {"mbid": mbid, "name": name, "score": 0.0,
|
||
"seeds": set(), "sources": set()})
|
||
a["score"] += score
|
||
a["seeds"].add(seed)
|
||
a["sources"].update(sources or [])
|
||
if name and not a["name"]:
|
||
a["name"] = name
|
||
|
||
surfaced: list[dict] = []
|
||
for a in sorted(agg.values(), key=lambda a: a["score"], reverse=True):
|
||
if a["score"] < cfg.min_score:
|
||
continue
|
||
cand = {"mbid": a["mbid"], "name": a["name"], "score": a["score"],
|
||
"seed_count": len(a["seeds"]), "sources": a["sources"]}
|
||
if _upsert_artist(conn, cand): # rowcount 1 => a live pending suggestion
|
||
surfaced.append(cand)
|
||
return surfaced
|
||
|
||
|
||
def run_discovery(conn: psycopg.Connection, sources: list[SimilaritySource],
|
||
browser: MbBrowser, cfg: DiscoveryConfig) -> DiscoveryResult:
|
||
"""Record each seed's similar-artist contributions, then recompute affected
|
||
suggestions from the summed contributions — so a candidate's score is the sum of
|
||
every seed's latest contribution, independent of chunk and sweep boundaries."""
|
||
seeds = _seed_artists(conn, cfg)
|
||
followed = _followed_mbids(conn)
|
||
live = _healthy_sources(sources)
|
||
|
||
# Evict contributions from seeds no longer followed — unfollow/removal has no
|
||
# FK cascade, so without this a removed seed's score would linger and inflate
|
||
# the aggregate forever. Seed `affected` with the touched candidates so their
|
||
# suggestion scores are recomputed from the remaining live seeds this sweep.
|
||
with conn.cursor() as cur:
|
||
cur.execute('SELECT DISTINCT "candidateMbid" FROM "DiscoverySeedContribution" '
|
||
'WHERE "seedMbid" NOT IN (SELECT mbid FROM "WatchedArtist")')
|
||
affected = {r[0] for r in cur.fetchall()}
|
||
cur.execute('DELETE FROM "DiscoverySeedContribution" '
|
||
'WHERE "seedMbid" NOT IN (SELECT mbid FROM "WatchedArtist")')
|
||
|
||
for seed_mbid, _seed_name in seeds:
|
||
affected |= _record_contributions(conn, seed_mbid, live, followed, cfg)
|
||
|
||
surfaced = _recompute_suggestions(conn, affected, cfg)
|
||
albums = _derive_albums(conn, browser, surfaced, cfg)
|
||
|
||
with conn.cursor() as cur:
|
||
for seed_mbid, _seed_name in seeds:
|
||
cur.execute('UPDATE "WatchedArtist" SET "lastDiscoveredAt" = now() WHERE mbid = %s',
|
||
(seed_mbid,))
|
||
conn.commit()
|
||
return DiscoveryResult(artists=len(surfaced), albums=albums, seeds=len(seeds))
|