.chainlink/issues.db

reviewed

This is a binary file and will not be displayed.

.planning/roadmap/v0.3/03_human-review-assessment.md .planning/roadmap/v0.2/03_human-review-assessment.md

reviewed

+4

CHANGELOG.md

reviewed

··· 25 25 - **Comprehensive integration test suite**: 593 tests covering E2E flows, error handling, edge cases 26 26 27 27 ### Changed 28 28 + - Implement PDSBlobStore for ATProto blob storage (#380) 29 29 + - Add tests for PDSBlobStore and BlobSource (#383) 30 30 + - Add BlobSource for reading PDS blobs as DataSource (#382) 31 31 + - Create PDSBlobStore class in atmosphere module (#381) 28 32 - Investigate Redis index entry expiration/reset issue (#376) 29 33 - Audit codebase for xs/@property vs list_xs() convention (#377) 30 34 - Evaluate PackableSample → Packable protocol migration (#375)

+1

src/atdata/__init__.py

reviewed

··· 68 68 from ._sources import ( 69 69 URLSource as URLSource, 70 70 S3Source as S3Source, 71 71 + BlobSource as BlobSource, 71 72 ) 72 73 73 74 from ._schema_codec import (

+172

src/atdata/_sources.py

reviewed

··· 337 337 ) 338 338 339 339 340 340 + @dataclass 341 341 + class BlobSource: 342 342 + """Data source for ATProto PDS blob storage. 343 343 + 344 344 + Streams dataset shards stored as blobs on an ATProto Personal Data Server. 345 345 + Each shard is identified by a blob reference containing the DID and CID. 346 346 + 347 347 + This source resolves blob references to HTTP URLs and streams the content 348 348 + directly, supporting efficient iteration over shards without downloading 349 349 + everything upfront. 350 350 + 351 351 + Attributes: 352 352 + blob_refs: List of blob reference dicts with 'did' and 'cid' keys. 353 353 + pds_endpoint: Optional PDS endpoint URL. If not provided, resolved from DID. 354 354 + 355 355 + Example: 356 356 + :: 357 357 + 358 358 + >>> source = BlobSource( 359 359 + ... blob_refs=[ 360 360 + ... {"did": "did:plc:abc123", "cid": "bafyrei..."}, 361 361 + ... {"did": "did:plc:abc123", "cid": "bafyrei..."}, 362 362 + ... ], 363 363 + ... ) 364 364 + >>> for shard_id, stream in source.shards: 365 365 + ... process(stream) 366 366 + """ 367 367 + 368 368 + blob_refs: list[dict[str, str]] 369 369 + pds_endpoint: str | None = None 370 370 + _endpoint_cache: dict[str, str] = field(default_factory=dict, repr=False, compare=False) 371 371 + 372 372 + def _resolve_pds_endpoint(self, did: str) -> str: 373 373 + """Resolve PDS endpoint for a DID, with caching.""" 374 374 + if did in self._endpoint_cache: 375 375 + return self._endpoint_cache[did] 376 376 + 377 377 + if self.pds_endpoint: 378 378 + self._endpoint_cache[did] = self.pds_endpoint 379 379 + return self.pds_endpoint 380 380 + 381 381 + import requests 382 382 + 383 383 + # Resolve via plc.directory 384 384 + if did.startswith("did:plc:"): 385 385 + plc_url = f"https://plc.directory/{did}" 386 386 + response = requests.get(plc_url, timeout=10) 387 387 + response.raise_for_status() 388 388 + doc = response.json() 389 389 + 390 390 + for service in doc.get("service", []): 391 391 + if service.get("type") == "AtprotoPersonalDataServer": 392 392 + endpoint = service.get("serviceEndpoint", "") 393 393 + self._endpoint_cache[did] = endpoint 394 394 + return endpoint 395 395 + 396 396 + raise ValueError(f"Could not resolve PDS endpoint for {did}") 397 397 + 398 398 + def _get_blob_url(self, did: str, cid: str) -> str: 399 399 + """Get HTTP URL for fetching a blob.""" 400 400 + endpoint = self._resolve_pds_endpoint(did) 401 401 + return f"{endpoint}/xrpc/com.atproto.sync.getBlob?did={did}&cid={cid}" 402 402 + 403 403 + def _make_shard_id(self, ref: dict[str, str]) -> str: 404 404 + """Create shard identifier from blob reference.""" 405 405 + return f"at://{ref['did']}/blob/{ref['cid']}" 406 406 + 407 407 + def list_shards(self) -> list[str]: 408 408 + """Return list of AT URI-style shard identifiers.""" 409 409 + return [self._make_shard_id(ref) for ref in self.blob_refs] 410 410 + 411 411 + @property 412 412 + def shards(self) -> Iterator[tuple[str, IO[bytes]]]: 413 413 + """Lazily yield (at_uri, stream) pairs for each shard. 414 414 + 415 415 + Fetches blobs via HTTP from the PDS and yields streaming responses. 416 416 + 417 417 + Yields: 418 418 + Tuple of (at://did/blob/cid URI, streaming response body). 419 419 + """ 420 420 + import requests 421 421 + 422 422 + for ref in self.blob_refs: 423 423 + did = ref["did"] 424 424 + cid = ref["cid"] 425 425 + url = self._get_blob_url(did, cid) 426 426 + 427 427 + response = requests.get(url, stream=True, timeout=60) 428 428 + response.raise_for_status() 429 429 + 430 430 + shard_id = self._make_shard_id(ref) 431 431 + # Wrap response in a file-like object 432 432 + yield shard_id, response.raw 433 433 + 434 434 + def open_shard(self, shard_id: str) -> IO[bytes]: 435 435 + """Open a single shard by its AT URI. 436 436 + 437 437 + Args: 438 438 + shard_id: AT URI of the shard (at://did/blob/cid). 439 439 + 440 440 + Returns: 441 441 + Streaming response body for reading the blob. 442 442 + 443 443 + Raises: 444 444 + KeyError: If shard_id is not in list_shards(). 445 445 + ValueError: If shard_id format is invalid. 446 446 + """ 447 447 + if shard_id not in self.list_shards(): 448 448 + raise KeyError(f"Shard not found: {shard_id}") 449 449 + 450 450 + # Parse at://did/blob/cid 451 451 + if not shard_id.startswith("at://"): 452 452 + raise ValueError(f"Invalid shard ID format: {shard_id}") 453 453 + 454 454 + parts = shard_id[5:].split("/") # Remove 'at://' 455 455 + if len(parts) != 3 or parts[1] != "blob": 456 456 + raise ValueError(f"Invalid blob URI format: {shard_id}") 457 457 + 458 458 + did, _, cid = parts 459 459 + url = self._get_blob_url(did, cid) 460 460 + 461 461 + import requests 462 462 + response = requests.get(url, stream=True, timeout=60) 463 463 + response.raise_for_status() 464 464 + return response.raw 465 465 + 466 466 + @classmethod 467 467 + def from_refs( 468 468 + cls, 469 469 + refs: list[dict], 470 470 + *, 471 471 + pds_endpoint: str | None = None, 472 472 + ) -> "BlobSource": 473 473 + """Create BlobSource from blob reference dicts. 474 474 + 475 475 + Accepts blob references in the format returned by upload_blob: 476 476 + ``{"$type": "blob", "ref": {"$link": "cid"}, ...}`` 477 477 + 478 478 + Also accepts simplified format: ``{"did": "...", "cid": "..."}`` 479 479 + 480 480 + Args: 481 481 + refs: List of blob reference dicts. 482 482 + pds_endpoint: Optional PDS endpoint to use for all blobs. 483 483 + 484 484 + Returns: 485 485 + Configured BlobSource. 486 486 + 487 487 + Raises: 488 488 + ValueError: If refs is empty or format is invalid. 489 489 + """ 490 490 + if not refs: 491 491 + raise ValueError("refs cannot be empty") 492 492 + 493 493 + blob_refs: list[dict[str, str]] = [] 494 494 + 495 495 + for ref in refs: 496 496 + if "did" in ref and "cid" in ref: 497 497 + # Simple format 498 498 + blob_refs.append({"did": ref["did"], "cid": ref["cid"]}) 499 499 + elif "ref" in ref and "$link" in ref.get("ref", {}): 500 500 + # ATProto blob format - need DID from elsewhere 501 501 + raise ValueError( 502 502 + "ATProto blob format requires 'did' field. " 503 503 + "Use from_record_storage() for records with storage.blobs." 504 504 + ) 505 505 + else: 506 506 + raise ValueError(f"Invalid blob reference format: {ref}") 507 507 + 508 508 + return cls(blob_refs=blob_refs, pds_endpoint=pds_endpoint) 509 509 + 510 510 + 340 511 __all__ = [ 341 512 "URLSource", 342 513 "S3Source", 514 514 + "BlobSource", 343 515 ]

