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" },