From 85a435b10409491e925c651774ce9981f8baaadd Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 19 May 2026 10:42:26 -0400 Subject: [PATCH] Write fulltext records to TIMDEX dataset Why these changes are being introduced: The final leg of the journey of harvesting fulltext for this app is writing fulltexts back to the TIMDEX dataset to live alongside the records they are associated with. We should still support writing locally to JSONLines for debugging. How this addresses that need: In direct support of this goal, - functions from `harvest.py` now yield `DatasetFulltext` instances which are expected by timdex-dataset-api (TDA) library for writing to the "fulltexts" data type - refactor `TIMDEXThesesRecords` to accept an instantiated `TIMDEXDataset` instance such that it can be created once, then reused for reading + writing - update `harvest` CLI command to default to writing back to the TIMDEX dataset Related to this work was some refactoring, - updated `record_and_fulltext_iter()` to use the `joblib` library to offload the complexities of parallel work to a well known library Side effects of this change: * By default, this CLI application will attempt to write the fulltext records back to the TIMDEX dataset Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-558 --- dfh/cli.py | 39 +++++++++--- dfh/config.py | 6 +- dfh/harvest.py | 89 ++++++++++++-------------- dfh/timdex_dataset.py | 35 ++++++----- pyproject.toml | 4 +- tests/conftest.py | 10 ++- tests/test_cli.py | 37 +++++++++++ tests/test_dataset.py | 6 +- tests/test_harvest.py | 141 ++++++++++++++++++++++++++++++++++++++++++ uv.lock | 31 +++++++++- 10 files changed, 313 insertions(+), 85 deletions(-) create mode 100644 tests/test_harvest.py diff --git a/dfh/cli.py b/dfh/cli.py index 009362a..b267f80 100644 --- a/dfh/cli.py +++ b/dfh/cli.py @@ -1,14 +1,16 @@ import logging import time -from datetime import timedelta +from collections.abc import Iterator +from datetime import datetime, timedelta import click import jsonlines +from timdex_dataset_api.data_types import DatasetFulltext from dfh.config import configure_logger, configure_sentry from dfh.dspace import warm_dspace_auth from dfh.harvest import record_and_fulltext_iter -from dfh.timdex_dataset import TIMDEXThesesRecords +from dfh.timdex_dataset import TIMDEXThesesRecords, get_timdex_dataset logger = logging.getLogger(__name__) @@ -115,26 +117,43 @@ def harvest( # warm DSpace API authentication warm_dspace_auth() + # init TIMDEXDataset instance, used for reading and writing + timdex_dataset = get_timdex_dataset(dataset_location) + # get iterator of target TIMDEX records + bitstream information ttr = TIMDEXThesesRecords( - dataset_location=dataset_location, + timdex_dataset=timdex_dataset, run_id=run_id, limit=record_limit, ) records_and_bitstreams = ttr.record_and_bitstream_metadata_iter() - # get iterator of records + fulltext, ready for writing + # get iterator of DatasetFulltext instances, ready for writing records_and_fulltexts = record_and_fulltext_iter( records_and_bitstreams, max_workers=workers, ) - if not output_jsonl: - raise ValueError( - "WIP: Until TIMDEX dataset has new 'fulltexts' " - "data type, only JSONL output is supported." - ) + # write to local JSONLines file for debugging + if output_jsonl: + write_jsonlines_output(output_jsonl, records_and_fulltexts) + # default: write to TIMDEX dataset + else: + timdex_dataset.fulltexts.write(records_and_fulltexts) + +def write_jsonlines_output( + output_jsonl: str, + records_and_fulltexts: Iterator[DatasetFulltext], +) -> None: + """Utility function to write to JSONLines for local debugging.""" with jsonlines.open(output_jsonl, "w") as writer: for record in records_and_fulltexts: - writer.write(record) + output_record = record.to_dict() + fulltext = output_record["fulltext"] + if isinstance(fulltext, bytes): + output_record["fulltext"] = fulltext.decode() + timestamp = output_record["fulltext_timestamp"] + if isinstance(timestamp, datetime): + output_record["fulltext_timestamp"] = timestamp.isoformat() + writer.write(output_record) diff --git a/dfh/config.py b/dfh/config.py index 5a36699..e98e553 100644 --- a/dfh/config.py +++ b/dfh/config.py @@ -8,14 +8,16 @@ def configure_logger(logger: logging.Logger, *, verbose: bool) -> str: if verbose: logging.basicConfig( format="%(asctime)s %(levelname)s %(name)s.%(funcName)s() line %(lineno)d: " - "%(message)s" + "%(message)s", + force=True, ) logger.setLevel(logging.DEBUG) for handler in logging.root.handlers: handler.addFilter(logging.Filter("dfh")) else: logging.basicConfig( - format="%(asctime)s %(levelname)s %(name)s.%(funcName)s(): %(message)s" + format="%(asctime)s %(levelname)s %(name)s.%(funcName)s(): %(message)s", + force=True, ) logger.setLevel(logging.INFO) return ( diff --git a/dfh/harvest.py b/dfh/harvest.py index fff3093..4a111dc 100644 --- a/dfh/harvest.py +++ b/dfh/harvest.py @@ -2,10 +2,11 @@ import threading import time from collections.abc import Iterator -from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait import requests from dspace_rest_client.client import DSpaceClient +from joblib import Parallel, delayed +from timdex_dataset_api.data_types import DatasetFulltext from dfh.dspace import get_dspace_client, get_presigned_url_for_bitstream @@ -18,52 +19,42 @@ def record_and_fulltext_iter( records: Iterator[dict], - *, max_workers: int = 10, - log_progress_interval: int = 10, -) -> Iterator[dict]: + log_progress_interval: int = 1000, +) -> Iterator[DatasetFulltext]: """Yield records with fulltext fetched in parallel. - Uses ThreadPoolExecutor (main IO bottleneck is network) to generate pre-signed URLs - and download bitstream content in parallel. - - The worker function _record_with_fulltext() has built-in retries. This orchestration - function with parallelizes the work is not aware of retries. + Uses a threaded worker to generate pre-signed URLs and download bitstream content in + parallel. The worker function _get_record_with_fulltext() has built-in retries. As + such, this orchestration function is not concerned with retries. """ - pending: set[Future[dict]] = set() - completed_count = 0 - - with ThreadPoolExecutor(max_workers=max_workers) as executor: - for record in records: - pending.add(executor.submit(_record_with_fulltext, record)) - - if len(pending) < max_workers: - continue - - done, pending = wait(pending, return_when=FIRST_COMPLETED) - for future in done: - completed_count += 1 - yield future.result() - if completed_count % log_progress_interval == 0: - logger.debug(f"Extracted fulltext for {completed_count} records.") - - while pending: - done, pending = wait(pending, return_when=FIRST_COMPLETED) - for future in done: - completed_count += 1 - yield future.result() - if completed_count % log_progress_interval == 0: - logger.debug(f"Extracted fulltext for {completed_count} records.") - - -def _record_with_fulltext( + parallel_client = Parallel( + n_jobs=max_workers, + prefer="threads", + return_as="generator_unordered", + ) + worker_func = delayed(_get_record_with_fulltext) + + results = parallel_client(worker_func(record) for record in records) + + # log results + count = 0 + for result in results: + count += 1 + if count % log_progress_interval == 0: + logger.info(f"Extracted fulltext for {count} records.") + yield result + logger.info(f"Extraction complete for {count} records.") + + +def _get_record_with_fulltext( record: dict, *, retry_attempts: int = 4, initial_backoff_seconds: float = 1.0, backoff_factor: float = 2.0, timeout_seconds: int = 60, -) -> dict: +) -> DatasetFulltext: """Return a TIMDEX fulltext record for one source record. The primary work here is two network requests: @@ -95,7 +86,7 @@ def _record_with_fulltext( # download bitstream content from S3 response = requests.get(pre_signed_url, timeout=timeout_seconds) response.raise_for_status() - fulltext = response.text + fulltext = response.content # break out of retries loop if successful break @@ -103,26 +94,28 @@ def _record_with_fulltext( except Exception as exc: if attempt == retry_attempts: logger.exception( - f"Max retries of {retry_attempts} encountered, failed to download " + f"Max retries of {retry_attempts} encountered, failed to download. " + f"""timdex_record_id '{record["timdex_record_id"]}', """ f"bitstream '{bitstream_uuid}'" ) break sleep_seconds = initial_backoff_seconds * (backoff_factor ** (attempt - 1)) logger.warning( - f"Retrying download for bitstream " - f"'{bitstream_uuid}' after attempt {attempt}/{retry_attempts} " + f"Retrying download for " + f"""timdex_record_id '{record["timdex_record_id"]}', """ + f"bitstream '{bitstream_uuid}'" + f"after attempt {attempt}/{retry_attempts} " f"failed; sleeping {sleep_seconds:.1f} seconds. Cause: {exc}" ) time.sleep(sleep_seconds) - return { - "timdex_record_id": record["timdex_record_id"], - "run_id": record["run_id"], - "run_record_offset": record["run_record_offset"], - "fulltext_bitstream_uuid": bitstream_uuid, - "fulltext_bitstream_content": fulltext, - } + return DatasetFulltext( + timdex_record_id=record["timdex_record_id"], + run_id=record["run_id"], + run_record_offset=record["run_record_offset"], + fulltext=fulltext, + ) def _get_dspace_client_for_thread() -> DSpaceClient: diff --git a/dfh/timdex_dataset.py b/dfh/timdex_dataset.py index 9fa61a2..c60c6ce 100644 --- a/dfh/timdex_dataset.py +++ b/dfh/timdex_dataset.py @@ -24,6 +24,20 @@ class TextBitstreamInfo(TypedDict): mimetype: str | None +def get_timdex_dataset(dataset_location: str | None) -> TIMDEXDataset: + """Get an initialized TIMDEXDataset instance. + + If dataset_location is not passed, look to env var TIMDEX_DATASET_LOCATION. + """ + dataset_location = dataset_location or os.getenv("TIMDEX_DATASET_LOCATION") + if not dataset_location: + raise ValueError( + "'dataset_location' must be passed on init " + "or env var 'TIMDEX_DATASET_LOCATION' set" + ) + return TIMDEXDataset(dataset_location) + + class TIMDEXThesesRecords: """Class to retrieve TIMDEX dataset records of DSpace Theses. @@ -42,32 +56,21 @@ class TIMDEXThesesRecords: def __init__( self, - dataset_location: str | None = None, + timdex_dataset: TIMDEXDataset, run_id: str | None = None, limit: int | None = None, ) -> None: + self.timdex_dataset = timdex_dataset self.run_id = run_id self.limit = limit - self.timdex_dataset = self._init_timdex_dataset(dataset_location) - - def _init_timdex_dataset(self, dataset_location: str | None) -> TIMDEXDataset: - """Init TIMDEXDataset and compose to self. - - If dataset_location is not passed, look to env var TIMDEX_DATASET_LOCATION. - """ - dataset_location = dataset_location or os.getenv("TIMDEX_DATASET_LOCATION") - if not dataset_location: - raise ValueError( - "'dataset_location' must be passed on init " - "or env var 'TIMDEX_DATASET_LOCATION' set" - ) - return TIMDEXDataset(dataset_location) - def record_and_bitstream_metadata_iter( self, ) -> Iterator[dict]: """Yield TIMDEX DSpace thesis records with their fulltext bitstream info.""" + logger.debug( + "Preparing to yield record + bitstream metadata from TIMDEX dataset." + ) record_count = 0 error_count = 0 for record in self.theses_records(): diff --git a/pyproject.toml b/pyproject.toml index bdef0f2..ab756ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ requires-python = ">=3.13" dependencies = [ "click>=8.2.1", "dspace-rest-client>=0.1.17", + "joblib>=1.5.3", "jsonlines>=4.0.0", "requests>=2.33.1", "sentry-sdk>=2.34.1", @@ -20,6 +21,7 @@ dependencies = [ dev = [ "coveralls>=4.0.1", "ipython>=9.13.0", + "joblib-stubs>=1.5.3.1.20260117", "mypy>=1.17.1", "pip-audit>=2.9.0", "pre-commit>=4.3.0", @@ -107,7 +109,7 @@ dfh = "dfh.cli:main" include = ["dfh", "dfh.*"] [tool.uv.sources] -timdex-dataset-api = { git = "https://github.com/MITLibraries/timdex-dataset-api" } +timdex-dataset-api = { git = "https://github.com/MITLibraries/timdex-dataset-api", rev = "USE-531-fulltext-data-type-add" } [build-system] requires = ["setuptools>=61"] diff --git a/tests/conftest.py b/tests/conftest.py index 4ae83c0..020e97d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,7 @@ import pytest from click.testing import CliRunner +from timdex_dataset_api import TIMDEXDataset from dfh.timdex_dataset import TIMDEXThesesRecords @@ -90,5 +91,10 @@ def dspace_mets_text_bitstream_info(): @pytest.fixture -def timdex_theses_records(tmp_path): - return TIMDEXThesesRecords(dataset_location=str(tmp_path)) +def timdex_dataset(tmp_path): + return TIMDEXDataset(str(tmp_path)) + + +@pytest.fixture +def timdex_theses_records(timdex_dataset): + return TIMDEXThesesRecords(timdex_dataset=timdex_dataset) diff --git a/tests/test_cli.py b/tests/test_cli.py index e69de29..4867f27 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -0,0 +1,37 @@ +import json + +from timdex_dataset_api.data_types import DatasetFulltext + +from dfh.cli import write_jsonlines_output + +RUN_RECORD_OFFSET = 12 + + +def test_write_jsonlines_output_serializes_dataset_fulltext(tmp_path): + output_jsonl = tmp_path / "fulltexts.jsonl" + dataset_fulltexts = iter( + [ + DatasetFulltext( + timdex_record_id="dspace:123", + run_id="run-1", + run_record_offset=RUN_RECORD_OFFSET, + fulltext_timestamp="2026-05-18T12:34:56+00:00", + fulltext=b"fulltext content", + ), + ] + ) + + write_jsonlines_output(str(output_jsonl), dataset_fulltexts) + + output_record = json.loads(output_jsonl.read_text().strip()) + assert output_record == { + "timdex_record_id": "dspace:123", + "run_id": "run-1", + "run_record_offset": RUN_RECORD_OFFSET, + "fulltext_timestamp": "2026-05-18T12:34:56+00:00", + "fulltext_md5": "1fd0a1296d0665461f688201f63a37f1", + "fulltext": "fulltext content", + "year": "2026", + "month": "05", + "day": "18", + } diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 65f327b..fcafd96 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -8,15 +8,15 @@ from dfh.timdex_dataset import ( MissingTextBitstreamError, SourceRecordParseError, - TIMDEXThesesRecords, + get_timdex_dataset, ) -def test_init_requires_dataset_location(monkeypatch): +def test_get_timdex_dataset_requires_dataset_location(monkeypatch): monkeypatch.delenv("TIMDEX_DATASET_LOCATION", raising=False) with pytest.raises(ValueError, match="dataset_location"): - TIMDEXThesesRecords() + get_timdex_dataset(None) def test_theses_records_returns_only_records_with_thesis_content_type( diff --git a/tests/test_harvest.py b/tests/test_harvest.py new file mode 100644 index 0000000..09648e6 --- /dev/null +++ b/tests/test_harvest.py @@ -0,0 +1,141 @@ +from unittest import mock + +import pytest +import requests +from timdex_dataset_api.data_types import DatasetFulltext + +from dfh.harvest import _get_record_with_fulltext, record_and_fulltext_iter + +RECORD_COUNT = 3 +RETRY_ATTEMPTS = 2 + + +@pytest.fixture +def record(): + return { + "timdex_record_id": "dspace:123", + "run_id": "run-1", + "run_record_offset": 42, + "fulltext_bitstream": {"uuid": "abc123"}, + } + + +@pytest.fixture +def presigned_url(): + return "https://s3.aws/mit-dspace/abc123" + + +@pytest.fixture +def successful_fulltext_response(): + response = mock.Mock(content=b"fulltext content") + response.raise_for_status.return_value = None + return response + + +def test_record_and_fulltext_iter_yields_fulltext_dataset_row( + record, + successful_fulltext_response, +): + records = iter([record] * RECORD_COUNT) + + with ( + mock.patch( + "dfh.harvest.get_dspace_client", + return_value="client", + ), + mock.patch( + "dfh.harvest.get_presigned_url_for_bitstream", + side_effect=lambda _client, uuid: f"https://s3.aws/mit-dspace/{uuid}", + ), + mock.patch( + "dfh.harvest.requests.get", + return_value=successful_fulltext_response, + ), + ): + results = list(record_and_fulltext_iter(records, max_workers=1)) + + assert len(results) == RECORD_COUNT + + result = results[0] + assert isinstance(result, DatasetFulltext) + assert result.timdex_record_id == record["timdex_record_id"] + assert result.run_id == record["run_id"] + assert result.run_record_offset == record["run_record_offset"] + assert result.fulltext == b"fulltext content" + + +def test_get_record_with_fulltext_retries_presigned_url_429( + record, + presigned_url, + successful_fulltext_response, +): + rate_limited_response = mock.Mock(status_code=429, text="Too Many Requests") + successful_presign_response = mock.Mock(status_code=200) + successful_presign_response.json.return_value = {"presignedUrl": presigned_url} + dspace_client = mock.Mock(API_ENDPOINT="https://dspace.mit.edu/server/api") + dspace_client.api_get.side_effect = [ + rate_limited_response, # first request fails + successful_presign_response, # second is success + ] + + with ( + mock.patch( + "dfh.harvest._get_dspace_client_for_thread", + return_value=dspace_client, + ), + mock.patch( + "dfh.harvest.requests.get", + return_value=successful_fulltext_response, + ) as s3_get, + mock.patch("dfh.harvest.time.sleep") as sleep, + ): + result = _get_record_with_fulltext( + record, + retry_attempts=RETRY_ATTEMPTS, + initial_backoff_seconds=0, + ) + + assert result.fulltext == b"fulltext content" + assert dspace_client.api_get.call_count == RETRY_ATTEMPTS + s3_get.assert_called_once() + sleep.assert_called_once_with(0) + + +def test_get_record_with_fulltext_retries_s3_500( + record, + presigned_url, + successful_fulltext_response, +): + failed_response = mock.Mock() + failed_response.raise_for_status.side_effect = requests.exceptions.HTTPError( + "500 Server Error" + ) + + with ( + mock.patch( + "dfh.harvest._get_dspace_client_for_thread", + return_value="client", + ), + mock.patch( + "dfh.harvest.get_presigned_url_for_bitstream", + return_value=presigned_url, + ) as get_presigned_url, + mock.patch( + "dfh.harvest.requests.get", + side_effect=[ + failed_response, # first request fails + successful_fulltext_response, # second is success + ], + ) as s3_get, + mock.patch("dfh.harvest.time.sleep") as sleep, + ): + result = _get_record_with_fulltext( + record, + retry_attempts=RETRY_ATTEMPTS, + initial_backoff_seconds=0, + ) + + assert result.fulltext == b"fulltext content" + assert get_presigned_url.call_count == RETRY_ATTEMPTS + assert s3_get.call_count == RETRY_ATTEMPTS + sleep.assert_called_once_with(0) diff --git a/uv.lock b/uv.lock index affc26f..f81134e 100644 --- a/uv.lock +++ b/uv.lock @@ -555,6 +555,7 @@ source = { editable = "." } dependencies = [ { name = "click" }, { name = "dspace-rest-client" }, + { name = "joblib" }, { name = "jsonlines" }, { name = "requests" }, { name = "sentry-sdk" }, @@ -565,6 +566,7 @@ dependencies = [ dev = [ { name = "coveralls" }, { name = "ipython" }, + { name = "joblib-stubs" }, { name = "mypy" }, { name = "pip-audit" }, { name = "pre-commit" }, @@ -576,16 +578,18 @@ dev = [ requires-dist = [ { name = "click", specifier = ">=8.2.1" }, { name = "dspace-rest-client", specifier = ">=0.1.17" }, + { name = "joblib", specifier = ">=1.5.3" }, { name = "jsonlines", specifier = ">=4.0.0" }, { name = "requests", specifier = ">=2.33.1" }, { name = "sentry-sdk", specifier = ">=2.34.1" }, - { name = "timdex-dataset-api", git = "https://github.com/MITLibraries/timdex-dataset-api" }, + { name = "timdex-dataset-api", git = "https://github.com/MITLibraries/timdex-dataset-api?rev=USE-531-fulltext-data-type-add" }, ] [package.metadata.requires-dev] dev = [ { name = "coveralls", specifier = ">=4.0.1" }, { name = "ipython", specifier = ">=9.13.0" }, + { name = "joblib-stubs", specifier = ">=1.5.3.1.20260117" }, { name = "mypy", specifier = ">=1.17.1" }, { name = "pip-audit", specifier = ">=2.9.0" }, { name = "pre-commit", specifier = ">=4.3.0" }, @@ -894,6 +898,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, ] +[[package]] +name = "joblib" +version = "1.5.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/f2/d34e8b3a08a9cc79a50b2208a93dce981fe615b64d5a4d4abee421d898df/joblib-1.5.3.tar.gz", hash = "sha256:8561a3269e6801106863fd0d6d84bb737be9e7631e33aaed3fb9ce5953688da3", size = 331603, upload-time = "2025-12-15T08:41:46.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl", hash = "sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713", size = 309071, upload-time = "2025-12-15T08:41:44.973Z" }, +] + +[[package]] +name = "joblib-stubs" +version = "1.5.3.1.20260117" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5e/4e/a4801406630955fe935d9d99efef72fd7398543f51fe5a787be80c235b4f/joblib_stubs-1.5.3.1.20260117.tar.gz", hash = "sha256:d420f9e26636df01e1bda851bf5bce9a8dbd6db84f08afe1dfdb963e6651e724", size = 14147, upload-time = "2026-01-17T11:16:43.666Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/7c/67ada9621b3c93caa75aba1a0b651dae988b22cdea93b33e7e49261364f5/joblib_stubs-1.5.3.1.20260117-py3-none-any.whl", hash = "sha256:ef9f570617b232263c8166f5b4351a5dcc90669d8b07e7e2cd7b25c3b0f26d5f", size = 21358, upload-time = "2026-01-17T11:16:44.327Z" }, +] + [[package]] name = "jsonlines" version = "4.0.0" @@ -1826,8 +1851,8 @@ wheels = [ [[package]] name = "timdex-dataset-api" -version = "5.0.0" -source = { git = "https://github.com/MITLibraries/timdex-dataset-api#6c1caa87c6c04a9f01996f7f74bec94ed72934c4" } +version = "5.1.0" +source = { git = "https://github.com/MITLibraries/timdex-dataset-api?rev=USE-531-fulltext-data-type-add#7dd5850ce2babdaf8414611268c4e5bcc700461b" } dependencies = [ { name = "attrs" }, { name = "boto3" },