USE 557 - Parallelized input record read and bitstream download

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)

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,

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

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

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