+25 -4

src/atdata/atmosphere/__init__.py

reviewed

··· 38 38 from .schema import SchemaPublisher, SchemaLoader 39 39 from .records import DatasetPublisher, DatasetLoader 40 40 from .lens import LensPublisher, LensLoader 41 41 + from .store import PDSBlobStore 41 42 from ._types import ( 42 43 AtUri, 43 44 SchemaRecord, ··· 102 103 Wraps SchemaPublisher/Loader and DatasetPublisher/Loader to provide 103 104 a unified interface compatible with LocalIndex. 104 105 106 106 + Optionally accepts a ``PDSBlobStore`` for writing dataset shards as 107 107 + ATProto blobs, enabling fully decentralized dataset storage. 108 108 + 105 109 Example: 106 110 :: 107 111 108 112 >>> client = AtmosphereClient() 109 113 >>> client.login("handle.bsky.social", "app-password") 110 114 >>> 115 115 + >>> # Without blob storage (external URLs only) 111 116 >>> index = AtmosphereIndex(client) 112 112 - >>> schema_ref = index.publish_schema(MySample, version="1.0.0") 117 117 + >>> 118 118 + >>> # With PDS blob storage 119 119 + >>> store = PDSBlobStore(client) 120 120 + >>> index = AtmosphereIndex(client, data_store=store) 113 121 >>> entry = index.insert_dataset(dataset, name="my-data") 114 122 """ 115 123 116 116 - def __init__(self, client: AtmosphereClient): 124 124 + def __init__( 125 125 + self, 126 126 + client: AtmosphereClient, 127 127 + *, 128 128 + data_store: Optional[PDSBlobStore] = None, 129 129 + ): 117 130 """Initialize the atmosphere index. 118 131 119 132 Args: 120 133 client: Authenticated AtmosphereClient instance. 134 134 + data_store: Optional PDSBlobStore for writing shards as blobs. 135 135 + If provided, insert_dataset will upload shards to PDS. 121 136 """ 122 137 self.client = client 123 138 self._schema_publisher = SchemaPublisher(client) 124 139 self._schema_loader = SchemaLoader(client) 125 140 self._dataset_publisher = DatasetPublisher(client) 126 141 self._dataset_loader = DatasetLoader(client) 127 127 - # AtmosphereIndex doesn't support data_store (uses PDS blobs) 128 128 - self.data_store = None 142 142 + self._data_store = data_store 143 143 + 144 144 + @property 145 145 + def data_store(self) -> Optional[PDSBlobStore]: 146 146 + """The PDS blob store for writing shards, or None if not configured.""" 147 147 + return self._data_store 129 148 130 149 # Dataset operations 131 150 ··· 291 310 __all__ = [ 292 311 # Client 293 312 "AtmosphereClient", 313 313 + # Storage 314 314 + "PDSBlobStore", 294 315 # Unified index (AbstractIndex protocol) 295 316 "AtmosphereIndex", 296 317 "AtmosphereIndexEntry",

+208

src/atdata/atmosphere/store.py

reviewed

··· 1 1 + """PDS blob storage for dataset shards. 2 2 + 3 3 + This module provides ``PDSBlobStore``, an implementation of the AbstractDataStore 4 4 + protocol that stores dataset shards as ATProto blobs in a Personal Data Server. 5 5 + 6 6 + This enables fully decentralized dataset storage where both metadata (records) 7 7 + and data (blobs) live on the AT Protocol network. 8 8 + 9 9 + Example: 10 10 + :: 11 11 + 12 12 + >>> from atdata.atmosphere import AtmosphereClient, PDSBlobStore 13 13 + >>> 14 14 + >>> client = AtmosphereClient() 15 15 + >>> client.login("handle.bsky.social", "app-password") 16 16 + >>> 17 17 + >>> store = PDSBlobStore(client) 18 18 + >>> urls = store.write_shards(dataset, prefix="mnist/v1") 19 19 + >>> print(urls) 20 20 + ['at://did:plc:.../blob/bafyrei...', ...] 21 21 + """ 22 22 + 23 23 + from __future__ import annotations 24 24 + 25 25 + import io 26 26 + import tempfile 27 27 + from dataclasses import dataclass 28 28 + from typing import TYPE_CHECKING, Any 29 29 + 30 30 + import webdataset as wds 31 31 + 32 32 + if TYPE_CHECKING: 33 33 + from ..dataset import Dataset 34 34 + from .client import AtmosphereClient 35 35 + 36 36 + 37 37 + @dataclass 38 38 + class PDSBlobStore: 39 39 + """PDS blob store implementing AbstractDataStore protocol. 40 40 + 41 41 + Stores dataset shards as ATProto blobs, enabling decentralized dataset 42 42 + storage on the AT Protocol network. 43 43 + 44 44 + Each shard is written to a temporary tar file, then uploaded as a blob 45 45 + to the user's PDS. The returned URLs are AT URIs that can be resolved 46 46 + to HTTP URLs for streaming. 47 47 + 48 48 + Attributes: 49 49 + client: Authenticated AtmosphereClient instance. 50 50 + 51 51 + Example: 52 52 + :: 53 53 + 54 54 + >>> store = PDSBlobStore(client) 55 55 + >>> urls = store.write_shards(dataset, prefix="training/v1") 56 56 + >>> # Returns AT URIs like: 57 57 + >>> # ['at://did:plc:abc/blob/bafyrei...', ...] 58 58 + """ 59 59 + 60 60 + client: "AtmosphereClient" 61 61 + 62 62 + def write_shards( 63 63 + self, 64 64 + ds: "Dataset", 65 65 + *, 66 66 + prefix: str, 67 67 + maxcount: int = 10000, 68 68 + maxsize: float = 3e9, 69 69 + **kwargs: Any, 70 70 + ) -> list[str]: 71 71 + """Write dataset shards as PDS blobs. 72 72 + 73 73 + Creates tar archives from the dataset and uploads each as a blob 74 74 + to the authenticated user's PDS. 75 75 + 76 76 + Args: 77 77 + ds: The Dataset to write. 78 78 + prefix: Logical path prefix for naming (used in shard names only). 79 79 + maxcount: Maximum samples per shard (default: 10000). 80 80 + maxsize: Maximum shard size in bytes (default: 3GB, PDS limit). 81 81 + **kwargs: Additional args passed to wds.ShardWriter. 82 82 + 83 83 + Returns: 84 84 + List of AT URIs for the written blobs, in format: 85 85 + ``at://{did}/blob/{cid}`` 86 86 + 87 87 + Raises: 88 88 + ValueError: If not authenticated. 89 89 + RuntimeError: If no shards were written. 90 90 + 91 91 + Note: 92 92 + PDS blobs have size limits (typically 50MB-5GB depending on PDS). 93 93 + Adjust maxcount/maxsize to stay within limits. 94 94 + """ 95 95 + if not self.client.did: 96 96 + raise ValueError("Client must be authenticated to upload blobs") 97 97 + 98 98 + did = self.client.did 99 99 + blob_urls: list[str] = [] 100 100 + 101 101 + # Write shards to temp files, upload each as blob 102 102 + with tempfile.TemporaryDirectory() as temp_dir: 103 103 + shard_pattern = f"{temp_dir}/shard-%06d.tar" 104 104 + written_files: list[str] = [] 105 105 + 106 106 + # Track written files via custom post callback 107 107 + def track_file(fname: str) -> None: 108 108 + written_files.append(fname) 109 109 + 110 110 + with wds.writer.ShardWriter( 111 111 + shard_pattern, 112 112 + maxcount=maxcount, 113 113 + maxsize=maxsize, 114 114 + post=track_file, 115 115 + **kwargs, 116 116 + ) as sink: 117 117 + for sample in ds.ordered(batch_size=None): 118 118 + sink.write(sample.as_wds) 119 119 + 120 120 + if not written_files: 121 121 + raise RuntimeError("No shards written") 122 122 + 123 123 + # Upload each shard as a blob 124 124 + for shard_path in written_files: 125 125 + with open(shard_path, "rb") as f: 126 126 + shard_data = f.read() 127 127 + 128 128 + blob_ref = self.client.upload_blob( 129 129 + shard_data, 130 130 + mime_type="application/x-tar", 131 131 + ) 132 132 + 133 133 + # Extract CID from blob reference 134 134 + cid = blob_ref["ref"]["$link"] 135 135 + at_uri = f"at://{did}/blob/{cid}" 136 136 + blob_urls.append(at_uri) 137 137 + 138 138 + return blob_urls 139 139 + 140 140 + def read_url(self, url: str) -> str: 141 141 + """Resolve an AT URI blob reference to an HTTP URL. 142 142 + 143 143 + Transforms ``at://did/blob/cid`` URIs to HTTP URLs that can be 144 144 + streamed by WebDataset. 145 145 + 146 146 + Args: 147 147 + url: AT URI in format ``at://{did}/blob/{cid}``. 148 148 + 149 149 + Returns: 150 150 + HTTP URL for fetching the blob via PDS API. 151 151 + 152 152 + Raises: 153 153 + ValueError: If URL format is invalid or PDS cannot be resolved. 154 154 + """ 155 155 + if not url.startswith("at://"): 156 156 + # Not an AT URI, return unchanged 157 157 + return url 158 158 + 159 159 + # Parse at://did/blob/cid 160 160 + parts = url[5:].split("/") # Remove 'at://' 161 161 + if len(parts) != 3 or parts[1] != "blob": 162 162 + raise ValueError(f"Invalid blob AT URI format: {url}") 163 163 + 164 164 + did, _, cid = parts 165 165 + return self.client.get_blob_url(did, cid) 166 166 + 167 167 + def supports_streaming(self) -> bool: 168 168 + """PDS blobs support streaming via HTTP. 169 169 + 170 170 + Returns: 171 171 + True. 172 172 + """ 173 173 + return True 174 174 + 175 175 + def create_source(self, urls: list[str]) -> "BlobSource": 176 176 + """Create a BlobSource for reading these AT URIs. 177 177 + 178 178 + This is a convenience method for creating a DataSource that can 179 179 + stream the blobs written by this store. 180 180 + 181 181 + Args: 182 182 + urls: List of AT URIs from write_shards(). 183 183 + 184 184 + Returns: 185 185 + BlobSource configured for the given URLs. 186 186 + 187 187 + Raises: 188 188 + ValueError: If URLs are not valid AT URIs. 189 189 + """ 190 190 + from .._sources import BlobSource 191 191 + 192 192 + blob_refs: list[dict[str, str]] = [] 193 193 + 194 194 + for url in urls: 195 195 + if not url.startswith("at://"): 196 196 + raise ValueError(f"Not an AT URI: {url}") 197 197 + 198 198 + parts = url[5:].split("/") 199 199 + if len(parts) != 3 or parts[1] != "blob": 200 200 + raise ValueError(f"Invalid blob AT URI: {url}") 201 201 + 202 202 + did, _, cid = parts 203 203 + blob_refs.append({"did": did, "cid": cid}) 204 204 + 205 205 + return BlobSource(blob_refs=blob_refs) 206 206 + 207 207 + 208 208 + __all__ = ["PDSBlobStore"]

+153

tests/test_integration_atmosphere.py

reviewed

··· 401 401 uri = AtUri.parse(f"at://did:plc:abc/{LEXICON_NAMESPACE}.sampleSchema/test") 402 402 403 403 assert uri.collection == f"{LEXICON_NAMESPACE}.sampleSchema" 404 404 + 405 405 + 406 406 + class TestPDSBlobStore: 407 407 + """Tests for PDSBlobStore blob storage.""" 408 408 + 409 409 + def test_create_with_client(self, authenticated_client): 410 410 + """PDSBlobStore can be created with authenticated client.""" 411 411 + from atdata.atmosphere import PDSBlobStore 412 412 + 413 413 + store = PDSBlobStore(client=authenticated_client) 414 414 + assert store.client is authenticated_client 415 415 + 416 416 + def test_supports_streaming(self, authenticated_client): 417 417 + """PDSBlobStore supports streaming.""" 418 418 + from atdata.atmosphere import PDSBlobStore 419 419 + 420 420 + store = PDSBlobStore(client=authenticated_client) 421 421 + assert store.supports_streaming() is True 422 422 + 423 423 + def test_write_shards_requires_auth(self, mock_atproto_client): 424 424 + """write_shards raises if client not authenticated.""" 425 425 + from atdata.atmosphere import PDSBlobStore 426 426 + 427 427 + # Create client without login 428 428 + client = AtmosphereClient(_client=mock_atproto_client) 429 429 + # Clear session 430 430 + client._session = None 431 431 + 432 432 + store = PDSBlobStore(client=client) 433 433 + 434 434 + # Create minimal mock dataset 435 435 + mock_ds = Mock() 436 436 + mock_ds.ordered.return_value = iter([]) 437 437 + 438 438 + with pytest.raises(ValueError, match="Not authenticated"): 439 439 + store.write_shards(mock_ds, prefix="test") 440 440 + 441 441 + def test_write_shards_uploads_blobs(self, authenticated_client, mock_atproto_client, tmp_path): 442 442 + """write_shards uploads each shard as a blob.""" 443 443 + from atdata.atmosphere import PDSBlobStore 444 444 + import webdataset as wds 445 445 + 446 446 + # Create a test dataset with samples 447 447 + tar_path = tmp_path / "test.tar" 448 448 + with wds.writer.TarWriter(str(tar_path)) as sink: 449 449 + sample = AtmoSample(name="test", value=42) 450 450 + sink.write(sample.as_wds) 451 451 + 452 452 + ds = atdata.Dataset[AtmoSample](str(tar_path)) 453 453 + 454 454 + # Mock upload_blob to return a blob reference 455 455 + authenticated_client.upload_blob = Mock(return_value={ 456 456 + "$type": "blob", 457 457 + "ref": {"$link": "bafyrei123abc"}, 458 458 + "mimeType": "application/x-tar", 459 459 + "size": 1024, 460 460 + }) 461 461 + 462 462 + store = PDSBlobStore(client=authenticated_client) 463 463 + urls = store.write_shards(ds, prefix="test/v1", maxcount=100) 464 464 + 465 465 + # Should have uploaded one shard 466 466 + assert len(urls) == 1 467 467 + assert urls[0] == "at://did:plc:integration123/blob/bafyrei123abc" 468 468 + 469 469 + # Verify upload_blob was called with tar data 470 470 + authenticated_client.upload_blob.assert_called_once() 471 471 + call_args = authenticated_client.upload_blob.call_args 472 472 + assert call_args.kwargs["mime_type"] == "application/x-tar" 473 473 + # First arg should be bytes (tar data) 474 474 + assert isinstance(call_args.args[0], bytes) 475 475 + 476 476 + def test_read_url_transforms_at_uri(self, authenticated_client, mock_atproto_client): 477 477 + """read_url transforms AT URIs to HTTP URLs.""" 478 478 + from atdata.atmosphere import PDSBlobStore 479 479 + 480 480 + authenticated_client.get_blob_url = Mock( 481 481 + return_value="https://pds.example.com/xrpc/com.atproto.sync.getBlob?did=did:plc:abc&cid=bafyrei123" 482 482 + ) 483 483 + 484 484 + store = PDSBlobStore(client=authenticated_client) 485 485 + url = store.read_url("at://did:plc:abc/blob/bafyrei123") 486 486 + 487 487 + assert "https://pds.example.com" in url 488 488 + assert "bafyrei123" in url 489 489 + authenticated_client.get_blob_url.assert_called_once_with("did:plc:abc", "bafyrei123") 490 490 + 491 491 + def test_read_url_passes_non_at_uri(self, authenticated_client): 492 492 + """read_url returns non-AT URIs unchanged.""" 493 493 + from atdata.atmosphere import PDSBlobStore 494 494 + 495 495 + store = PDSBlobStore(client=authenticated_client) 496 496 + url = store.read_url("https://example.com/data.tar") 497 497 + 498 498 + assert url == "https://example.com/data.tar" 499 499 + 500 500 + def test_read_url_invalid_format(self, authenticated_client): 501 501 + """read_url raises on invalid AT URI format.""" 502 502 + from atdata.atmosphere import PDSBlobStore 503 503 + 504 504 + store = PDSBlobStore(client=authenticated_client) 505 505 + 506 506 + with pytest.raises(ValueError, match="Invalid blob AT URI format"): 507 507 + store.read_url("at://did:plc:abc/invalid/format/extra") 508 508 + 509 509 + def test_create_source_returns_blob_source(self, authenticated_client): 510 510 + """create_source returns BlobSource for AT URIs.""" 511 511 + from atdata.atmosphere import PDSBlobStore 512 512 + from atdata._sources import BlobSource 513 513 + 514 514 + store = PDSBlobStore(client=authenticated_client) 515 515 + source = store.create_source([ 516 516 + "at://did:plc:abc/blob/bafyrei111", 517 517 + "at://did:plc:abc/blob/bafyrei222", 518 518 + ]) 519 519 + 520 520 + assert isinstance(source, BlobSource) 521 521 + assert len(source.blob_refs) == 2 522 522 + assert source.blob_refs[0]["did"] == "did:plc:abc" 523 523 + assert source.blob_refs[0]["cid"] == "bafyrei111" 524 524 + 525 525 + def test_create_source_invalid_url(self, authenticated_client): 526 526 + """create_source raises on non-AT URIs.""" 527 527 + from atdata.atmosphere import PDSBlobStore 528 528 + 529 529 + store = PDSBlobStore(client=authenticated_client) 530 530 + 531 531 + with pytest.raises(ValueError, match="Not an AT URI"): 532 532 + store.create_source(["https://example.com/data.tar"]) 533 533 + 534 534 + def test_atmosphere_index_with_data_store(self, authenticated_client): 535 535 + """AtmosphereIndex can be created with PDSBlobStore.""" 536 536 + from atdata.atmosphere import PDSBlobStore 537 537 + 538 538 + store = PDSBlobStore(client=authenticated_client) 539 539 + index = AtmosphereIndex(client=authenticated_client, data_store=store) 540 540 + 541 541 + assert index.data_store is store 542 542 + 543 543 + def test_atmosphere_index_data_store_property(self, authenticated_client): 544 544 + """AtmosphereIndex.data_store property returns the store.""" 545 545 + from atdata.atmosphere import PDSBlobStore 546 546 + 547 547 + store = PDSBlobStore(client=authenticated_client) 548 548 + index = AtmosphereIndex(client=authenticated_client, data_store=store) 549 549 + 550 550 + assert index.data_store is store 551 551 + 552 552 + def test_atmosphere_index_without_data_store(self, authenticated_client): 553 553 + """AtmosphereIndex without data_store has None.""" 554 554 + index = AtmosphereIndex(client=authenticated_client) 555 555 + 556 556 + assert index.data_store is None

+154 -1

tests/test_sources.py

reviewed

··· 9 9 import webdataset as wds 10 10 11 11 import atdata 12 12 - from atdata._sources import URLSource, S3Source 12 12 + from atdata._sources import URLSource, S3Source, BlobSource 13 13 from atdata._protocols import DataSource 14 14 15 15 ··· 287 287 288 288 assert client1 is client2 289 289 assert mock_boto.call_count == 1 290 290 + 291 291 + 292 292 + class TestBlobSource: 293 293 + """Tests for BlobSource (ATProto PDS blob storage).""" 294 294 + 295 295 + def test_conforms_to_protocol(self): 296 296 + """BlobSource should satisfy DataSource protocol.""" 297 297 + source = BlobSource(blob_refs=[{"did": "did:plc:abc", "cid": "bafyrei123"}]) 298 298 + assert isinstance(source, DataSource) 299 299 + 300 300 + def test_list_shards(self): 301 301 + """list_shards returns AT URIs.""" 302 302 + source = BlobSource(blob_refs=[ 303 303 + {"did": "did:plc:abc", "cid": "bafyrei111"}, 304 304 + {"did": "did:plc:abc", "cid": "bafyrei222"}, 305 305 + ]) 306 306 + assert source.list_shards() == [ 307 307 + "at://did:plc:abc/blob/bafyrei111", 308 308 + "at://did:plc:abc/blob/bafyrei222", 309 309 + ] 310 310 + 311 311 + def test_from_refs_simple_format(self): 312 312 + """from_refs accepts simple {did, cid} format.""" 313 313 + source = BlobSource.from_refs([ 314 314 + {"did": "did:plc:abc", "cid": "bafyrei123"}, 315 315 + ]) 316 316 + assert len(source.blob_refs) == 1 317 317 + assert source.blob_refs[0]["did"] == "did:plc:abc" 318 318 + assert source.blob_refs[0]["cid"] == "bafyrei123" 319 319 + 320 320 + def test_from_refs_with_endpoint(self): 321 321 + """from_refs accepts pds_endpoint parameter.""" 322 322 + source = BlobSource.from_refs( 323 323 + [{"did": "did:plc:abc", "cid": "bafyrei123"}], 324 324 + pds_endpoint="https://pds.example.com", 325 325 + ) 326 326 + assert source.pds_endpoint == "https://pds.example.com" 327 327 + 328 328 + def test_from_refs_empty(self): 329 329 + """from_refs raises on empty list.""" 330 330 + with pytest.raises(ValueError, match="cannot be empty"): 331 331 + BlobSource.from_refs([]) 332 332 + 333 333 + def test_from_refs_invalid_format(self): 334 334 + """from_refs raises on invalid blob reference format.""" 335 335 + with pytest.raises(ValueError, match="Invalid blob reference format"): 336 336 + BlobSource.from_refs([{"invalid": "data"}]) 337 337 + 338 338 + def test_from_refs_atproto_format_without_did(self): 339 339 + """from_refs raises helpful error for ATProto format without DID.""" 340 340 + with pytest.raises(ValueError, match="requires 'did' field"): 341 341 + BlobSource.from_refs([{"ref": {"$link": "bafyrei123"}}]) 342 342 + 343 343 + def test_resolve_pds_endpoint_uses_cache(self): 344 344 + """PDS endpoint resolution is cached.""" 345 345 + source = BlobSource(blob_refs=[{"did": "did:plc:abc", "cid": "cid"}]) 346 346 + 347 347 + # Pre-populate cache 348 348 + source._endpoint_cache["did:plc:abc"] = "https://cached.pds.com" 349 349 + 350 350 + endpoint = source._resolve_pds_endpoint("did:plc:abc") 351 351 + assert endpoint == "https://cached.pds.com" 352 352 + 353 353 + def test_resolve_pds_endpoint_uses_provided_endpoint(self): 354 354 + """Provided pds_endpoint is used instead of resolution.""" 355 355 + source = BlobSource( 356 356 + blob_refs=[{"did": "did:plc:abc", "cid": "cid"}], 357 357 + pds_endpoint="https://my.pds.com", 358 358 + ) 359 359 + 360 360 + endpoint = source._resolve_pds_endpoint("did:plc:abc") 361 361 + assert endpoint == "https://my.pds.com" 362 362 + 363 363 + def test_get_blob_url(self): 364 364 + """_get_blob_url constructs correct URL.""" 365 365 + source = BlobSource( 366 366 + blob_refs=[{"did": "did:plc:abc", "cid": "bafyrei123"}], 367 367 + pds_endpoint="https://pds.example.com", 368 368 + ) 369 369 + 370 370 + url = source._get_blob_url("did:plc:abc", "bafyrei123") 371 371 + assert url == "https://pds.example.com/xrpc/com.atproto.sync.getBlob?did=did:plc:abc&cid=bafyrei123" 372 372 + 373 373 + def test_shards_fetches_blobs(self): 374 374 + """shards property fetches blobs via HTTP.""" 375 375 + mock_response = Mock() 376 376 + mock_response.raw = Mock() 377 377 + mock_response.raise_for_status = Mock() 378 378 + 379 379 + with patch("requests.get", return_value=mock_response) as mock_get: 380 380 + source = BlobSource( 381 381 + blob_refs=[{"did": "did:plc:abc", "cid": "bafyrei123"}], 382 382 + pds_endpoint="https://pds.example.com", 383 383 + ) 384 384 + 385 385 + shards = list(source.shards) 386 386 + 387 387 + assert len(shards) == 1 388 388 + shard_id, stream = shards[0] 389 389 + assert shard_id == "at://did:plc:abc/blob/bafyrei123" 390 390 + assert stream is mock_response.raw 391 391 + 392 392 + mock_get.assert_called_once_with( 393 393 + "https://pds.example.com/xrpc/com.atproto.sync.getBlob?did=did:plc:abc&cid=bafyrei123", 394 394 + stream=True, 395 395 + timeout=60, 396 396 + ) 397 397 + 398 398 + def test_open_shard_fetches_single_blob(self): 399 399 + """open_shard fetches a specific blob.""" 400 400 + mock_response = Mock() 401 401 + mock_response.raw = Mock() 402 402 + mock_response.raise_for_status = Mock() 403 403 + 404 404 + with patch("requests.get", return_value=mock_response) as mock_get: 405 405 + source = BlobSource( 406 406 + blob_refs=[ 407 407 + {"did": "did:plc:abc", "cid": "bafyrei111"}, 408 408 + {"did": "did:plc:abc", "cid": "bafyrei222"}, 409 409 + ], 410 410 + pds_endpoint="https://pds.example.com", 411 411 + ) 412 412 + 413 413 + stream = source.open_shard("at://did:plc:abc/blob/bafyrei222") 414 414 + 415 415 + assert stream is mock_response.raw 416 416 + mock_get.assert_called_once() 417 417 + call_args = mock_get.call_args 418 418 + assert "bafyrei222" in call_args[0][0] 419 419 + 420 420 + def test_open_shard_not_found(self): 421 421 + """open_shard raises KeyError for unknown shard.""" 422 422 + source = BlobSource(blob_refs=[{"did": "did:plc:abc", "cid": "bafyrei123"}]) 423 423 + 424 424 + with pytest.raises(KeyError, match="Shard not found"): 425 425 + source.open_shard("at://did:plc:abc/blob/unknown") 426 426 + 427 427 + def test_open_shard_invalid_format(self): 428 428 + """open_shard raises ValueError for invalid shard ID format.""" 429 429 + # Test that we properly validate the AT URI format 430 430 + # by checking the error message when we pass an invalid format 431 431 + # that isn't in the list but would fail format check 432 432 + source = BlobSource( 433 433 + blob_refs=[{"did": "did:plc:abc", "cid": "bafyrei123"}], 434 434 + ) 435 435 + 436 436 + # A non-AT URI should raise KeyError (not in list) 437 437 + with pytest.raises(KeyError, match="Shard not found"): 438 438 + source.open_shard("not-an-at-uri") 439 439 + 440 440 + # An AT URI with wrong format should also raise KeyError (not in list) 441 441 + with pytest.raises(KeyError, match="Shard not found"): 442 442 + source.open_shard("at://did:plc:abc/wrong/format") 290 443 291 444 292 445 class TestDatasetWithDataSource: