decompress_objects.py
python
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef
debug(push/stream): instrument O-frame decode path with INF…
Sonnet 4.6
patch
121 days ago
| 1 | #!/usr/bin/env python3 |
| 2 | """One-time backfill: decompress zlib-stored objects in R2 and replace with plain bytes. |
| 3 | |
| 4 | Objects pushed via the old wire path were stored zlib-compressed in R2 under the |
| 5 | SHA-256 of their *plain* content. This violates content-addressing: the declared |
| 6 | identity (SHA-256 of plain bytes) does not match the stored bytes (compressed). |
| 7 | |
| 8 | This script corrects all such objects: |
| 9 | |
| 10 | 1. Pages through musehub_objects rows where storage_uri starts with "s3://". |
| 11 | 2. Fetches each object from R2. |
| 12 | 3. Skips objects that are already plain bytes. |
| 13 | 4. For zlib-compressed objects: |
| 14 | a. Decompresses. |
| 15 | b. Verifies SHA-256(decompressed) == object_id. Skips on mismatch. |
| 16 | c. Re-uploads plain bytes to R2 (same key — idempotent). |
| 17 | d. Updates size_bytes in DB (content_cache stays NULL). |
| 18 | 5. Reports totals and any errors. |
| 19 | |
| 20 | After a successful run, decompress_if_needed() is no longer needed on the read |
| 21 | path — all objects in R2 are guaranteed to be plain bytes. |
| 22 | |
| 23 | Run inside Docker on the target instance: |
| 24 | |
| 25 | docker exec musehub-blue python3 /app/deploy/decompress_objects.py [--dry-run] [--batch 200] [--concurrency 16] |
| 26 | |
| 27 | Options: |
| 28 | --dry-run Print what would be changed without touching R2 or the DB. |
| 29 | --batch N DB page size (default 200). |
| 30 | --concurrency N Parallel R2 fetches per batch (default 16). |
| 31 | --repo-id UUID Limit to a single repo (targeted fix). |
| 32 | """ |
| 33 | from __future__ import annotations |
| 34 | |
| 35 | import argparse |
| 36 | import asyncio |
| 37 | import sys |
| 38 | import time |
| 39 | import zlib |
| 40 | |
| 41 | import sqlalchemy as sa |
| 42 | from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine |
| 43 | from sqlalchemy.orm import sessionmaker |
| 44 | |
| 45 | from muse.core.types import blob_id, split_id |
| 46 | from musehub.config import settings |
| 47 | from musehub.db import musehub_models as db |
| 48 | from musehub.storage import get_backend |
| 49 | |
| 50 | _ZLIB_MAGIC = (b"\x78\x01", b"\x78\x9c", b"\x78\xda", b"\x78\x5e") |
| 51 | |
| 52 | |
| 53 | def _is_zlib(data: bytes) -> bool: |
| 54 | return len(data) >= 2 and data[:2] in _ZLIB_MAGIC |
| 55 | |
| 56 | |
| 57 | def _decompress(data: bytes) -> bytes | None: |
| 58 | try: |
| 59 | return zlib.decompress(data) |
| 60 | except zlib.error: |
| 61 | return None |
| 62 | |
| 63 | |
| 64 | def _fmt_eta(seconds: float) -> str: |
| 65 | if seconds < 60: |
| 66 | return f"{seconds:.0f}s" |
| 67 | if seconds < 3600: |
| 68 | return f"{seconds / 60:.1f}m" |
| 69 | return f"{seconds / 3600:.1f}h" |
| 70 | |
| 71 | |
| 72 | async def _get_header(backend: object, object_id: str) -> bytes | None: |
| 73 | """Return the first 2 bytes of an object using a Range GET where possible.""" |
| 74 | from musehub.storage.backends import S3Backend, LocalBackend |
| 75 | |
| 76 | if isinstance(backend, S3Backend): |
| 77 | client = backend._get_client() |
| 78 | key = backend._key(object_id) |
| 79 | |
| 80 | def _range_get() -> bytes | None: |
| 81 | try: |
| 82 | resp = client.get_object( |
| 83 | Bucket=backend._bucket, Key=key, Range="bytes=0-1" |
| 84 | ) |
| 85 | return resp["Body"].read(2) |
| 86 | except Exception: |
| 87 | return None |
| 88 | |
| 89 | return await asyncio.to_thread(_range_get) |
| 90 | |
| 91 | if isinstance(backend, LocalBackend): |
| 92 | path = backend._path(object_id) |
| 93 | |
| 94 | def _read_header() -> bytes | None: |
| 95 | try: |
| 96 | with open(path, "rb") as fh: |
| 97 | return fh.read(2) |
| 98 | except Exception: |
| 99 | return None |
| 100 | |
| 101 | return await asyncio.to_thread(_read_header) |
| 102 | |
| 103 | data = await backend.get(object_id) # type: ignore[attr-defined] |
| 104 | return data[:2] if data else None |
| 105 | |
| 106 | |
| 107 | async def backfill(dry_run: bool, quiet: bool, batch_size: int, concurrency: int, repo_id: str | None) -> int: |
| 108 | engine = create_async_engine(settings.database_url, echo=False) |
| 109 | async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) |
| 110 | backend = get_backend() |
| 111 | |
| 112 | # ── count total objects up front so we know the denominator ────────────── |
| 113 | async with async_session() as session: |
| 114 | count_stmt = ( |
| 115 | sa.select(sa.func.count()).select_from(db.MusehubObject) |
| 116 | .where( |
| 117 | db.MusehubObject.storage_uri.like("s3://%"), |
| 118 | db.MusehubObject.deleted_at.is_(None), |
| 119 | ) |
| 120 | ) |
| 121 | if repo_id: |
| 122 | count_stmt = ( |
| 123 | sa.select(sa.func.count()).select_from(db.MusehubObject) |
| 124 | .join(db.MusehubObjectRef, db.MusehubObject.object_id == db.MusehubObjectRef.object_id) |
| 125 | .where( |
| 126 | db.MusehubObjectRef.repo_id == repo_id, |
| 127 | db.MusehubObject.storage_uri.like("s3://%"), |
| 128 | db.MusehubObject.deleted_at.is_(None), |
| 129 | ) |
| 130 | ) |
| 131 | total_objects: int = (await session.execute(count_stmt)).scalar_one() |
| 132 | |
| 133 | scope = f"repo_id={repo_id}" if repo_id else "all repos" |
| 134 | print(f"Backfill scope: {scope}") |
| 135 | print(f"Total objects to scan: {total_objects:,}") |
| 136 | if total_objects == 0: |
| 137 | print("Nothing to do.") |
| 138 | return 0 |
| 139 | print() |
| 140 | |
| 141 | # ── shared progress state (updated inside asyncio tasks) ───────────────── |
| 142 | done_count = 0 |
| 143 | plain_count = 0 |
| 144 | decompressed_count = 0 |
| 145 | hash_mismatch_count = 0 |
| 146 | error_count = 0 |
| 147 | start_time = time.monotonic() |
| 148 | progress_lock = asyncio.Lock() |
| 149 | |
| 150 | def _progress_line(extra: str = "") -> None: |
| 151 | if quiet: |
| 152 | return |
| 153 | elapsed = time.monotonic() - start_time |
| 154 | rate = done_count / elapsed if elapsed > 0 else 0 |
| 155 | remaining = total_objects - done_count |
| 156 | eta_str = _fmt_eta(remaining / rate) if rate > 0 else "?" |
| 157 | pct = 100 * done_count / total_objects if total_objects else 100 |
| 158 | print( |
| 159 | f"\r [{done_count:>{len(str(total_objects))}}/{total_objects}]" |
| 160 | f" {pct:5.1f}%" |
| 161 | f" {remaining:,} remaining" |
| 162 | f" {rate:.1f} obj/s" |
| 163 | f" ETA {eta_str}" |
| 164 | + (f" {extra}" if extra else ""), |
| 165 | end="", |
| 166 | flush=True, |
| 167 | ) |
| 168 | |
| 169 | total_checked = 0 |
| 170 | total_errors = 0 |
| 171 | |
| 172 | async with async_session() as session: |
| 173 | offset = 0 |
| 174 | while True: |
| 175 | obj_stmt = ( |
| 176 | sa.select(db.MusehubObject.object_id) |
| 177 | .where( |
| 178 | db.MusehubObject.storage_uri.like("s3://%"), |
| 179 | db.MusehubObject.deleted_at.is_(None), |
| 180 | ) |
| 181 | ) |
| 182 | if repo_id: |
| 183 | obj_stmt = ( |
| 184 | sa.select(db.MusehubObject.object_id) |
| 185 | .join(db.MusehubObjectRef, db.MusehubObject.object_id == db.MusehubObjectRef.object_id) |
| 186 | .where( |
| 187 | db.MusehubObjectRef.repo_id == repo_id, |
| 188 | db.MusehubObject.storage_uri.like("s3://%"), |
| 189 | db.MusehubObject.deleted_at.is_(None), |
| 190 | ) |
| 191 | ) |
| 192 | |
| 193 | rows = (await session.execute( |
| 194 | obj_stmt |
| 195 | .order_by(db.MusehubObject.object_id) |
| 196 | .offset(offset) |
| 197 | .limit(batch_size) |
| 198 | )).scalars().all() |
| 199 | |
| 200 | if not rows: |
| 201 | break |
| 202 | |
| 203 | offset += len(rows) |
| 204 | total_checked += len(rows) |
| 205 | sem = asyncio.Semaphore(concurrency) |
| 206 | |
| 207 | async def _process(oid: str) -> tuple[str, str, int]: |
| 208 | nonlocal done_count, plain_count, decompressed_count |
| 209 | nonlocal hash_mismatch_count, error_count |
| 210 | |
| 211 | async with sem: |
| 212 | status = "plain" |
| 213 | new_size = 0 |
| 214 | detail = "" |
| 215 | |
| 216 | try: |
| 217 | header = await _get_header(backend, oid) |
| 218 | except Exception as exc: |
| 219 | print(f"\n ERROR fetching header {oid}: {exc}", file=sys.stderr) |
| 220 | async with progress_lock: |
| 221 | done_count += 1 |
| 222 | error_count += 1 |
| 223 | _progress_line() |
| 224 | return oid, "error", 0 |
| 225 | |
| 226 | if header is None or not _is_zlib(header): |
| 227 | async with progress_lock: |
| 228 | done_count += 1 |
| 229 | plain_count += 1 |
| 230 | _progress_line() |
| 231 | return oid, "plain", 0 |
| 232 | |
| 233 | # Has zlib header — fetch full object. |
| 234 | try: |
| 235 | data = await backend.get(oid) |
| 236 | except Exception as exc: |
| 237 | print(f"\n ERROR fetching {oid}: {exc}", file=sys.stderr) |
| 238 | async with progress_lock: |
| 239 | done_count += 1 |
| 240 | error_count += 1 |
| 241 | _progress_line() |
| 242 | return oid, "error", 0 |
| 243 | |
| 244 | if data is None: |
| 245 | async with progress_lock: |
| 246 | done_count += 1 |
| 247 | plain_count += 1 |
| 248 | _progress_line() |
| 249 | return oid, "plain", 0 |
| 250 | |
| 251 | decompressed = _decompress(data) |
| 252 | if decompressed is None: |
| 253 | detail = f"zlib header but decompress failed — skipping" |
| 254 | async with progress_lock: |
| 255 | done_count += 1 |
| 256 | error_count += 1 |
| 257 | _progress_line(f"WARN {oid} {detail}") |
| 258 | return oid, "error", 0 |
| 259 | |
| 260 | if blob_id(decompressed) != oid: |
| 261 | _, bare_oid = split_id(oid) |
| 262 | _, actual = split_id(blob_id(decompressed)) |
| 263 | detail = f"hash mismatch (declared={bare_oid[:12]}… actual={actual[:12]}…)" |
| 264 | async with progress_lock: |
| 265 | done_count += 1 |
| 266 | hash_mismatch_count += 1 |
| 267 | _progress_line(f"WARN {oid} {detail}") |
| 268 | return oid, "hash_mismatch", 0 |
| 269 | |
| 270 | new_size = len(decompressed) |
| 271 | verb = "[dry] decompress" if dry_run else "decompress" |
| 272 | async with progress_lock: |
| 273 | done_count += 1 |
| 274 | decompressed_count += 1 |
| 275 | _progress_line(f"{verb} {oid} ({len(data)} → {new_size} bytes)") |
| 276 | |
| 277 | if dry_run: |
| 278 | return oid, "decompressed", new_size |
| 279 | |
| 280 | try: |
| 281 | await backend.put(bare_oid, decompressed) |
| 282 | except Exception as exc: |
| 283 | print(f"\n ERROR re-uploading {oid}: {exc}", file=sys.stderr) |
| 284 | async with progress_lock: |
| 285 | error_count += 1 |
| 286 | return oid, "error", 0 |
| 287 | |
| 288 | return oid, "decompressed", new_size |
| 289 | |
| 290 | r2_results = await asyncio.gather(*(_process(oid) for oid in rows)) |
| 291 | |
| 292 | # DB updates for successfully decompressed objects. |
| 293 | if not dry_run: |
| 294 | for oid, status, new_size in r2_results: |
| 295 | if status != "decompressed": |
| 296 | continue |
| 297 | try: |
| 298 | await session.execute( |
| 299 | sa.update(db.MusehubObject) |
| 300 | .where(db.MusehubObject.object_id == oid) |
| 301 | .values(size_bytes=new_size) |
| 302 | ) |
| 303 | await session.commit() |
| 304 | except Exception as exc: |
| 305 | print(f"\n ERROR updating DB for {oid}: {exc}", file=sys.stderr) |
| 306 | await session.rollback() |
| 307 | async with progress_lock: |
| 308 | error_count += 1 |
| 309 | |
| 310 | # Final newline after the inline progress line. |
| 311 | if not quiet: |
| 312 | print() |
| 313 | |
| 314 | elapsed = time.monotonic() - start_time |
| 315 | prefix = "[dry-run] " if dry_run else "" |
| 316 | print( |
| 317 | f"\n{prefix}Backfill complete ({elapsed:.1f}s):\n" |
| 318 | f" {total_checked:6,} objects checked\n" |
| 319 | f" {plain_count:6,} already plain (skipped)\n" |
| 320 | f" {decompressed_count:6,} decompressed and re-uploaded\n" |
| 321 | f" {hash_mismatch_count:6,} skipped (hash mismatch after decompress)\n" |
| 322 | f" {error_count:6,} errors" |
| 323 | ) |
| 324 | return error_count |
| 325 | |
| 326 | |
| 327 | def main() -> None: |
| 328 | parser = argparse.ArgumentParser( |
| 329 | description=__doc__, |
| 330 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 331 | ) |
| 332 | parser.add_argument("--dry-run", action="store_true", |
| 333 | help="Print what would change without touching R2 or the DB") |
| 334 | parser.add_argument("--quiet", action="store_true", |
| 335 | help="Suppress per-object progress; only print final summary") |
| 336 | parser.add_argument("--batch", type=int, default=200, metavar="N", |
| 337 | help="DB page size (default 200)") |
| 338 | parser.add_argument("--concurrency", type=int, default=16, metavar="N", |
| 339 | help="Parallel R2 fetches per batch (default 16)") |
| 340 | parser.add_argument("--repo-id", default=None, metavar="UUID", |
| 341 | help="Limit to a single repo_id (for targeted testing)") |
| 342 | args = parser.parse_args() |
| 343 | |
| 344 | errors = asyncio.run(backfill( |
| 345 | dry_run=args.dry_run, |
| 346 | quiet=args.quiet, |
| 347 | batch_size=args.batch, |
| 348 | concurrency=args.concurrency, |
| 349 | repo_id=args.repo_id, |
| 350 | )) |
| 351 | sys.exit(1 if errors else 0) |
| 352 | |
| 353 | |
| 354 | if __name__ == "__main__": |
| 355 | main() |
File History
1 commit
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef
debug(push/stream): instrument O-frame decode path with INF…
Sonnet 4.6
patch
121 days ago