diff --git a/src/dynavec/ingest.py b/src/dynavec/ingest.py index e08ef85..c4d1277 100644 --- a/src/dynavec/ingest.py +++ b/src/dynavec/ingest.py @@ -18,6 +18,7 @@ import datetime import hashlib +import logging from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field from pathlib import Path @@ -30,6 +31,7 @@ from .utils import chunked Metadata = dict[str, Any] +logger = logging.getLogger(__name__) def _normalize_front_matter(value: Any) -> Any: @@ -496,12 +498,16 @@ def __init__( suffix: str | tuple[str, ...] | None = None, boto_session: Any = None, s3_client: Any = None, + on_error: str = "raise", ) -> None: + if on_error not in ("raise", "skip"): + raise ValueError("on_error must be 'raise' or 'skip'") self.bucket = bucket self.prefix = prefix self.suffix = suffix self._session = boto_session self._client = s3_client + self.on_error = on_error def _get_client(self) -> Any: if self._client is not None: @@ -529,10 +535,7 @@ def _parse_object(self, key: str, data: bytes, content_type: str) -> Iterator[Re # 1. Markdown if ext == ".md" or "text/markdown" in content_type: - try: - text = data.decode("utf-8-sig") - except UnicodeDecodeError: - return + text = data.decode("utf-8-sig") clean_text, front_meta = MarkdownSource._front_matter(text, Path(key)) if clean_text.strip(): yield Record( @@ -544,10 +547,7 @@ def _parse_object(self, key: str, data: bytes, content_type: str) -> Iterator[Re # 2. CSV if ext == ".csv" or "text/csv" in content_type: - try: - text = data.decode("utf-8-sig") - except UnicodeDecodeError: - return + text = data.decode("utf-8-sig") lines = text.splitlines() csv_reader = csv.reader(lines) try: @@ -699,7 +699,17 @@ def __iter__(self) -> Iterator[Record]: body = obj_resp["Body"] data = body.read() if hasattr(body, "read") else bytes(body) - yield from self._parse_object(key, data, content_type) + try: + yield from self._parse_object(key, data, content_type) + except Exception as exc: + if self.on_error == "skip": + logger.warning( + "Failed to parse S3 object %s: %s", + key, + exc, + ) + continue + raise if resp.get("IsTruncated"): continuation_token = resp.get("NextContinuationToken") diff --git a/tests/test_s3_source.py b/tests/test_s3_source.py index 57684e5..ad74207 100644 --- a/tests/test_s3_source.py +++ b/tests/test_s3_source.py @@ -1,6 +1,8 @@ import io from unittest.mock import MagicMock +import pytest + from dynavec.ingest import S3Source, ingest @@ -168,3 +170,154 @@ def test_s3_source_ingest_integration(): assert doc.id == "s3://test-bucket/article.txt#chunk0" assert doc.metadata["source_id"] == "s3://test-bucket/article.txt" assert doc.metadata["source"] == "s3" + + +def test_s3_source_raises_on_invalid_utf8_by_default(): + client = MagicMock() + + client.list_objects_v2.return_value = { + "Contents": [{"Key": "documents/broken.md"}], + "IsTruncated": False, + } + + client.get_object.return_value = { + "ContentType": "text/markdown", + "Body": io.BytesIO(b"\xff\xfe\xfa"), + } + + source = S3Source( + "test-bucket", + s3_client=client, + ) + + with pytest.raises(UnicodeDecodeError): + list(source) + + +def test_s3_source_skips_invalid_object_when_on_error_skip(): + client = MagicMock() + + client.list_objects_v2.return_value = { + "Contents": [ + {"Key": "documents/broken.md"}, + {"Key": "documents/good.md"}, + ], + "IsTruncated": False, + } + + def mock_get_object(Bucket, Key): + if Key == "documents/broken.md": + return { + "ContentType": "text/markdown", + "Body": io.BytesIO(b"\xff\xfe\xfa"), + } + + return { + "ContentType": "text/markdown", + "Body": io.BytesIO(b"# Valid document"), + } + + client.get_object.side_effect = mock_get_object + + source = S3Source( + "test-bucket", + s3_client=client, + on_error="skip", + ) + + records = list(source) + + assert len(records) == 1 + assert records[0].id == "s3://test-bucket/documents/good.md" + assert records[0].text == "# Valid document" + + +def test_s3_source_rejects_invalid_on_error(): + with pytest.raises(ValueError, match="on_error"): + S3Source( + "test-bucket", + on_error="invalid", + ) + + +def test_s3_source_logs_warning_when_skipping(caplog): + client = MagicMock() + + client.list_objects_v2.return_value = { + "Contents": [{"Key": "documents/broken.md"}], + "IsTruncated": False, + } + + client.get_object.return_value = { + "ContentType": "text/markdown", + "Body": io.BytesIO(b"\xff\xfe\xfa"), + } + + source = S3Source( + "test-bucket", + s3_client=client, + on_error="skip", + ) + + records = list(source) + + assert records == [] + assert "documents/broken.md" in caplog.text + + +def test_s3_source_does_not_skip_s3_get_object_errors(): + client = MagicMock() + + client.list_objects_v2.return_value = { + "Contents": [{"Key": "documents/guide.md"}], + "IsTruncated": False, + } + + client.get_object.side_effect = RuntimeError("S3 access denied") + + source = S3Source( + "test-bucket", + s3_client=client, + on_error="skip", + ) + + with pytest.raises(RuntimeError, match="S3 access denied"): + list(source) + + +def test_s3_source_skips_invalid_markdown_front_matter(caplog): + client = MagicMock() + + client.list_objects_v2.return_value = { + "Contents": [ + {"Key": "documents/broken.md"}, + {"Key": "documents/good.md"}, + ], + "IsTruncated": False, + } + + def mock_get_object(Bucket, Key): + if Key == "documents/broken.md": + return { + "ContentType": "text/markdown", + "Body": io.BytesIO(b"---\ntitle: [invalid\n---\nBroken document"), + } + + return { + "ContentType": "text/markdown", + "Body": io.BytesIO(b"# Good document"), + } + + client.get_object.side_effect = mock_get_object + + source = S3Source( + "test-bucket", + s3_client=client, + on_error="skip", + ) + + records = list(source) + + assert len(records) == 1 + assert records[0].id == "s3://test-bucket/documents/good.md" + assert "documents/broken.md" in caplog.text