backends.py
python
sha256:77fc45e703f90c0d603ecb1a0ce21ff21095728ca7dd0e146eb5e966c8f9fcc9
more passing tests from full test suite fun
Human
patch
35 days ago
| 1 | """Storage backend implementations. |
| 2 | |
| 3 | All blobs (objects) are content-addressed by ``object_id``. The only backend |
| 4 | is ``BlobBackend``, which works against any S3-compatible store (Cloudflare R2, |
| 5 | MinIO, AWS S3). Configure via ``BLOB_STORAGE_BUCKET`` / ``BLOB_STORAGE_ENDPOINT`` |
| 6 | env vars. |
| 7 | |
| 8 | URI scheme: |
| 9 | ``s3://<bucket>/objects/<object_id>`` — AWS S3 / S3-compatible |
| 10 | """ |
| 11 | |
| 12 | import base64 |
| 13 | import concurrent.futures |
| 14 | import logging |
| 15 | from collections.abc import AsyncIterator |
| 16 | from typing import TYPE_CHECKING, Protocol |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from musehub.db.musehub_repo_models import MusehubObject |
| 20 | |
| 21 | from musehub.config import settings |
| 22 | from musehub.types.json_types import StrDict |
| 23 | |
| 24 | logger = logging.getLogger(__name__) |
| 25 | |
| 26 | # --------------------------------------------------------------------------- |
| 27 | # S3 concurrency constants |
| 28 | # --------------------------------------------------------------------------- |
| 29 | |
| 30 | # Size of the dedicated thread pool used for S3/R2 blocking I/O. |
| 31 | # boto3 is not async-native, so every put/get/delete goes through |
| 32 | # asyncio.get_running_loop().run_in_executor(). The default executor has |
| 33 | # min(32, cpu+4) threads — on a 4-CPU server that's 8, which throttles |
| 34 | # concurrent R2 puts to 8 even when the asyncio semaphore allows 100. |
| 35 | # Using a dedicated pool with S3_THREAD_POOL_SIZE workers lets us saturate |
| 36 | # the R2 connection pool (S3_POOL_CONNECTIONS) without queuing in the |
| 37 | # scheduler. |
| 38 | S3_THREAD_POOL_SIZE: int = 100 |
| 39 | |
| 40 | # Maximum simultaneous TCP connections per boto3 S3 client. |
| 41 | # The default (10) serializes 100 concurrent asyncio-to-thread dispatches |
| 42 | # through a 10-slot queue. 100 connections matches the thread pool size so |
| 43 | # every in-flight thread can make progress immediately. |
| 44 | S3_POOL_CONNECTIONS: int = 100 |
| 45 | |
| 46 | # Lazy-initialized thread pool — created on first S3 put, shared for the |
| 47 | # lifetime of the process. |
| 48 | _s3_thread_pool: concurrent.futures.ThreadPoolExecutor | None = None |
| 49 | |
| 50 | def _get_s3_thread_pool() -> concurrent.futures.ThreadPoolExecutor: |
| 51 | global _s3_thread_pool |
| 52 | if _s3_thread_pool is None: |
| 53 | _s3_thread_pool = concurrent.futures.ThreadPoolExecutor( |
| 54 | max_workers=S3_THREAD_POOL_SIZE, |
| 55 | thread_name_prefix="muse-s3", |
| 56 | ) |
| 57 | return _s3_thread_pool |
| 58 | |
| 59 | # Named type for the object map returned by get_batch. |
| 60 | type ObjectMap = dict[str, bytes] |
| 61 | |
| 62 | # ── Boto3 client Protocol ───────────────────────────────────────────────────── |
| 63 | |
| 64 | class _StreamingBody(Protocol): |
| 65 | def read(self, size: int = -1) -> bytes: ... |
| 66 | |
| 67 | class _S3Response(Protocol): |
| 68 | def __getitem__(self, key: str) -> _StreamingBody: ... |
| 69 | |
| 70 | class _S3Client(Protocol): |
| 71 | def head_object(self, *, Bucket: str, Key: str) -> _S3Response: ... |
| 72 | def put_object(self, *, Bucket: str, Key: str, Body: bytes) -> _S3Response: ... |
| 73 | def get_object(self, *, Bucket: str, Key: str) -> _S3Response: ... |
| 74 | def delete_object(self, *, Bucket: str, Key: str) -> _S3Response: ... |
| 75 | def generate_presigned_url( |
| 76 | self, |
| 77 | ClientMethod: str, |
| 78 | Params: StrDict, |
| 79 | ExpiresIn: int, |
| 80 | ) -> str: ... |
| 81 | |
| 82 | # ── Storage backend Protocol ────────────────────────────────────────────────── |
| 83 | |
| 84 | class StorageBackend(Protocol): |
| 85 | """Content-addressed binary store protocol. |
| 86 | |
| 87 | Storage is globally content-addressed — every method operates on |
| 88 | ``object_id`` alone; there is no per-repo namespace. |
| 89 | |
| 90 | Backends must be safe to call from async code. I/O-bound implementations |
| 91 | should use ``asyncio.to_thread`` internally rather than blocking the event |
| 92 | loop. |
| 93 | """ |
| 94 | |
| 95 | async def put(self, object_id: str, data: bytes) -> str: |
| 96 | """Persist ``data`` and return a ``storage_uri``.""" |
| 97 | ... |
| 98 | |
| 99 | async def get(self, object_id: str) -> bytes | None: |
| 100 | """Return raw bytes for the object, or ``None`` if not found.""" |
| 101 | ... |
| 102 | |
| 103 | async def get_batch(self, object_ids: list[str]) -> ObjectMap: |
| 104 | """Return ``{object_id: bytes}`` for all *object_ids* that exist. |
| 105 | |
| 106 | Default: sequential fallback via ``get()``. Implementations should |
| 107 | override this to do all reads in a single thread dispatch. |
| 108 | """ |
| 109 | result: ObjectMap = {} |
| 110 | for oid in object_ids: |
| 111 | data = await self.get(oid) |
| 112 | if data is not None: |
| 113 | result[oid] = data |
| 114 | return result |
| 115 | |
| 116 | async def exists(self, object_id: str) -> bool: |
| 117 | """Return True if the object is already stored (avoids re-upload).""" |
| 118 | ... |
| 119 | |
| 120 | async def delete(self, object_id: str) -> None: |
| 121 | """Remove an object (used on repo deletion).""" |
| 122 | ... |
| 123 | |
| 124 | def uri_for(self, object_id: str) -> str: |
| 125 | """Return the canonical storage URI without necessarily persisting.""" |
| 126 | ... |
| 127 | |
| 128 | async def stream( |
| 129 | self, object_id: str, chunk_size: int = 65536 |
| 130 | ) -> AsyncIterator[bytes]: |
| 131 | """Yield the object's bytes in chunks without buffering the full file. |
| 132 | |
| 133 | Implementations must yield an empty iterator (not raise) when the object |
| 134 | does not exist. Callers should check ``exists()`` first if a 404 is |
| 135 | needed; this method is optimised for the streaming-download hot path. |
| 136 | """ |
| 137 | ... # pragma: no cover |
| 138 | yield b"" # satisfy type checker — real impls override this |
| 139 | |
| 140 | async def presign_batch( |
| 141 | self, |
| 142 | object_ids: list[str], |
| 143 | direction: str, |
| 144 | ttl_seconds: int, |
| 145 | ) -> StrDict: |
| 146 | """Return presigned URLs keyed by object_id. Default: empty (not supported).""" |
| 147 | return {} |
| 148 | |
| 149 | class BlobBackend: |
| 150 | """AWS S3 (or S3-compatible) storage backend. |
| 151 | |
| 152 | Requires ``boto3`` to be installed and AWS credentials to be configured |
| 153 | (via environment variables, IAM instance profile, or ``~/.aws/credentials``). |
| 154 | |
| 155 | Objects are stored at: |
| 156 | ``s3://<bucket>/objects/<object_id>`` |
| 157 | |
| 158 | Storage is globally content-addressed — identical bytes are stored once |
| 159 | regardless of which repo pushed them. |
| 160 | |
| 161 | Uploads are idempotent — ``put`` calls PutObject directly without a prior |
| 162 | HeadObject check. Deduplication is handled at the database layer via the |
| 163 | ``filter-objects`` endpoint: only objects missing from the DB reach this |
| 164 | method. Skipping HeadObject halves the R2 round-trips per object and |
| 165 | doubles effective throughput. R2 PutObject is idempotent (same content → |
| 166 | same key → no-op at the storage layer), so the double-write case is safe. |
| 167 | """ |
| 168 | |
| 169 | supports_presign: bool = True |
| 170 | |
| 171 | def __init__( |
| 172 | self, |
| 173 | bucket: str | None = None, |
| 174 | region: str | None = None, |
| 175 | endpoint_url: str | None = None, |
| 176 | access_key_id: str | None = None, |
| 177 | secret_access_key: str | None = None, |
| 178 | public_endpoint_url: str | None = None, |
| 179 | ) -> None: |
| 180 | self._bucket = bucket or settings.blob_storage_bucket or settings.aws_s3_asset_bucket or "" |
| 181 | self._region = region or (settings.blob_storage_region if settings.blob_storage_bucket else settings.aws_region) |
| 182 | self._endpoint_url = endpoint_url or settings.blob_storage_endpoint |
| 183 | self._public_endpoint_url = public_endpoint_url or settings.blob_storage_public_endpoint |
| 184 | self._cdn_base_url = settings.blob_storage_cdn_base_url |
| 185 | self._access_key_id = access_key_id or settings.blob_storage_access_key_id |
| 186 | self._secret_access_key = secret_access_key or settings.blob_storage_secret_access_key |
| 187 | self._client: _S3Client | None = None |
| 188 | |
| 189 | def _get_client(self) -> _S3Client: |
| 190 | if self._client is None: |
| 191 | import boto3 |
| 192 | from botocore.config import Config |
| 193 | kwargs = { |
| 194 | "region_name": self._region or "", |
| 195 | # Increase the connection pool beyond the 10-connection default. |
| 196 | # With S3_THREAD_POOL_SIZE worker threads all making concurrent |
| 197 | # put_object calls, a 10-connection pool serializes 90 % of them |
| 198 | # into a wait queue. S3_POOL_CONNECTIONS connections lets every |
| 199 | # in-flight thread make immediate progress. |
| 200 | "config": Config( |
| 201 | max_pool_connections=S3_POOL_CONNECTIONS, |
| 202 | # Suppress optional CRC32 checksums — MinIO echoes them back, |
| 203 | # adding response header bytes. Keeping these off reduces wire |
| 204 | # overhead and keeps the 200 OK response compact. |
| 205 | request_checksum_calculation="when_required", |
| 206 | response_checksum_validation="when_required", |
| 207 | ), |
| 208 | } |
| 209 | if self._endpoint_url: |
| 210 | kwargs["endpoint_url"] = self._endpoint_url |
| 211 | if self._access_key_id and self._secret_access_key: |
| 212 | kwargs["aws_access_key_id"] = self._access_key_id |
| 213 | kwargs["aws_secret_access_key"] = self._secret_access_key |
| 214 | self._client = boto3.client("s3", **kwargs) |
| 215 | # Python 3.14 regression: when botocore sends Expect: 100-continue, |
| 216 | # MinIO sends 100 Continue then 200 OK in the same TCP segment. |
| 217 | # Python 3.14's http.client sees the 200 OK bytes as unparsed data |
| 218 | # after the 100 Continue headers → MissingHeaderBodySeparatorDefect. |
| 219 | # urllib3 catches the defect but returns the connection to the pool |
| 220 | # without draining the response body (dirty connection). The next PUT |
| 221 | # on the same pool reuses this connection; MinIO reads the leftover |
| 222 | # bytes as a new request, becomes confused, and closes the connection |
| 223 | # after its 30s keep-alive timeout. Fix: remove Expect: 100-continue |
| 224 | # so botocore sends headers+body in one shot → single 200 OK → clean. |
| 225 | def _remove_expect_continue(request: "botocore.awsrequest.AWSPreparedRequest", **kwargs: str) -> None: |
| 226 | request.headers.pop("Expect", None) |
| 227 | self._client.meta.events.register( |
| 228 | "before-send.s3.PutObject", |
| 229 | _remove_expect_continue, |
| 230 | ) |
| 231 | return self._client |
| 232 | |
| 233 | def _key(self, object_id: str) -> str: |
| 234 | # object_id is canonical algo:hex (e.g. "sha256:<64-hex>"). |
| 235 | # Colons are valid in S3/R2 keys — no substitution needed. |
| 236 | return f"objects/{object_id}" |
| 237 | |
| 238 | def uri_for(self, object_id: str) -> str: |
| 239 | return f"s3://{self._bucket}/{self._key(object_id)}" |
| 240 | |
| 241 | async def put(self, object_id: str, data: bytes) -> str: |
| 242 | import asyncio |
| 243 | key = self._key(object_id) |
| 244 | import logging as _log |
| 245 | _log.getLogger(__name__).debug("BlobBackend.put: oid=%s key=%s size=%d", object_id, key, len(data)) |
| 246 | client = self._get_client() |
| 247 | loop = asyncio.get_running_loop() |
| 248 | await loop.run_in_executor(_get_s3_thread_pool(), self._s3_put, client, key, data) |
| 249 | return self.uri_for(object_id) |
| 250 | |
| 251 | def _s3_put(self, client: _S3Client, key: str, data: bytes) -> None: |
| 252 | # No HeadObject check — deduplication is handled upstream by the |
| 253 | # filter-objects endpoint. Only objects confirmed missing from the DB |
| 254 | # reach this call, so the check is redundant and wastes one R2 |
| 255 | # round-trip per object. PutObject is idempotent: uploading the same |
| 256 | # content-addressed bytes twice is a safe no-op at the storage layer. |
| 257 | client.put_object(Bucket=self._bucket, Key=key, Body=data) |
| 258 | |
| 259 | async def get(self, object_id: str) -> bytes | None: |
| 260 | import asyncio |
| 261 | c = self._get_client() |
| 262 | key = self._key(object_id) |
| 263 | try: |
| 264 | def _do_get() -> _S3Response: |
| 265 | return c.get_object(Bucket=self._bucket, Key=key) |
| 266 | result = await asyncio.to_thread(_do_get) |
| 267 | data: bytes = result["Body"].read() |
| 268 | return data |
| 269 | except Exception: |
| 270 | return None |
| 271 | |
| 272 | async def get_batch(self, object_ids: list[str]) -> ObjectMap: |
| 273 | """Fetch all objects in parallel via asyncio.gather. |
| 274 | |
| 275 | The base-class fallback is sequential — one S3 round-trip per object. |
| 276 | On R2/S3 in production that means a 100+ object clone hits Cloudflare's |
| 277 | 100-second origin timeout. Overriding here dispatches all fetches |
| 278 | concurrently so total latency ≈ max(single object latency). |
| 279 | """ |
| 280 | import asyncio |
| 281 | |
| 282 | results = await asyncio.gather(*(self.get(oid) for oid in object_ids)) |
| 283 | return { |
| 284 | oid: data |
| 285 | for oid, data in zip(object_ids, results) |
| 286 | if data is not None |
| 287 | } |
| 288 | |
| 289 | async def exists(self, object_id: str) -> bool: |
| 290 | import asyncio |
| 291 | c = self._get_client() |
| 292 | key = self._key(object_id) |
| 293 | |
| 294 | def _head() -> bool: |
| 295 | try: |
| 296 | c.head_object(Bucket=self._bucket, Key=key) |
| 297 | return True |
| 298 | except Exception: |
| 299 | return False |
| 300 | |
| 301 | loop = asyncio.get_running_loop() |
| 302 | return await loop.run_in_executor(_get_s3_thread_pool(), _head) |
| 303 | |
| 304 | async def delete(self, object_id: str) -> None: |
| 305 | import asyncio |
| 306 | key = self._key(object_id) |
| 307 | c = self._get_client() |
| 308 | def _do_delete() -> None: |
| 309 | c.delete_object(Bucket=self._bucket, Key=key) |
| 310 | await asyncio.to_thread(_do_delete) |
| 311 | |
| 312 | async def stream( |
| 313 | self, object_id: str, chunk_size: int = 65536 |
| 314 | ) -> AsyncIterator[bytes]: |
| 315 | """Yield S3 object content in ``chunk_size`` byte chunks. |
| 316 | |
| 317 | Uses boto3's streaming body so the full object is never loaded into |
| 318 | memory — the HTTP response body is consumed incrementally. |
| 319 | """ |
| 320 | import asyncio |
| 321 | |
| 322 | c = self._get_client() |
| 323 | key = self._key(object_id) |
| 324 | body: _StreamingBody | None = None |
| 325 | try: |
| 326 | def _get_stream() -> _StreamingBody: |
| 327 | return c.get_object(Bucket=self._bucket, Key=key)["Body"] |
| 328 | body = await asyncio.to_thread(_get_stream) |
| 329 | except Exception: |
| 330 | return # object not found — yield nothing |
| 331 | |
| 332 | assert body is not None |
| 333 | def _read_chunk() -> bytes: |
| 334 | return body.read(chunk_size) |
| 335 | |
| 336 | while True: |
| 337 | chunk = await asyncio.to_thread(_read_chunk) |
| 338 | if not chunk: |
| 339 | break |
| 340 | yield chunk |
| 341 | |
| 342 | def _rewrite_presign_url(self, url: str) -> str: |
| 343 | """Replace the internal endpoint host with the public endpoint host. |
| 344 | |
| 345 | boto3 embeds whatever endpoint_url the client was initialised with into |
| 346 | every presigned URL. In local dev that's http://minio:9000 (Docker |
| 347 | internal), which is unreachable from the host machine. When |
| 348 | BLOB_STORAGE_PUBLIC_ENDPOINT is set, swap in that host so external clients can |
| 349 | PUT/GET directly. |
| 350 | """ |
| 351 | if not self._public_endpoint_url or not self._endpoint_url: |
| 352 | return url |
| 353 | return url.replace(self._endpoint_url, self._public_endpoint_url, 1) |
| 354 | |
| 355 | def _mpack_key(self, mpack_key: str) -> str: |
| 356 | return f"mpacks/{mpack_key}" |
| 357 | |
| 358 | async def presign_mpack_put(self, mpack_key: str, ttl_seconds: int = 3600) -> str: |
| 359 | """Return a presigned PUT URL for an mpack blob (stored under mpacks/ prefix).""" |
| 360 | import asyncio |
| 361 | c = self._get_client() |
| 362 | key = self._mpack_key(mpack_key) |
| 363 | logger.warning("[presign_mpack_put] mpack_key=%s s3_key=%s bucket=%s", mpack_key, key, self._bucket) |
| 364 | |
| 365 | def _presign() -> str: |
| 366 | raw_url = c.generate_presigned_url( |
| 367 | "put_object", |
| 368 | Params={"Bucket": self._bucket, "Key": key, "ContentType": "application/x-muse-pack"}, |
| 369 | ExpiresIn=ttl_seconds, |
| 370 | ) |
| 371 | logger.warning("[presign_mpack_put] raw_url (before rewrite) = %s", raw_url) |
| 372 | return raw_url |
| 373 | |
| 374 | loop = asyncio.get_running_loop() |
| 375 | final_url = self._rewrite_presign_url(await loop.run_in_executor(_get_s3_thread_pool(), _presign)) |
| 376 | logger.warning("[presign_mpack_put] final_url (after rewrite) = %s", final_url) |
| 377 | return final_url |
| 378 | |
| 379 | async def presign_mpack_get(self, mpack_key: str, ttl_seconds: int = 3600) -> str: |
| 380 | """Return a presigned GET URL for an mpack blob (stored under mpacks/ prefix). |
| 381 | |
| 382 | When ``blob_storage_cdn_base_url`` is configured the presigned URL host is |
| 383 | replaced with the CDN base URL so repeat clones hit the Cloudflare edge |
| 384 | cache instead of R2/MinIO directly. |
| 385 | """ |
| 386 | import asyncio |
| 387 | c = self._get_client() |
| 388 | key = self._mpack_key(mpack_key) |
| 389 | |
| 390 | def _presign() -> str: |
| 391 | return c.generate_presigned_url( |
| 392 | "get_object", |
| 393 | Params={"Bucket": self._bucket, "Key": key}, |
| 394 | ExpiresIn=ttl_seconds, |
| 395 | ) |
| 396 | |
| 397 | loop = asyncio.get_running_loop() |
| 398 | url = self._rewrite_presign_url(await loop.run_in_executor(_get_s3_thread_pool(), _presign)) |
| 399 | if self._cdn_base_url: |
| 400 | from urllib.parse import urlparse, urlunparse |
| 401 | parsed = urlparse(url) |
| 402 | cdn = urlparse(self._cdn_base_url) |
| 403 | url = urlunparse(parsed._replace(scheme=cdn.scheme, netloc=cdn.netloc)) |
| 404 | return url |
| 405 | |
| 406 | async def get_range(self, mpack_key: str, byte_offset: int, byte_length: int) -> bytes | None: |
| 407 | """Fetch a byte range from an mpack using an S3 Range GET. |
| 408 | |
| 409 | Issues ``Range: bytes=<byte_offset>-<byte_offset+byte_length-1>`` so |
| 410 | only the requested slice is transferred — O(object_size) instead of |
| 411 | O(mpack_size). Returns ``None`` if the mpack or range is not found. |
| 412 | """ |
| 413 | import asyncio |
| 414 | c = self._get_client() |
| 415 | key = self._mpack_key(mpack_key) |
| 416 | range_header = f"bytes={byte_offset}-{byte_offset + byte_length - 1}" |
| 417 | |
| 418 | def _get() -> bytes | None: |
| 419 | try: |
| 420 | resp = c.get_object(Bucket=self._bucket, Key=key, Range=range_header) |
| 421 | return resp["Body"].read() |
| 422 | except Exception: |
| 423 | return None |
| 424 | |
| 425 | return await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _get) |
| 426 | |
| 427 | async def exists_mpack(self, mpack_key: str) -> bool: |
| 428 | """Return True if the mpack exists in storage.""" |
| 429 | import asyncio |
| 430 | c = self._get_client() |
| 431 | key = self._mpack_key(mpack_key) |
| 432 | |
| 433 | def _head() -> bool: |
| 434 | try: |
| 435 | c.head_object(Bucket=self._bucket, Key=key) |
| 436 | return True |
| 437 | except Exception: |
| 438 | return False |
| 439 | |
| 440 | return await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _head) |
| 441 | |
| 442 | async def get_mpack(self, mpack_key: str) -> bytes | None: |
| 443 | """Fetch the raw mpack bytes from the mpacks/ prefix.""" |
| 444 | import asyncio |
| 445 | import hashlib as _hashlib |
| 446 | c = self._get_client() |
| 447 | key = self._mpack_key(mpack_key) |
| 448 | logger.warning("[get_mpack] mpack_key=%s s3_key=%s bucket=%s", mpack_key, key, self._bucket) |
| 449 | |
| 450 | def _get() -> bytes | None: |
| 451 | try: |
| 452 | resp = c.get_object(Bucket=self._bucket, Key=key) |
| 453 | etag = resp.get("ETag", "") |
| 454 | content_length = resp.get("ContentLength", "?") |
| 455 | data = resp["Body"].read() |
| 456 | actual_sha256 = "sha256:" + _hashlib.sha256(data).hexdigest() |
| 457 | logger.warning( |
| 458 | "[get_mpack] fetched %d bytes etag=%s content_length=%s sha256(data)=%s matches_key=%s", |
| 459 | len(data), etag, content_length, actual_sha256, actual_sha256 == mpack_key, |
| 460 | ) |
| 461 | return data |
| 462 | except Exception as exc: |
| 463 | logger.warning("[get_mpack] FAILED key=%s exc=%s", key, exc) |
| 464 | return None |
| 465 | |
| 466 | return await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _get) |
| 467 | |
| 468 | async def put_mpack(self, mpack_key: str, data: bytes) -> None: |
| 469 | """Write raw mpack bytes directly under the mpacks/ prefix. |
| 470 | |
| 471 | Mpacks are content-addressed (key == sha256(bytes)) so the key never |
| 472 | changes after creation. ``Cache-Control: public, max-age=31536000, immutable`` |
| 473 | tells Cloudflare it is safe to cache the response indefinitely at the edge. |
| 474 | """ |
| 475 | import asyncio |
| 476 | c = self._get_client() |
| 477 | key = self._mpack_key(mpack_key) |
| 478 | |
| 479 | def _put() -> None: |
| 480 | c.put_object( |
| 481 | Bucket=self._bucket, |
| 482 | Key=key, |
| 483 | Body=data, |
| 484 | ContentType="application/x-muse-pack", |
| 485 | CacheControl="public, max-age=31536000, immutable", |
| 486 | ) |
| 487 | |
| 488 | await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _put) |
| 489 | |
| 490 | async def quarantine_mpack(self, mpack_key: str) -> None: |
| 491 | """Move an mpack blob from mpacks/ to quarantine/ prefix. |
| 492 | |
| 493 | Called by process_mpack_index_job on MPackValidationError to |
| 494 | isolate suspect content from normal mpack storage. |
| 495 | """ |
| 496 | import asyncio |
| 497 | c = self._get_client() |
| 498 | src_key = self._mpack_key(mpack_key) |
| 499 | dst_key = f"quarantine/{mpack_key}" |
| 500 | |
| 501 | def _move() -> None: |
| 502 | try: |
| 503 | c.copy_object( |
| 504 | Bucket=self._bucket, |
| 505 | CopySource={"Bucket": self._bucket, "Key": src_key}, |
| 506 | Key=dst_key, |
| 507 | ) |
| 508 | c.delete_object(Bucket=self._bucket, Key=src_key) |
| 509 | except Exception: |
| 510 | pass # best-effort; do not mask the validation error |
| 511 | |
| 512 | await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _move) |
| 513 | |
| 514 | async def quarantine_object(self, object_id: str) -> None: |
| 515 | """Move an object blob from objects/ to quarantine/ prefix. |
| 516 | |
| 517 | Called by the DMCA takedown endpoint to remove a blob from public access. |
| 518 | Best-effort: silently no-ops if the object is not present. |
| 519 | """ |
| 520 | import asyncio |
| 521 | c = self._get_client() |
| 522 | src_key = self._key(object_id) |
| 523 | dst_key = f"quarantine/{object_id}" |
| 524 | |
| 525 | def _move() -> None: |
| 526 | try: |
| 527 | c.copy_object( |
| 528 | Bucket=self._bucket, |
| 529 | CopySource={"Bucket": self._bucket, "Key": src_key}, |
| 530 | Key=dst_key, |
| 531 | ) |
| 532 | c.delete_object(Bucket=self._bucket, Key=src_key) |
| 533 | except Exception: |
| 534 | pass |
| 535 | |
| 536 | await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _move) |
| 537 | |
| 538 | async def presign_put(self, object_id: str, ttl_seconds: int = 3600) -> str: |
| 539 | """Return a presigned PUT URL for *object_id* valid for *ttl_seconds*. |
| 540 | |
| 541 | The client PUTs raw bytes directly to R2 using this URL, bypassing the |
| 542 | API server and Cloudflare entirely. After all uploads complete, the |
| 543 | client calls ``POST /push/objects/confirm`` to register the objects in |
| 544 | the MuseHub DB. |
| 545 | """ |
| 546 | import asyncio |
| 547 | |
| 548 | c = self._get_client() |
| 549 | key = self._key(object_id) |
| 550 | |
| 551 | def _presign() -> str: |
| 552 | return c.generate_presigned_url( |
| 553 | "put_object", |
| 554 | Params={"Bucket": self._bucket, "Key": key, "ContentType": "application/octet-stream"}, |
| 555 | ExpiresIn=ttl_seconds, |
| 556 | ) |
| 557 | |
| 558 | loop = asyncio.get_running_loop() |
| 559 | return self._rewrite_presign_url(await loop.run_in_executor(_get_s3_thread_pool(), _presign)) |
| 560 | |
| 561 | async def presign_get(self, object_id: str, ttl_seconds: int = 3600) -> str: |
| 562 | """Return a presigned GET URL for *object_id* valid for *ttl_seconds*. |
| 563 | |
| 564 | The client GETs raw bytes directly from R2 using this URL, bypassing |
| 565 | the API server. Used in the fetch path when the backend is R2. |
| 566 | """ |
| 567 | import asyncio |
| 568 | |
| 569 | c = self._get_client() |
| 570 | key = self._key(object_id) |
| 571 | |
| 572 | def _presign() -> str: |
| 573 | return c.generate_presigned_url( |
| 574 | "get_object", |
| 575 | Params={"Bucket": self._bucket, "Key": key}, |
| 576 | ExpiresIn=ttl_seconds, |
| 577 | ) |
| 578 | |
| 579 | loop = asyncio.get_running_loop() |
| 580 | return self._rewrite_presign_url( |
| 581 | await loop.run_in_executor(_get_s3_thread_pool(), _presign) |
| 582 | ) |
| 583 | |
| 584 | async def presign_batch( |
| 585 | self, |
| 586 | object_ids: list[str], |
| 587 | direction: str, |
| 588 | ttl_seconds: int, |
| 589 | ) -> StrDict: |
| 590 | """Return presigned URLs for *object_ids* in a single thread dispatch.""" |
| 591 | import asyncio |
| 592 | |
| 593 | c = self._get_client() |
| 594 | method = "put_object" if direction == "put" else "get_object" |
| 595 | |
| 596 | def _presign_all() -> list[str]: |
| 597 | return [ |
| 598 | c.generate_presigned_url( |
| 599 | method, |
| 600 | Params={"Bucket": self._bucket, "Key": self._key(oid)}, |
| 601 | ExpiresIn=ttl_seconds, |
| 602 | ) |
| 603 | for oid in object_ids |
| 604 | ] |
| 605 | |
| 606 | urls = await asyncio.to_thread(_presign_all) |
| 607 | return {oid: self._rewrite_presign_url(url) for oid, url in zip(object_ids, urls)} |
| 608 | |
| 609 | class MemoryBackend: |
| 610 | """In-memory blob backend for tests — no external services required.""" |
| 611 | |
| 612 | supports_presign: bool = False |
| 613 | |
| 614 | def __init__(self) -> None: |
| 615 | self._store: dict[str, bytes] = {} |
| 616 | self._mpacks: dict[str, bytes] = {} |
| 617 | |
| 618 | def uri_for(self, object_id: str) -> str: |
| 619 | # Use s3:// prefix so read_object_bytes routes through get_backend().get(), |
| 620 | # which MemoryBackend implements in-memory without external services. |
| 621 | return f"s3://memory-bucket/objects/{object_id}" |
| 622 | |
| 623 | async def put(self, object_id: str, data: bytes) -> str: |
| 624 | self._store[object_id] = data |
| 625 | return self.uri_for(object_id) |
| 626 | |
| 627 | async def get(self, object_id: str) -> bytes | None: |
| 628 | return self._store.get(object_id) |
| 629 | |
| 630 | async def get_batch(self, object_ids: list[str]) -> "ObjectMap": |
| 631 | return {oid: self._store[oid] for oid in object_ids if oid in self._store} |
| 632 | |
| 633 | async def exists(self, object_id: str) -> bool: |
| 634 | return object_id in self._store |
| 635 | |
| 636 | async def delete(self, object_id: str) -> None: |
| 637 | self._store.pop(object_id, None) |
| 638 | |
| 639 | async def stream(self, object_id: str, chunk_size: int = 65536) -> "AsyncIterator[bytes]": |
| 640 | data = self._store.get(object_id) |
| 641 | if data: |
| 642 | yield data |
| 643 | |
| 644 | async def presign_batch(self, object_ids: list[str], direction: str, ttl_seconds: int) -> "StrDict": |
| 645 | return {} |
| 646 | |
| 647 | async def presign_get(self, object_id: str, ttl_seconds: int = 3600) -> str: |
| 648 | return f"memory://objects/{object_id}?op=get&ttl={ttl_seconds}" |
| 649 | |
| 650 | # ── mpack interface ─────────────────────────────────────────────────────── |
| 651 | |
| 652 | async def put_mpack(self, mpack_key: str, data: bytes) -> None: |
| 653 | self._mpacks[mpack_key] = data |
| 654 | |
| 655 | async def get_mpack(self, mpack_key: str) -> bytes | None: |
| 656 | return self._mpacks.get(mpack_key) |
| 657 | |
| 658 | async def exists_mpack(self, mpack_key: str) -> bool: |
| 659 | return mpack_key in self._mpacks |
| 660 | |
| 661 | async def get_range(self, mpack_key: str, byte_offset: int, byte_length: int) -> bytes | None: |
| 662 | data = self._mpacks.get(mpack_key) |
| 663 | if data is None: |
| 664 | return None |
| 665 | return data[byte_offset: byte_offset + byte_length] |
| 666 | |
| 667 | async def presign_mpack_put(self, mpack_key: str, ttl_seconds: int = 3600) -> str: |
| 668 | return f"memory://mpacks/{mpack_key}?op=put" |
| 669 | |
| 670 | async def presign_mpack_get(self, mpack_key: str, ttl_seconds: int = 3600) -> str: |
| 671 | return f"memory://mpacks/{mpack_key}?op=get" |
| 672 | |
| 673 | async def quarantine_mpack(self, mpack_key: str) -> None: |
| 674 | self._mpacks.pop(mpack_key, None) |
| 675 | |
| 676 | async def quarantine_object(self, object_id: str) -> None: |
| 677 | self._store.pop(object_id, None) |
| 678 | |
| 679 | |
| 680 | def _get_backend_impl() -> StorageBackend: |
| 681 | """Real implementation — not patched by conftest. Tests that verify the |
| 682 | selection logic call this directly.""" |
| 683 | if settings.blob_storage_bucket: |
| 684 | return BlobBackend() |
| 685 | if settings.aws_s3_asset_bucket: |
| 686 | return BlobBackend() |
| 687 | raise RuntimeError( |
| 688 | "No storage backend configured. Set BLOB_STORAGE_BUCKET " |
| 689 | "or AWS_S3_ASSET_BUCKET." |
| 690 | ) |
| 691 | |
| 692 | |
| 693 | def get_backend() -> StorageBackend: |
| 694 | """Return the configured BlobBackend.""" |
| 695 | return _get_backend_impl() |
| 696 | |
| 697 | async def read_object_bytes( |
| 698 | obj: "MusehubObject", |
| 699 | session: "AsyncSession | None" = None, |
| 700 | ) -> bytes | None: |
| 701 | """Return raw bytes for any MusehubObject. |
| 702 | |
| 703 | Resolution order: |
| 704 | 1. ``content_cache`` — in-memory, no I/O |
| 705 | 2. ``storage_uri`` starts with ``s3://`` → BlobBackend.get() |
| 706 | 3. ``storage_uri`` starts with ``mpack://`` → byte-range GET via |
| 707 | musehub_mpack_index (fast) or full mpack download (slow fallback) |
| 708 | 4. Nothing available → None |
| 709 | |
| 710 | Pass ``session`` when available so byte ranges can be looked up from |
| 711 | musehub_mpack_index — this avoids downloading the full mpack. |
| 712 | """ |
| 713 | content_cache = getattr(obj, "content_cache", None) |
| 714 | if content_cache is not None: |
| 715 | return content_cache |
| 716 | |
| 717 | storage_uri: str = getattr(obj, "storage_uri", "") or "" |
| 718 | if storage_uri.startswith("s3://"): |
| 719 | return await get_backend().get(getattr(obj, "object_id", "")) |
| 720 | |
| 721 | if storage_uri.startswith("mpack://"): |
| 722 | mpack_id = storage_uri[len("mpack://"):] |
| 723 | oid: str = getattr(obj, "object_id", "") |
| 724 | byte_offset: int | None = getattr(obj, "byte_offset", None) |
| 725 | byte_length: int | None = getattr(obj, "byte_length", None) |
| 726 | |
| 727 | # Look up byte range from musehub_mpack_index when session is available. |
| 728 | if (byte_offset is None or byte_length is None) and oid and session is not None: |
| 729 | from musehub.db.musehub_repo_models import MusehubMPackIndex |
| 730 | from sqlalchemy import select as _sa_select |
| 731 | # Prefer the specific push mpack this object was stored in. |
| 732 | _idx = (await session.execute( |
| 733 | _sa_select(MusehubMPackIndex).where( |
| 734 | MusehubMPackIndex.entity_id == oid, |
| 735 | MusehubMPackIndex.mpack_id == mpack_id, |
| 736 | MusehubMPackIndex.entity_type == "object", |
| 737 | ).limit(1) |
| 738 | )).scalar_one_or_none() |
| 739 | if _idx is not None: |
| 740 | byte_offset = _idx.byte_offset |
| 741 | byte_length = _idx.byte_length |
| 742 | else: |
| 743 | # The original push mpack may not have index entries (e.g. the |
| 744 | # prebuilt clone mpack superseded it in the index). Fall back |
| 745 | # to any indexed location for this object so we can serve it |
| 746 | # via range GET without downloading the entire push mpack. |
| 747 | _any_idx = (await session.execute( |
| 748 | _sa_select(MusehubMPackIndex).where( |
| 749 | MusehubMPackIndex.entity_id == oid, |
| 750 | MusehubMPackIndex.entity_type == "object", |
| 751 | MusehubMPackIndex.byte_offset.isnot(None), |
| 752 | MusehubMPackIndex.byte_length.isnot(None), |
| 753 | ).limit(1) |
| 754 | )).scalar_one_or_none() |
| 755 | if _any_idx is not None: |
| 756 | mpack_id = _any_idx.mpack_id |
| 757 | byte_offset = _any_idx.byte_offset |
| 758 | byte_length = _any_idx.byte_length |
| 759 | |
| 760 | # Fast path: byte-range GET — O(object_size), not O(mpack_size). |
| 761 | if byte_offset is not None and byte_length is not None: |
| 762 | return await get_backend().get_range(mpack_id, byte_offset, byte_length) |
| 763 | |
| 764 | mpack_raw = await get_backend().get_mpack(mpack_id) |
| 765 | if mpack_raw is None: |
| 766 | return None |
| 767 | import zstandard as _zstd |
| 768 | _dctx = _zstd.ZstdDecompressor() |
| 769 | if mpack_raw[:4] != b"MUSE": |
| 770 | raise ValueError(f"mpack is not MUSE binary format (got {mpack_raw[:4]!r})") |
| 771 | from muse.core.mpack import parse_wire_mpack as _parse_wire_mpack |
| 772 | payload = _parse_wire_mpack(mpack_raw) |
| 773 | for entry in payload.get("blobs", []): |
| 774 | if entry.get("object_id") == oid: |
| 775 | content = entry.get("content") or b"" |
| 776 | if not isinstance(content, bytes): |
| 777 | content = bytes(content) |
| 778 | _ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" |
| 779 | if (entry.get("encoding") == "zstd" or content[:4] == _ZSTD_MAGIC) and content: |
| 780 | content = _dctx.decompress(content) |
| 781 | return content |
| 782 | return None |
| 783 | |
| 784 | return None |
| 785 | |
| 786 | def decode_b64(b64_str: str) -> bytes: |
| 787 | """Decode a base64 string (with or without padding) into bytes.""" |
| 788 | padded = b64_str + "=" * (-len(b64_str) % 4) |
| 789 | return base64.b64decode(padded) |
File History
8 commits
sha256:77fc45e703f90c0d603ecb1a0ce21ff21095728ca7dd0e146eb5e966c8f9fcc9
more passing tests from full test suite fun
Human
patch
35 days ago
sha256:92528ae07d0e1239d87fd5fd1f439e8fbb49c9778a9a400bc4a736073fb28316
feat: byte-range blob reads, file attribution DAG walk, bra…
Sonnet 4.6
minor
⚠
43 days ago
sha256:f3995ec2c05c9c34b0e4d6e96349a811d0117a1c51d78096d757998ccb3c0520
fix: blobs only in S3/mpack — remove commit/snapshot indivi…
Sonnet 4.6
patch
45 days ago
sha256:d33e88196634b22b27811fa156c7d32a2311173b9daa7c7be91e2829372a24f5
fix: sign ContentType into presigned PUT URLs + wire fetch …
Sonnet 4.6
minor
⚠
47 days ago
sha256:302574ddba13c9a20694c0fb051176eef4896f943b63bc458df886633b1bfcd6
feat: mpack byte-range index — store byte_offset/byte_lengt…
Sonnet 4.6
minor
⚠
48 days ago
sha256:39e9c4e6f2134da0732e6983268a218178973936f8d7ca03c91f2b5ad42133c8
fix: use read_object_bytes in blob viewer; add zstd magic d…
Sonnet 4.6
patch
48 days ago
sha256:ef10830ce231e0a20efcb0e2586cb879471247e916616e6fdd0d51df459e2595
fix: typing audit — 0 violations, 0 untyped defs across all…
Sonnet 4.6
minor
⚠
49 days ago
sha256:1b52d53ca7bae6c2b32297aaa798cd9f25e8241f457e8de4186aded84ab9c4a1
debug: log full presign URL and get_mpack key/bytes to isol…
Sonnet 4.6
minor
⚠
50 days ago