Something went wrong. Try again.
Something went wrong. Try again.
15 kB · 367 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368#!/usr/bin/env python3"""Moves the item records of an account from `externalId` to `identifiers` and fills in what the app now stores (one-off).
Usage: migrate_identifiers.py --handle you.example.com --backup mediatracker-export-YYYY-MM-DD.zip \ [--pds https://pds.example.com] [--enrich] [--dry-run]
Phase 1 (always): the item records - movie.item, tvShow.item, tvShow.episode.item and listitem - held ONE `externalId`(`{ source, id }`); now they hold `identifiers` (a list of those), the `originSource` the entry was created from andmetadata. The old field is turned into `identifiers: [externalId]` and `originSource: externalId.source`, and removed.Nothing needs the network except the PDS.
Phase 2 (`--enrich`, needs a TMDB API key in the TMDB_API_KEY environment variable): for every item whose origin istmdb the script asks TMDB for what the app stores since:
- movies and shows: the ids at other sources (imdb, tvdb for shows, wikidata) and `releaseYear` (release date / first air date),- episodes: the ids at other sources (imdb, tvdb, wikidata) and `absoluteNumber`, the running number across the regular seasons (season 1 and up, specials have none), counted from the season layout like the app does,- list items: the same as the movie or show they name.
An episode needs one TMDB request of its own (there is no batch call; for the ~15,000 episodes of a large account thistakes about ten minutes at the pace the script keeps below TMDB's limit). Records that carry everything already areskipped, so both phases can be run again at any time and an interrupted run carries on where it stopped. A record keepsevery field and identifier it has (those of sources this script doesn't know included); only missing ones are added.
NOTHING IS WRITTEN WITHOUT A BACKUP: `--backup` is the ZIP of the app's export (Settings -> Account -> Export). It mustbelong to this account and hold at least as many records per collection as the PDS does now. A dry run (`--dry-run`)needs none and writes nothing (it still asks TMDB with `--enrich`, so it shows what would be added).
The app password is read from the ATPROTO_APP_PASSWORD environment variable, or asked for. Afterwards use "Read dataagain" in the app's settings, so the record index reads the new records."""import argparseimport getpassimport jsonimport osimport sysimport threadingimport timeimport urllib.errorimport urllib.requestfrom concurrent.futures import ThreadPoolExecutor
from import_to_pds import ( NAMESPACE, MOVIE_ITEM, SHOW_ITEM, EPISODE_ITEM, MAX_WRITES_PER_CALL, Pds, XrpcError, rkey_of, _chunks,)from migrate_strongrefs import check_backup
LIST_ITEM = f"{NAMESPACE}.listitem"COLLECTIONS = (MOVIE_ITEM, SHOW_ITEM, EPISODE_ITEM, LIST_ITEM)
TMDB_URL = "https://api.themoviedb.org/3"REQUESTS_PER_SECOND = 30WORKERS = 12
def reshape(value): """Phase 1: the new shape of a record with an `externalId`; None if it has none (it is in the new shape).""" external_id = value.get("externalId") if not isinstance(external_id, dict): return None
result = {key: entry for key, entry in value.items() if key != "externalId"} identifiers = list(result.get("identifiers", [])) if not any(same(identifier, external_id) for identifier in identifiers): identifiers.insert(0, {"source": external_id["source"], "id": external_id["id"]}) result["identifiers"] = identifiers result.setdefault("originSource", external_id["source"])
return result
def same(left, right): return left.get("source") == right.get("source") and left.get("id") == right.get("id")
def identifier_of(value, source): return next((identifier["id"] for identifier in value.get("identifiers", []) if identifier["source"] == source), None)
def merge(value, further, year=None, absolute=None): """The record with the missing identifiers, year and absolute number added; None if nothing was missing.""" result = dict(value) identifiers = list(value.get("identifiers", [])) known = {identifier["source"] for identifier in identifiers} changed = False
for source, id_ in further: if source not in known: identifiers.append({"source": source, "id": id_}) known.add(source) changed = True result["identifiers"] = identifiers
if year is not None and "releaseYear" not in value: result["releaseYear"] = year changed = True
if absolute is not None and "absoluteNumber" not in value: result["absoluteNumber"] = absolute changed = True
return result if changed else None
class RateLimiter: def __init__(self, per_second): self.interval = 1.0 / per_second self.lock = threading.Lock() self.next_time = time.monotonic()
def wait(self): with self.lock: now = time.monotonic() wait = max(0.0, self.next_time - now) self.next_time = max(now, self.next_time) + self.interval if wait: time.sleep(wait)
class Tmdb: """The few TMDB calls the enrichment needs; results are kept, so a title is asked for once."""
def __init__(self, api_key): self.api_key = api_key self.limiter = RateLimiter(REQUESTS_PER_SECOND) self.cache = {} self.lock = threading.Lock()
def get(self, path, **query): """The JSON of a TMDB path, None for 404; waits out the rate limit and retries a failing request.""" key = (path, tuple(sorted(query.items()))) with self.lock: if key in self.cache: return self.cache[key]
query["api_key"] = self.api_key url = f"{TMDB_URL}{path}?" + "&".join(f"{name}={value}" for name, value in query.items()) result = None for attempt in range(6): self.limiter.wait() try: with urllib.request.urlopen(url, timeout=30) as response: result = json.load(response) break except urllib.error.HTTPError as error: if error.code == 404: break if error.code == 429: time.sleep(float(error.headers.get("Retry-After", "2")) + 0.5) continue if error.code >= 500: time.sleep(1 + attempt) continue raise except (urllib.error.URLError, TimeoutError): time.sleep(1 + attempt) else: raise RuntimeError(f"TMDB did not answer for {path}")
with self.lock: self.cache[key] = result
return result
def movie(self, tmdb_id): movie = self.get(f"/movie/{tmdb_id}", append_to_response="external_ids") if movie is None: return [], None return ids_of(movie.get("external_ids")), year_of(movie.get("release_date"))
def show(self, tmdb_id): show = self.get(f"/tv/{tmdb_id}", append_to_response="external_ids") if show is None: return [], None return ids_of(show.get("external_ids")), year_of(show.get("first_air_date"))
def episode_ids(self, show_id, season, episode): external_ids = self.get(f"/tv/{show_id}/season/{season}/episode/{episode}/external_ids") return ids_of(external_ids)
def absolute_numbers(self, show_id): """Episode number -> absolute number, by (season, episode), counted like the app does.""" show = self.get(f"/tv/{show_id}") if show is None: return {} numbering, next_number = {}, 1 for season in sorted(entry["season_number"] for entry in show.get("seasons", []) if entry["season_number"] > 0): details = self.get(f"/tv/{show_id}/season/{season}") for episode in sorted(entry["episode_number"] for entry in (details or {}).get("episodes", [])): numbering[(season, episode)] = next_number next_number += 1 return numbering
def ids_of(external_ids): """(source, id) pairs of a TMDB external_ids answer, for the sources the lexicon knows.""" result = [] for source, field in (("imdb", "imdb_id"), ("tvdb", "tvdb_id"), ("wikidata", "wikidata_id")): value = (external_ids or {}).get(field) if value: result.append((source, str(value))) return result
def year_of(date): return int(date[:4]) if date and len(date) >= 4 and date[:4].isdigit() else None
def enrich(records, tmdb): """Phase 2: [(collection, rkey, new value)] for the records that miss something TMDB can tell.""" updates = {} numberings = {}
shows_by_uri = {} for record in records[SHOW_ITEM]: shows_by_uri[record["uri"]] = identifier_of(record["value"], "tmdb")
titles = {}
def title(media, tmdb_id): key = (media, tmdb_id) if key not in titles: titles[key] = tmdb.movie(tmdb_id) if media == "movie" else tmdb.show(tmdb_id) return titles[key]
def numbering(show_id): if show_id not in numberings: numberings[show_id] = tmdb.absolute_numbers(show_id) return numberings[show_id]
def apply(collection, record, ids, year=None, absolute=None): value = merge(record["value"], ids, year, absolute) if value is not None: updates[(collection, rkey_of(record["uri"]))] = value
tasks = [] for collection, media in ((MOVIE_ITEM, "movie"), (SHOW_ITEM, "tv")): for record in records[collection]: if record["value"].get("originSource") == "tmdb" and identifier_of(record["value"], "tmdb"): tasks.append(lambda c=collection, r=record, m=media: apply( c, r, *title(m, identifier_of(r["value"], "tmdb"))))
for record in records[LIST_ITEM]: value = record["value"] if value.get("originSource") == "tmdb" and identifier_of(value, "tmdb"): media = "movie" if value.get("mediaType") == "movie" else "tv" tasks.append(lambda r=record, m=media: apply(LIST_ITEM, r, *title(m, identifier_of(r["value"], "tmdb"))))
for record in records[EPISODE_ITEM]: value = record["value"] show_id = shows_by_uri.get(value.get("show")) if value.get("originSource") != "tmdb" or show_id is None: continue season, number = value.get("seasonNumber"), value.get("episodeNumber") if season is None or number is None: continue
def episode_task(r=record, show=show_id, s=season, n=number): absolute = numbering(show).get((s, n)) if s > 0 else None has_ids = len(r["value"].get("identifiers", [])) > 1 ids = [] if has_ids else tmdb.episode_ids(show, s, n) apply(EPISODE_ITEM, r, ids, None, absolute)
tasks.append(episode_task)
print(f"Asking TMDB for {len(tasks)} records (a title is asked for once) ...") done = 0 with ThreadPoolExecutor(max_workers=WORKERS) as pool: for _ in pool.map(lambda task: task(), tasks): done += 1 if done % 500 == 0: print(f" {done}/{len(tasks)}")
return [(collection, rkey, value) for (collection, rkey), value in updates.items()]
def migrate(pds, dry_run, backup, tmdb): records = {collection: pds.list_records(collection) for collection in COLLECTIONS} counts = {collection: len(items) for collection, items in records.items()} print("Read " + ", ".join(f"{count} of {collection.removeprefix(NAMESPACE + '.')}" for collection, count in counts.items()))
pending = {} for collection, items in records.items(): for record in items: value = reshape(record["value"]) if value is not None: record["value"] = value pending[(collection, rkey_of(record["uri"]))] = value
print(f"Phase 1: {len(pending)} records to move from externalId to identifiers.")
if tmdb is not None: for collection, rkey, value in enrich(records, tmdb): pending[(collection, rkey)] = value print(f"Phase 2: {len(pending)} records to write in all.")
if dry_run or not pending: return 0
check_backup(backup, pds.did, counts)
failed = written = 0 by_collection = {} for (collection, rkey), value in pending.items(): by_collection.setdefault(collection, []).append((rkey, value))
for collection, entries in by_collection.items(): short = collection.removeprefix(NAMESPACE + ".") for chunk in _chunks(entries, MAX_WRITES_PER_CALL): try: pds.apply_updates([(collection, rkey, value) for rkey, value in chunk]) except XrpcError as error: print(f" {len(chunk)} records of {short} failed: {error}") failed += len(chunk) continue written += len(chunk) print(f" {written}/{len(pending)} written ({short})")
print(f"Done: {written} written, {failed} failed.")
return failed
def main(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--handle", required=True, help="handle (or DID) of the account") parser.add_argument("--pds", default="https://bsky.social", help="URL of the account's PDS") parser.add_argument("--backup", help="ZIP of the app's export; required to write anything") parser.add_argument("--enrich", action="store_true", help="also add ids, release year and absolute numbers (TMDB)") parser.add_argument("--dry-run", action="store_true", help="only read and say what would be written") arguments = parser.parse_args()
if not arguments.dry_run and not arguments.backup: sys.exit("Export your data first (Settings -> Account -> Export) and pass the ZIP as --backup. " "Nothing is written without it.")
tmdb = None if arguments.enrich: api_key = os.environ.get("TMDB_API_KEY") if not api_key: sys.exit("--enrich needs the TMDB API key in the TMDB_API_KEY environment variable.") tmdb = Tmdb(api_key)
try: password = os.environ.get("ATPROTO_APP_PASSWORD") or getpass.getpass("App password: ") pds = Pds(arguments.pds, arguments.handle, password) print(f"Logged in as {pds.did}") failed = migrate(pds, arguments.dry_run, arguments.backup, tmdb) except XrpcError as error: sys.exit(f"The PDS answered with an error: {error}") except KeyboardInterrupt: sys.exit("\nStopped. Run it again to carry on - what is written already is skipped.")
if arguments.dry_run: print("Dry run: nothing was written.")
if failed: sys.exit(1)
print("Afterwards use \"Read data again\" in the app's settings, so the index reads the new records.")
if __name__ == "__main__": main()