diff --git a/dfh/cli.py b/dfh/cli.py index 729b45f..009362a 100644 --- a/dfh/cli.py +++ b/dfh/cli.py @@ -3,8 +3,12 @@ from datetime import timedelta import click +import jsonlines 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 logger = logging.getLogger(__name__) @@ -33,9 +37,8 @@ def main( def _log_command_elapsed_time() -> None: elapsed_time = time.perf_counter() - ctx.obj["start_time"] - logger.info( - "Total time to complete process: %s", str(timedelta(seconds=elapsed_time)) - ) + elapsed_time_display = timedelta(seconds=elapsed_time) + logger.info(f"Total time to complete process: {elapsed_time_display}") ctx.call_on_close(_log_command_elapsed_time) @@ -60,7 +63,7 @@ def test_dspace_connection( @click.pass_context @click.option( "--dataset-location", - required=True, + required=False, type=click.Path(), help=( "TIMDEX dataset location, e.g. 's3://timdex/dataset', " @@ -80,6 +83,13 @@ def test_dspace_connection( default=None, help="Maximum records of records to retrieve and harvest fulltext for.", ) +@click.option( + "--workers", + required=False, + type=int, + default=10, + help="Max number of parallel workers for downloading.", +) @click.option( "--output-jsonl", required=False, @@ -88,15 +98,43 @@ def test_dspace_connection( help="Write harvested fulltext to a local JSONLines file (primarily for testing).", ) def harvest( - ctx: click.Context, + _ctx: click.Context, dataset_location: str, run_id: str | None, record_limit: int | None, + workers: int, output_jsonl: str | None, ) -> None: """Harvest fulltext from DSpace and write to a TIMDEX dataset. - The argument --datast-location is required, as it provides the TIMDEX records to read - metadata from, and the location where fulltext will be written back to. + Flow: + 1. Get an iterator of source records from TIMDEX Dataset. + 2. Use to retrieve fulltext from DSpace, yielding as another iterator + 3. Write record + fulltext to TIMDEX dataset or JSONLines as output """ - raise NotImplementedError + # warm DSpace API authentication + warm_dspace_auth() + + # get iterator of target TIMDEX records + bitstream information + ttr = TIMDEXThesesRecords( + dataset_location=dataset_location, + run_id=run_id, + limit=record_limit, + ) + records_and_bitstreams = ttr.record_and_bitstream_metadata_iter() + + # get iterator of records + fulltext, 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." + ) + + with jsonlines.open(output_jsonl, "w") as writer: + for record in records_and_fulltexts: + writer.write(record) diff --git a/dfh/dspace.py b/dfh/dspace.py index ff5c683..3432d7a 100644 --- a/dfh/dspace.py +++ b/dfh/dspace.py @@ -1,10 +1,13 @@ import logging import os +import time from dspace_rest_client.client import DSpaceClient logger = logging.getLogger(__name__) +HTTP_OK = 200 + def get_dspace_client( *, @@ -27,6 +30,55 @@ def get_dspace_client( return client +def warm_dspace_auth( + *, + retry_attempts: int = 5, + initial_backoff_seconds: float = 2.0, + backoff_factor: float = 2.0, +) -> None: + """Warm and verify DSpace authentication before threaded downloads. + + Intermittent failures were observed for authentication against a cold API endpoint, + where after a successful authentication, all future requests were healthy. Could be + an artifact of server deploys and upgrades (this codebase is landing during a DSpace + migration), or an artifact of DSpace CRIS, unsure. Either way, this authentication + check + retry feels pretty harmless to layer on until we're 100% sure it's not needed. + """ + if retry_attempts < 1: + raise ValueError("retry_attempts must be at least 1") + + logger.info("Warming DSpace authentication for cold start") + last_error: Exception | None = None + for attempt in range(1, retry_attempts + 1): + try: + client = get_dspace_client(auth_on_init=True) + response = client.api_get(f"{client.API_ENDPOINT}/authn/status") + if ( + response.status_code == HTTP_OK + and response.json().get("authenticated") is True + ): + logger.info("DSpace authentication warm-up successful") + return + last_error = RuntimeError( + "DSpace authentication warm-up failed: " + f"{response.status_code} {response.text}" + ) + except Exception as exc: # noqa: BLE001 + last_error = exc + + if attempt < retry_attempts: + sleep_seconds = initial_backoff_seconds * (backoff_factor ** (attempt - 1)) + logger.warning( + "DSpace authentication warm-up attempt " + f"{attempt}/{retry_attempts} failed; sleeping " + f"{sleep_seconds:.1f} seconds before retry. Cause: {last_error}" + ) + time.sleep(sleep_seconds) + + msg = f"DSpace authentication warm-up failed after {retry_attempts} attempts" + raise RuntimeError(msg) from last_error + + def get_presigned_url_for_bitstream( dspace_client: DSpaceClient, bitstream_uuid: str, diff --git a/dfh/harvest.py b/dfh/harvest.py new file mode 100644 index 0000000..fff3093 --- /dev/null +++ b/dfh/harvest.py @@ -0,0 +1,134 @@ +import logging +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 dfh.dspace import get_dspace_client, get_presigned_url_for_bitstream + +logger = logging.getLogger(__name__) + +# module level, cross-thread object to hold thread specific DSpace client instances +# see: https://docs.python.org/3/library/threading.html#thread-local-data +threaded_dspace_clients = threading.local() + + +def record_and_fulltext_iter( + records: Iterator[dict], + *, + max_workers: int = 10, + log_progress_interval: int = 10, +) -> Iterator[dict]: + """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. + """ + 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( + record: dict, + *, + retry_attempts: int = 4, + initial_backoff_seconds: float = 1.0, + backoff_factor: float = 2.0, + timeout_seconds: int = 60, +) -> dict: + """Return a TIMDEX fulltext record for one source record. + + The primary work here is two network requests: + 1. Hit DSpace API for a presigned S3 URL for a bitstream UUID + 2. Use presigned S3 URL to download content + + This function is a worker function designed to be parallelized across threads. As + such, note the use of _get_dspace_client_for_thread() which ensures that it checks + the threading.local() object for a DSpaceClient instance to use that is unique and + reusable by this thread. + + Retries and backoffs are fairly simple: all exceptions, for either network request, + that bubble up are caught, logged, and increment the retry counter. + """ + bitstream_uuid = record["fulltext_bitstream"]["uuid"] + + fulltext = None + for attempt in range(1, retry_attempts + 1): + try: + # reuse dspace client for thread, or init if first time + dspace_client = _get_dspace_client_for_thread() + + # generate a pre-signed URL for a bitstream UUID + pre_signed_url = get_presigned_url_for_bitstream( + dspace_client, + bitstream_uuid, + ) + + # download bitstream content from S3 + response = requests.get(pre_signed_url, timeout=timeout_seconds) + response.raise_for_status() + fulltext = response.text + + # break out of retries loop if successful + break + + except Exception as exc: + if attempt == retry_attempts: + logger.exception( + f"Max retries of {retry_attempts} encountered, failed to download " + 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"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, + } + + +def _get_dspace_client_for_thread() -> DSpaceClient: + """Get thread's DSpaceClient from module level threaded_dspace_clients.""" + dspace_client = getattr(threaded_dspace_clients, "dspace_client", None) + if dspace_client is None: + dspace_client = get_dspace_client() + threaded_dspace_clients.dspace_client = dspace_client + return dspace_client diff --git a/pyproject.toml b/pyproject.toml index a1d656a..bdef0f2 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", + "jsonlines>=4.0.0", "requests>=2.33.1", "sentry-sdk>=2.34.1", "timdex-dataset-api", diff --git a/tests/test_dspace.py b/tests/test_dspace.py new file mode 100644 index 0000000..f0b7f67 --- /dev/null +++ b/tests/test_dspace.py @@ -0,0 +1,15 @@ +from unittest import mock + +from dfh.dspace import warm_dspace_auth + + +def test_warm_dspace_auth_retries_after_401_and_succeeds(): + fail_auth = mock.Mock(status_code=401, text="Unauthorized") + success_auth = mock.Mock(status_code=200, text="OK") + success_auth.json.return_value = {"authenticated": True} + + client = mock.Mock() + client.api_get.side_effect = [fail_auth, success_auth] + + with mock.patch("dfh.dspace.get_dspace_client", return_value=client): + assert warm_dspace_auth(initial_backoff_seconds=0.1) is None diff --git a/uv.lock b/uv.lock index 38fc036..affc26f 100644 --- a/uv.lock +++ b/uv.lock @@ -555,6 +555,7 @@ source = { editable = "." } dependencies = [ { name = "click" }, { name = "dspace-rest-client" }, + { name = "jsonlines" }, { name = "requests" }, { name = "sentry-sdk" }, { name = "timdex-dataset-api" }, @@ -575,6 +576,7 @@ dev = [ requires-dist = [ { name = "click", specifier = ">=8.2.1" }, { name = "dspace-rest-client", specifier = ">=0.1.17" }, + { 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" }, @@ -892,6 +894,18 @@ 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 = "jsonlines" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/87/bcda8e46c88d0e34cad2f09ee2d0c7f5957bccdb9791b0b934ec84d84be4/jsonlines-4.0.0.tar.gz", hash = "sha256:0c6d2c09117550c089995247f605ae4cf77dd1533041d366351f6f298822ea74", size = 11359, upload-time = "2023-09-01T12:34:44.187Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/62/d9ba6323b9202dd2fe166beab8a86d29465c41a0288cbe229fac60c1ab8d/jsonlines-4.0.0-py3-none-any.whl", hash = "sha256:185b334ff2ca5a91362993f42e83588a360cf95ce4b71a73548502bda52a7c55", size = 8701, upload-time = "2023-09-01T12:34:42.563Z" }, +] + [[package]] name = "librt" version = "0.10.0"