Something went wrong. Try again.
Something went wrong. Try again.
16 kB · 369 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370#!/usr/bin/env python3"""Imports the JSON of bonukai_to_json.py into a PDS as dev.erican.mediatracker.* records.
Usage: import_to_pds.py import.json --handle you.example.com [--pds https://pds.example.com] [--dry-run]
The app password is read from the ATPROTO_APP_PASSWORD environment variable, or asked for.
The import can be run again at any time: what already exists in the PDS (items by external id, ratings and trackedrecords by item, watches by item and date) is skipped. That makes it safe to interrupt, and it is what lets itcarry on after a rate limit. A PDS that limits writes is waited out (until the time it reports), not failed."""import argparseimport getpassimport jsonimport osimport sysimport timeimport urllib.errorimport urllib.parseimport urllib.requestfrom collections import Counterfrom datetime import datetime, timezone
NAMESPACE = "dev.erican.mediatracker"MOVIE_ITEM = f"{NAMESPACE}.movie.item"MOVIE_TRACKED = f"{NAMESPACE}.movie.tracked"MOVIE_RATING = f"{NAMESPACE}.movie.rating"MOVIE_WATCH = f"{NAMESPACE}.movie.watch"SHOW_ITEM = f"{NAMESPACE}.tvShow.item"SHOW_TRACKED = f"{NAMESPACE}.tvShow.tracked"SHOW_RATING = f"{NAMESPACE}.tvShow.rating"SHOW_WATCH = f"{NAMESPACE}.tvShow.watch"EPISODE_ITEM = f"{NAMESPACE}.tvShow.episode.item"EPISODE_RATING = f"{NAMESPACE}.tvShow.episode.rating"
CLIENT_AGENT = "mediatracker-import/bonukai"# A PDS accepts 200 writes per call at most, but also caps the request body (150 KB in the reference PDS).MAX_WRITES_PER_CALL = 100LIST_PAGE_SIZE = 100TID_ALPHABET = "234567abcdefghijklmnopqrstuvwxyz"
class XrpcError(Exception): def __init__(self, status, error, message, headers): super().__init__(f"{status} {error}: {message}") self.status = status self.error = error self.headers = headers
class TidGenerator: """atproto TIDs: 53 bits of microseconds and a 10 bit clock id in 13 base32-sortable characters."""
def __init__(self): self._last = 0
def next(self): value = max(time.time_ns() // 1000, self._last + 1) self._last = value number = (value << 10) | 1 characters = [] for _ in range(13): characters.append(TID_ALPHABET[number & 31]) number >>= 5 return "".join(reversed(characters))
class Pds: """The few XRPC calls of a logged-in repo, with token refresh and waiting out rate limits."""
def __init__(self, url, identifier, password): self.url = url.rstrip("/") self.session = self._call("POST", "com.atproto.server.createSession", {"identifier": identifier, "password": password}, authorize=False) self.did = self.session["did"]
def list_records(self, collection): records, cursor = [], None while True: query = {"repo": self.did, "collection": collection, "limit": LIST_PAGE_SIZE} if cursor: query["cursor"] = cursor page = self._call("GET", "com.atproto.repo.listRecords", query=query) records.extend(page["records"]) cursor = page.get("cursor") if not cursor or not page["records"]: return records
def put_record(self, collection, rkey, value): """Creates or overwrites the record under `rkey`; returns its (uri, cid).""" response = self._call("POST", "com.atproto.repo.putRecord", {"repo": self.did, "collection": collection, "rkey": rkey, "record": value}) return response["uri"], response["cid"]
def apply_writes(self, writes): """Creates the records of `writes` in one call; returns their (uri, cid) in order.""" body = {"repo": self.did, "writes": [ {"$type": "com.atproto.repo.applyWrites#create", "collection": collection, "rkey": rkey, "value": value} for collection, rkey, value in writes ]} response = self._call("POST", "com.atproto.repo.applyWrites", body) return [(result["uri"], result["cid"]) for result in response["results"]]
def apply_updates(self, writes): """Overwrites the existing records of `writes` ((collection, rkey, value)) in one call.""" body = {"repo": self.did, "writes": [ {"$type": "com.atproto.repo.applyWrites#update", "collection": collection, "rkey": rkey, "value": value} for collection, rkey, value in writes ]} self._call("POST", "com.atproto.repo.applyWrites", body)
def _call(self, method, nsid, body=None, query=None, authorize=True): while True: url = f"{self.url}/xrpc/{nsid}" + (f"?{urllib.parse.urlencode(query)}" if query else "") headers = {"Content-Type": "application/json"} if body is not None else {} if authorize: headers["Authorization"] = f"Bearer {self.session['accessJwt']}" request = urllib.request.Request( url, method=method, headers=headers, data=None if body is None else json.dumps(body).encode() ) try: with urllib.request.urlopen(request, timeout=120) as response: payload = response.read() return json.loads(payload) if payload else {} except urllib.error.HTTPError as error: problem = _read_problem(error) if error.code == 429: self._wait_for_rate_limit(error.headers) continue if authorize and problem.get("error") == "ExpiredToken": self._refresh() continue raise XrpcError(error.code, problem.get("error"), problem.get("message"), error.headers) from None
def _refresh(self): print(" refreshing the access token") request = urllib.request.Request( f"{self.url}/xrpc/com.atproto.server.refreshSession", method="POST", headers={"Authorization": f"Bearer {self.session['refreshJwt']}"}, ) with urllib.request.urlopen(request, timeout=120) as response: self.session = json.loads(response.read())
@staticmethod def _wait_for_rate_limit(headers): reset = headers.get("ratelimit-reset") seconds = max(int(reset) - time.time(), 0) + 2 if reset else 60 until = datetime.fromtimestamp(time.time() + seconds).strftime("%H:%M:%S") print(f" the PDS rate limits writes - waiting {int(seconds)}s (until {until}), Ctrl+C stops safely") time.sleep(seconds)
def _read_problem(error): try: return json.loads(error.read()) except (ValueError, OSError): return {}
def normalize_time(value): """A datetime string as UTC milliseconds, so the app's and this script's spellings of a time compare equal.""" if not value: return None moment = datetime.fromisoformat(value) if moment.tzinfo is None: moment = moment.replace(tzinfo=timezone.utc) return round(moment.timestamp() * 1000)
def now(): moment = datetime.now(timezone.utc) return moment.strftime("%Y-%m-%dT%H:%M:%S.") + f"{moment.microsecond // 1000:03d}Z"
def external_key(external_id): return external_id["source"], external_id["id"]
def identity(external_id): """The `identifiers` and `originSource` of an item record created from this source's id.""" return {"identifiers": [{"source": external_id["source"], "id": external_id["id"]}], "originSource": external_id["source"]}
def reference_uri(reference): """The AT-URI of a reference to a record: the string, or the `uri` of the strongRef of records from before.""" return reference["uri"] if isinstance(reference, dict) else reference
def rkey_of(uri): return uri.rsplit("/", 1)[1]
class Import: def __init__(self, pds, data, dry_run): self.pds = pds self.data = data self.dry_run = dry_run self.tids = TidGenerator() self.created_at = now() # external id -> uri, per collection of items self.items = {MOVIE_ITEM: {}, SHOW_ITEM: {}, EPISODE_ITEM: {}} # subject uri of the tracked/rating records that exist self.subjects = {MOVIE_TRACKED: set(), MOVIE_RATING: set(), SHOW_TRACKED: set(), SHOW_RATING: set(), EPISODE_RATING: set()} # (subject uri, watchedAt in ms) -> how many such watches exist self.watches = Counter() self.written = 0 self.total = 0
def run(self): self._load_existing()
self._plan_and_write_items() self._plan_and_write_records()
# --- what is there already -------------------------------------------------------------------------------------
def _load_existing(self): print("Reading what the PDS already has ...") for collection in self.items: for record in self.pds.list_records(collection): for identifier in record["value"].get("identifiers", []): self.items[collection].setdefault(external_key(identifier), record["uri"]) for collection in self.subjects: for record in self.pds.list_records(collection): self.subjects[collection].add(reference_uri(record["value"]["subject"])) for collection in (MOVIE_WATCH, SHOW_WATCH): for record in self.pds.list_records(collection): self.watches[(reference_uri(record["value"]["subject"]), normalize_time(record["value"].get("watchedAt")))] += 1 print(f" {sum(len(items) for items in self.items.values())} items, " f"{sum(len(subjects) for subjects in self.subjects.values())} tracked/rating records, " f"{sum(self.watches.values())} watches")
# --- items: movies and shows first, then the episodes that point at their show -----------------------------------
def _plan_and_write_items(self): first = [(MOVIE_ITEM, movie) for movie in self.data["movies"]] + \ [(SHOW_ITEM, show) for show in self.data["tvShows"]] self._create_items(first, lambda collection, entry: { "$type": collection, "title": entry["title"], **identity(entry["externalId"]) }, "movie and show items")
episodes = [(EPISODE_ITEM, (show, episode)) for show in self.data["tvShows"] for episode in show["episodes"]] self._create_items(episodes, self._episode_item, "episode items")
def _episode_item(self, collection, entry): show, episode = entry show_uri = self.items[SHOW_ITEM][external_key(show["externalId"])] value = { "$type": collection, "show": show_uri, "seasonNumber": episode["seasonNumber"], "episodeNumber": episode["episodeNumber"], **identity(episode["externalId"]), } if episode.get("title"): value["title"] = episode["title"] return value
def _create_items(self, entries, build, label): pending = [] for collection, entry in entries: external_id = (entry[1] if collection == EPISODE_ITEM else entry)["externalId"] if external_key(external_id) not in self.items[collection]: pending.append((collection, external_key(external_id), entry)) print(f"{label}: {len(pending)} to create") if self.dry_run: self.total += len(pending) return for chunk in _chunks(pending, MAX_WRITES_PER_CALL): writes = [(collection, self.tids.next(), build(collection, entry)) for collection, _, entry in chunk] results = self.pds.apply_writes(writes) for (collection, key, _), result in zip(chunk, results): self.items[collection][key] = result[0] self._progress(len(chunk))
# --- tracked, ratings and watches, which point at items ----------------------------------------------------------
def _plan_and_write_records(self): writes = []
for movie in self.data["movies"]: self._plan_entry(writes, movie, MOVIE_ITEM, MOVIE_TRACKED, MOVIE_RATING, MOVIE_WATCH)
for show in self.data["tvShows"]: self._plan_entry(writes, show, SHOW_ITEM, SHOW_TRACKED, SHOW_RATING, SHOW_WATCH) for episode in show["episodes"]: self._plan_entry(writes, episode, EPISODE_ITEM, None, EPISODE_RATING, SHOW_WATCH)
print(f"tracked, rating and watch records: {len(writes)} to create") if self.dry_run: self.total += len(writes) return for chunk in _chunks(writes, MAX_WRITES_PER_CALL): self.pds.apply_writes(chunk) self._progress(len(chunk))
def _plan_entry(self, writes, entry, item_collection, tracked_collection, rating_collection, watch_collection): if self.dry_run and external_key(entry["externalId"]) not in self.items[item_collection]: # Not created in a dry run, so give it a stand-in URI: everything of it counts as new. ref = f"at://dry-run/{item_collection}/{external_key(entry['externalId'])}" else: ref = self.items[item_collection][external_key(entry["externalId"])] rkey = rkey_of(ref)
if entry.get("tracked") and ref not in self.subjects[tracked_collection]: writes.append((tracked_collection, rkey, { "$type": tracked_collection, "subject": ref, "createdAt": entry["tracked"]["createdAt"] }))
if entry.get("rating") and ref not in self.subjects[rating_collection]: writes.append((rating_collection, rkey, { "$type": rating_collection, "subject": ref, "rating": entry["rating"]["rating"], "createdAt": entry["rating"]["createdAt"] }))
for watch in entry["watches"]: key = (ref, normalize_time(watch.get("watchedAt"))) if self.watches[key] > 0: self.watches[key] -= 1 continue value = {"$type": watch_collection, "subject": ref, "createdAt": self.created_at, "submissionClientAgent": CLIENT_AGENT} if watch.get("watchedAt"): value["watchedAt"] = watch["watchedAt"] writes.append((watch_collection, self.tids.next(), value))
def _progress(self, count): self.written += count print(f" {self.written} records written")
def _chunks(items, size): for start in range(0, len(items), size): yield items[start:start + size]
def main(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("file", help="the JSON written by bonukai_to_json.py") parser.add_argument("--handle", required=True, help="handle (or DID) of the account to import into") parser.add_argument("--pds", default="https://bsky.social", help="URL of the account's PDS") parser.add_argument("--dry-run", action="store_true", help="only read the PDS and count what would be written") arguments = parser.parse_args()
with open(arguments.file, encoding="utf-8") as file: data = json.load(file)
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}") importer = Import(pds, data, arguments.dry_run) importer.run() 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(f"Dry run: {importer.total} records would be created.") else: print(f"Done: {importer.written} records created.")
if __name__ == "__main__": main()