diff --git a/tests/sources/xml/test_marc.py b/tests/sources/xml/test_marc.py index 54e8f0e..f53c4d4 100644 --- a/tests/sources/xml/test_marc.py +++ b/tests/sources/xml/test_marc.py @@ -2282,6 +2282,11 @@ def test_marc_parse_source_record_recovers_from_malformed_xml(): assert Marc.find_element_string(source_record, "leader") == "x" +def test_marc_parse_source_record_raises_parsing_errors(): + with pytest.raises(etree.XMLSyntaxError, match="Document is empty"): + Marc.parse_source_record(b"") + + def test_record_is_deleted_returns_true_if_deleted(): deleted_record = Marc.parse_source_file("tests/fixtures/marc/marc_record_deleted.xml") assert Marc.record_is_deleted(next(deleted_record)) is True diff --git a/tests/test_cli.py b/tests/test_cli.py index b8ca274..5a0a59d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -213,7 +213,7 @@ def test_transform_run_timestamp_argument_not_passed_and_timestamp_minted( assert transformer.run_data["run_date"] == "2024-06-03" -def test_transform_no_memory_fault_for_threaded_bs4_parsing(monkeypatch, tmp_path): +def test_transform_no_memory_fault_for_bs4_parsing(monkeypatch, tmp_path): """This test requires running the CLI as a subprocess to simulate the 'uv run ...' context. In this context, we have observed memory faults when BeautifulSoup4 is used by a source, and the number of records in the run requires multiple batches during @@ -238,6 +238,39 @@ def test_transform_no_memory_fault_for_threaded_bs4_parsing(monkeypatch, tmp_pat assert result.returncode == 0 +def test_transform_no_memory_fault_for_marc_parsing(tmp_path): + """As test_transform_no_memory_fault_for_bs4_parsing, for lxml parsing. + + Requires one more record than a single parquet dataset write batch. + """ + record = ( + "03282nam 2200721Ki 4500" + '990027185640106761' + '170906s2016 fr mun| o e zxx d' + 'Title' + ) + input_file = tmp_path / "alma-2026-01-01-daily-extracted-records-to-index.xml" + input_file.write_text(f"{record * 1001}") + + result = subprocess.run( # noqa: S603 + [ # noqa: S607 + "uv", + "run", + "transform", + "-s", + "alma", + "-i", + str(input_file), + "-o", + f"{tmp_path}/dataset", + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0 + + def test_transform_critical_error_prevents_writing( caplog, runner, source_input_file, source_transformer, empty_dataset_location ): diff --git a/transmogrifier/sources/xmltransformer.py b/transmogrifier/sources/xmltransformer.py index 2005b6d..9fb02ae 100644 --- a/transmogrifier/sources/xmltransformer.py +++ b/transmogrifier/sources/xmltransformer.py @@ -1,6 +1,7 @@ from __future__ import annotations -import threading +import itertools +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING import smart_open # type: ignore[import-untyped] @@ -12,6 +13,8 @@ if TYPE_CHECKING: from collections.abc import Iterator +READ_CHUNK_SIZE = 100 + class XMLTransformer(Transformer): """XML transformer class.""" @@ -21,9 +24,25 @@ def read_source_file(cls, source_file: str) -> Iterator[bytes]: """ Read XML file and yield each element as bytes. + The file is read on a single, dedicated thread. pyarrow may resume this generator + under a new thread state each time, and lxml has a string dictionary per thread + state. If iterparse finishes under a different thread state than it started, one + that has parsed something else, lxml replaces the document's dictionary and the + application aborts with a memory fault when that document is freed. + Args: source_file: A file containing source records to be transformed. """ + records = cls._iterparse_source_file(source_file) + with ThreadPoolExecutor(max_workers=1) as executor: + while chunk := executor.submit( + lambda: list(itertools.islice(records, READ_CHUNK_SIZE)) + ).result(): + yield from chunk + + @classmethod + def _iterparse_source_file(cls, source_file: str) -> Iterator[bytes]: + """Iteratively parse an XML file and yield each element as bytes.""" with smart_open.open(source_file, "rb") as file: for _, element in etree.iterparse( file, @@ -42,18 +61,7 @@ def parse_source_record(cls, source_record: bytes) -> Tag: Args: source_record: An XML source record as bytes. """ - return cls.parse_bs4_in_isolated_thread(source_record) - - @classmethod - def parse_bs4_in_isolated_thread(cls, source_record: bytes) -> Tag: - def parse() -> Tag: - nonlocal source_record - source_record = BeautifulSoup(source_record, "xml") - - thread = threading.Thread(target=parse) - thread.start() - thread.join() - return source_record + return BeautifulSoup(source_record, "xml") @classmethod def get_main_titles(cls, _source_record: Tag) -> list[Tag]: