zzstoatzz.io
search for standard sites
pub-search.waow.tech
search
zig
blog
atproto
10 kB
298 lines
1#!/usr/bin/env -S uv run --script --quiet
2# /// script
3# requires-python = ">=3.12"
4# dependencies = ["httpx", "pydantic-settings"]
5# ///
6"""
7Backfill `documents.content_type` from PDS records.
8
9Reads `content.$type` from each record on its PDS and writes it back to Turso.
10Skips whitewind (no content.$type) and rows already populated.
11
12Does NOT bump `indexed_at` — content_type isn't in the local replica schema yet,
13so we don't want to force a full re-sync of every touched row.
14
15Usage:
16 TURSO_URL=libsql://... TURSO_TOKEN=... ./scripts/backfill-content-type
17 ./scripts/backfill-content-type --dry-run
18 ./scripts/backfill-content-type --limit 100 # cap distinct (did, collection) groups
19 ./scripts/backfill-content-type --workers 10 # PDS concurrency
20"""
21
22from __future__ import annotations
23
24import argparse
25import os
26import sys
27import time
28from concurrent.futures import ThreadPoolExecutor, as_completed
29from dataclasses import dataclass
30
31import httpx
32from pydantic_settings import BaseSettings, SettingsConfigDict
33
34
35class Settings(BaseSettings):
36 model_config = SettingsConfigDict(
37 env_file=os.environ.get("ENV_FILE", ".env"), extra="ignore"
38 )
39
40 turso_url: str
41 turso_token: str
42
43 @property
44 def turso_host(self) -> str:
45 url = self.turso_url
46 if url.startswith("libsql://"):
47 url = url[len("libsql://") :]
48 return url
49
50
51# ---- Turso HTTP pipeline -----------------------------------------------------
52
53
54def turso_pipeline(settings: Settings, requests: list[dict]) -> list[dict]:
55 """POST a batch of statements; return the per-statement results."""
56 last_err: Exception | None = None
57 for attempt in range(3):
58 try:
59 resp = httpx.post(
60 f"https://{settings.turso_host}/v2/pipeline",
61 headers={
62 "Authorization": f"Bearer {settings.turso_token}",
63 "Content-Type": "application/json",
64 },
65 json={"requests": [*requests, {"type": "close"}]},
66 timeout=180,
67 )
68 resp.raise_for_status()
69 return resp.json()["results"]
70 except (httpx.ReadTimeout, httpx.ConnectError, httpx.RemoteProtocolError) as e:
71 last_err = e
72 print(f" turso pipeline attempt {attempt + 1} failed: {type(e).__name__}; retrying...", file=sys.stderr)
73 time.sleep(2 * (attempt + 1))
74 assert last_err is not None
75 raise last_err
76
77
78def turso_query(settings: Settings, sql: str, args: list | None = None) -> list[list]:
79 stmt: dict = {"sql": sql}
80 if args:
81 stmt["args"] = [
82 {"type": "null"} if a is None else {"type": "text", "value": str(a)}
83 for a in args
84 ]
85 results = turso_pipeline(settings, [{"type": "execute", "stmt": stmt}])
86 r = results[0]
87 if r["type"] == "error":
88 raise RuntimeError(f"Turso query error: {r['error']}")
89 rows = r["response"]["result"]["rows"]
90 return [[c["value"] if isinstance(c, dict) and "value" in c else None for c in row] for row in rows]
91
92
93def turso_batch_update(settings: Settings, updates: list[tuple[str, str, str]]) -> int:
94 """Apply UPDATE documents SET content_type=? WHERE did=? AND rkey=? in batches."""
95 BATCH = 50
96 written = 0
97 for i in range(0, len(updates), BATCH):
98 chunk = updates[i : i + BATCH]
99 reqs = [
100 {
101 "type": "execute",
102 "stmt": {
103 "sql": "UPDATE documents SET content_type = ? WHERE did = ? AND rkey = ?",
104 "args": [
105 {"type": "text", "value": ct},
106 {"type": "text", "value": did},
107 {"type": "text", "value": rkey},
108 ],
109 },
110 }
111 for (ct, did, rkey) in chunk
112 ]
113 results = turso_pipeline(settings, reqs)
114 for r in results:
115 if r["type"] == "error":
116 print(f" turso UPDATE error: {r['error']}", file=sys.stderr)
117 continue
118 response = r.get("response") or {}
119 if response.get("type") != "execute":
120 continue # skip the trailing close response
121 written += int(response.get("result", {}).get("affected_row_count", 0))
122 return written
123
124
125# ---- PDS resolution + fetch --------------------------------------------------
126
127
128_pds_cache: dict[str, str | None] = {}
129
130
131def resolve_pds(did: str) -> str | None:
132 """Return PDS endpoint URL for a DID, or None if unresolvable."""
133 if did in _pds_cache:
134 return _pds_cache[did]
135 url: str | None = None
136 try:
137 if did.startswith("did:plc:"):
138 resp = httpx.get(f"https://plc.directory/{did}", timeout=15)
139 resp.raise_for_status()
140 for svc in resp.json().get("service", []):
141 if svc.get("type") == "AtprotoPersonalDataServer":
142 url = svc["serviceEndpoint"]
143 break
144 elif did.startswith("did:web:"):
145 host = did[len("did:web:") :].replace(":", "/")
146 resp = httpx.get(f"https://{host}/.well-known/did.json", timeout=15, follow_redirects=True)
147 resp.raise_for_status()
148 for svc in resp.json().get("service", []):
149 if svc.get("type") == "AtprotoPersonalDataServer":
150 url = svc["serviceEndpoint"]
151 break
152 except Exception:
153 url = None
154 _pds_cache[did] = url
155 return url
156
157
158def list_records(pds: str, did: str, collection: str) -> list[dict]:
159 records: list[dict] = []
160 cursor: str | None = None
161 while True:
162 params: dict = {"repo": did, "collection": collection, "limit": 100}
163 if cursor:
164 params["cursor"] = cursor
165 resp = httpx.get(
166 f"{pds}/xrpc/com.atproto.repo.listRecords", params=params, timeout=30
167 )
168 resp.raise_for_status()
169 data = resp.json()
170 records.extend(data.get("records", []))
171 cursor = data.get("cursor")
172 if not cursor:
173 break
174 return records
175
176
177# ---- Backfill --------------------------------------------------------------
178
179
180@dataclass
181class Group:
182 did: str
183 collection: str
184 expected: int # how many rows we expect to update in this group
185
186
187def gather_groups(settings: Settings, limit: int | None) -> list[Group]:
188 sql = (
189 "SELECT did, source_collection, COUNT(*) "
190 "FROM documents "
191 "WHERE content_type IS NULL "
192 " AND source_collection NOT LIKE 'com.whtwnd.%' "
193 "GROUP BY did, source_collection "
194 "ORDER BY COUNT(*) DESC"
195 )
196 if limit:
197 sql += f" LIMIT {int(limit)}"
198 rows = turso_query(settings, sql)
199 return [Group(did=r[0], collection=r[1], expected=int(r[2])) for r in rows]
200
201
202def fetch_group(group: Group) -> tuple[Group, list[tuple[str, str, str]] | None, str]:
203 """Return (group, [(content_type, did, rkey), ...] | None, status)."""
204 pds = resolve_pds(group.did)
205 if not pds:
206 return group, None, "no-pds"
207 try:
208 records = list_records(pds, group.did, group.collection)
209 except httpx.HTTPStatusError as e:
210 if e.response.status_code == 400:
211 return group, [], "empty"
212 return group, None, f"pds-{e.response.status_code}"
213 except Exception as e:
214 return group, None, f"pds-err:{type(e).__name__}"
215
216 updates: list[tuple[str, str, str]] = []
217 for rec in records:
218 uri = rec.get("uri", "")
219 rkey = uri.rsplit("/", 1)[-1] if uri else ""
220 if not rkey:
221 continue
222 value = rec.get("value") or {}
223 content = value.get("content")
224 content_type = ""
225 if isinstance(content, dict):
226 ct = content.get("$type")
227 if isinstance(ct, str):
228 content_type = ct
229 updates.append((content_type, group.did, rkey))
230 return group, updates, "ok"
231
232
233def main() -> int:
234 parser = argparse.ArgumentParser(description="Backfill documents.content_type from PDS records")
235 parser.add_argument("--dry-run", action="store_true", help="Fetch from PDSes but don't write")
236 parser.add_argument("--limit", type=int, default=None, help="Cap distinct (did, collection) groups (debug)")
237 parser.add_argument("--workers", type=int, default=8, help="PDS concurrency")
238 args = parser.parse_args()
239
240 try:
241 settings = Settings() # type: ignore
242 except Exception as e:
243 print(f"error loading settings: {e}", file=sys.stderr)
244 print("required env vars: TURSO_URL, TURSO_TOKEN", file=sys.stderr)
245 return 1
246
247 print("loading groups...")
248 groups = gather_groups(settings, args.limit)
249 total_expected = sum(g.expected for g in groups)
250 print(f" {len(groups)} (did, collection) groups, {total_expected} rows to attempt")
251
252 if not groups:
253 print("nothing to backfill")
254 return 0
255
256 pending_writes: list[tuple[str, str, str]] = []
257 started = time.time()
258 done = 0
259 status_counts: dict[str, int] = {}
260
261 with ThreadPoolExecutor(max_workers=args.workers) as pool:
262 futures = {pool.submit(fetch_group, g): g for g in groups}
263 for fut in as_completed(futures):
264 group, updates, status = fut.result()
265 status_counts[status] = status_counts.get(status, 0) + 1
266 done += 1
267 if updates:
268 pending_writes.extend(updates)
269
270 if done % 50 == 0 or done == len(groups):
271 elapsed = time.time() - started
272 rate = done / elapsed if elapsed > 0 else 0
273 print(
274 f" [{done}/{len(groups)}] {rate:.1f} groups/s, "
275 f"{len(pending_writes)} pending writes, statuses={status_counts}"
276 )
277
278 print(f"\nfetched {len(pending_writes)} record-level updates")
279 if args.dry_run:
280 print("dry-run: skipping writes")
281 # show a sample
282 by_ct: dict[str, int] = {}
283 for ct, _, _ in pending_writes:
284 by_ct[ct] = by_ct.get(ct, 0) + 1
285 print("content_type distribution (dry-run):")
286 for ct, n in sorted(by_ct.items(), key=lambda kv: -kv[1])[:20]:
287 print(f" {n:>6} {ct!r}")
288 return 0
289
290 print("writing...")
291 written = turso_batch_update(settings, pending_writes)
292 print(f" wrote {written} rows")
293
294 return 0
295
296
297if __name__ == "__main__":
298 sys.exit(main())