From 8d11fc6dc3e4deae7debebded2eb688cc4fb6f68 Mon Sep 17 00:00:00 2001 From: jaleman-vdr-wikimedia Date: Wed, 29 Oct 2025 12:25:15 -0600 Subject: [PATCH 1/2] Add parallel decompression, streaming reads, and stoppable callbacks - Replace standard gzip decompression in read_all with zlib_ng_threaded using threads=-1 for parallel processing, improving performance on multi-core systems. - Modify read_all to use tarfile's streaming mode (r|) and iterate with tar.next(), removing the need to read the full member list and preventing seeks on the non-seekable stream. - Implement true end-to-end streaming in read_all by removing the io.BytesIO(f.read()) buffering, making it safe for large files within archives. - Add the _TarfileStreamWrapper helper class to bridge compatibility issues between tarfile's streaming file objects and io.TextIOWrapper (missing .seekable(), .closed, .flush()). - Update ReadCallback type hint to Callable[[dict], bool]. - Modify _read_loop, _subscribe_to_entity, and read_all to check the boolean return value of the callback, allowing users to gracefully stop processing early. - Added try... block to teardown class in integration suite, so clean up process is safer after tests are executed - Removed hardcoded date from integration suite and replaced it with automatic date determination --- modules/api/api_client.py | 105 +++++++++++++++--- modules/api/test/api_client_test.py | 2 +- .../api/test/test_api_client_integration.py | 23 ++-- 3 files changed, 102 insertions(+), 28 deletions(-) diff --git a/modules/api/api_client.py b/modules/api/api_client.py index a878068..1c696dd 100644 --- a/modules/api/api_client.py +++ b/modules/api/api_client.py @@ -15,9 +15,14 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Callable, Dict, List, Optional, Union +import typing +import gzip +from zlib_ng import gzip_ng_threaded import httpx from .exceptions import APIDataError, APIRequestError, APIStatusError +ReadCallback = Callable[[dict], bool] + logger = logging.getLogger(__name__) DATE_FORMAT = "%Y-%m-%d" @@ -25,12 +30,12 @@ class Filter: """Represents a simple key-value filter for an API query.""" - def __init__(self, field: str, value: str): + def __init__(self, field: str, value: Any): """Initializes a Filter object. Args: field (str): The name of the field to filter on. - value (str): The value to filter for. + value (Any): The value to filter for. """ self.field = field self.value = value @@ -117,6 +122,50 @@ def to_json(self): # Remove keys with None or empty values return {k: v for k, v in result.items() if v not in [None, [], {}, '']} +class _TarfileStreamWrapper: + """ + Wraps the non-seekable file object from tarfile.extractfile() + to make it compatible with io.TextIOWrapper, which expects + a .seekable() method to exist. + """ + def __init__(self, tarfile_stream: typing.IO[bytes]): + self._stream = tarfile_stream + + @property + def closed(self) -> bool: + """Returns True if the underlying stream is closed, False otherwise.""" + return self._stream.closed + + def read(self, *args, **kwargs): + """Reads and returns data from the underlying stream, passing along any arguments.""" + return self._stream.read(*args, **kwargs) + + def readable(self): + """Returns True to indicate the stream is readable.""" + return True + + def seekable(self): + """ + Returns False to indicate the stream is not seekable. + This is the primary purpose of this wrapper. + """ + return False + + def writable(self): + """Returns False to indicate the stream is not writable.""" + return False + + def close(self): + """Closes the underlying tarfile stream.""" + self._stream.close() + + def flush(self): + """ + Flushes the write buffers of the underlying stream, if applicable. + This method was missing. + """ + self._stream.flush() + class Client: """ The main client for interacting with the Wikimedia Enterprise API. @@ -248,7 +297,7 @@ def _get_entity(self, req: Optional[Request], path: str, val: Any): else: raise APIDataError("Mismatched types between expected container and JSON response.") - def _read_loop(self, rdr: io.BytesIO, cbk: Callable[[dict], Any]): + def _read_loop(self, rdr: typing.BinaryIO, cbk: Callable[[dict], Any]): """ Processes a byte stream of newline-delimited JSON (NDJSON). @@ -266,10 +315,13 @@ def _read_loop(self, rdr: io.BytesIO, cbk: Callable[[dict], Any]): continue try: article = json.loads(line) - cbk(article) + if not cbk(article): + return False except json.JSONDecodeError: logger.warning("Skipping line due to JSON decode error", exc_info=True) + return True + def _read_entity(self, path: str, cbk: Callable[[dict], Any]): """ Internal helper to fetch a resource and process it as NDJSON. @@ -411,7 +463,8 @@ def _subscribe_to_entity(self, path: str, req: Request, cbk: Callable[[dict], An if line: try: article=json.loads(line) - cbk(article) + if not cbk(article): + break except json.JSONDecodeError: logger.warning("Skipping malformed JSON line in stream: %s", line) except httpx.HTTPStatusError as e: @@ -419,7 +472,7 @@ def _subscribe_to_entity(self, path: str, req: Request, cbk: Callable[[dict], An except httpx.RequestError as e: raise APIRequestError(f"Stream Request Error: {e}", request=e.request) from e - def read_all(self, rdr: io.BytesIO, cbk: Callable[[dict], Any]): + def read_all(self, rdr: io.BytesIO, cbk: ReadCallback): """ Reads a .tar.gz archive containing NDJSON files. @@ -436,14 +489,32 @@ def read_all(self, rdr: io.BytesIO, cbk: Callable[[dict], Any]): APIDataError: If the archive is corrupt or cannot be read as a tarfile. """ try: - with tarfile.open(fileobj=rdr, mode='r:gz') as tar: - for member in tar.getmembers(): - f = tar.extractfile(member) - if f: - with f: - self._read_loop(io.BytesIO(f.read()), cbk) + with gzip_ng_threaded.open(rdr, mode="rb", threads=-1) as decompressed_stream: + + typed_stream = typing.cast(typing.BinaryIO, decompressed_stream) + + with tarfile.open(fileobj=typed_stream, mode='r|') as tar: + while True: + member = tar.next() + if member is None: + break + + if not member.isfile(): + continue + + f = tar.extractfile(member) + + if f: + with f: + wrapped_stream = _TarfileStreamWrapper(f) + typed_stream = typing.cast(typing.BinaryIO, wrapped_stream) + if not self._read_loop(typed_stream, cbk): + break except tarfile.TarError as e: raise APIDataError(f"Failed to read tar archive: {e}") from e + except gzip.BadGzipFile as e: + raise APIDataError(f"Failed to decompress Gzip archive: {e}") from e + def set_access_token(self, token: str): """ @@ -525,7 +596,7 @@ def head_batch(self, timestamp: datetime.datetime, idr: str) -> dict: """Retrieves metadata for a specific data batch.""" return self._head_entity(f"{self._get_batches_prefix(timestamp)}/{idr}/download") - def read_batch(self, timestamp: datetime.datetime, idr: str, cbk: Callable[[dict], Any]): + def read_batch(self, timestamp: datetime.datetime, idr: str, cbk: ReadCallback): """Reads and processes the content of a specific data batch via a callback.""" self._read_entity(f"b{self._get_batches_prefix(timestamp)}/{idr}/download", cbk) @@ -549,7 +620,7 @@ def head_snapshot(self, idr: str) -> dict: """Retrieves metadata for a snapshot""" return self._head_entity(f"snapshots/{idr}/download") - def read_snapshot(self, idr: str, cbk: Callable[[dict], Any]): + def read_snapshot(self, idr: str, cbk: ReadCallback): """Reads a snapshot""" self._read_entity(f"snapshots/{idr}/download", cbk) @@ -573,7 +644,7 @@ def head_chunk(self, sid: str, idr: str) -> dict: """Retrieves a chunk's metadata""" return self._head_entity(f"snapshots/{sid}/chunks/{idr}/download") - def read_chunk(self, sid: str, idr: str, cbk: Callable[[dict], Any]): + def read_chunk(self, sid: str, idr: str, cbk: ReadCallback): """Reads a chunk""" self._read_entity(f"snapshots/{sid}/chunks/{idr}/download", cbk) @@ -609,7 +680,7 @@ def head_structured_snapshot(self, idr: str) -> dict: """Retrieves a structured snapshot's metadata""" return self._head_entity(f"snapshots/structured-contents/{idr}/download") - def read_structured_snapshot(self, idr: str, cbk: Callable[[dict], Any]): + def read_structured_snapshot(self, idr: str, cbk: ReadCallback): """Reads a structured snapshot""" self._read_entity(f"snapshots/structured-contents/{idr}/download", cbk) @@ -617,6 +688,6 @@ def download_structured_snapshot(self, idr: str, writer: io.BytesIO): """Downloads a structured snapshot""" self._download_entity(f"snapshots/structured-contents/{idr}/download", writer) - def stream_articles(self, req: Request, cbk: Callable[[dict], Any]): + def stream_articles(self, req: Request, cbk: ReadCallback): """Streams rt articles""" self._subscribe_to_entity("articles", req, cbk) diff --git a/modules/api/test/api_client_test.py b/modules/api/test/api_client_test.py index b14023c..71f81a6 100644 --- a/modules/api/test/api_client_test.py +++ b/modules/api/test/api_client_test.py @@ -564,7 +564,7 @@ def test_read_all_raises_on_corrupt_tar_archive(self): mock_cbk = MagicMock() - with self.assertRaisesRegex(APIDataError, "Failed to read tar archive"): + with self.assertRaisesRegex(APIDataError, "Failed to decompress Gzip archive"): self.client.read_all(corrupt_data, mock_cbk) mock_cbk.assert_not_called() diff --git a/modules/api/test/test_api_client_integration.py b/modules/api/test/test_api_client_integration.py index cb1398d..795b93d 100644 --- a/modules/api/test/test_api_client_integration.py +++ b/modules/api/test/test_api_client_integration.py @@ -17,7 +17,7 @@ import unittest import os import io -from datetime import datetime +from datetime import datetime, timedelta, timezone from modules.api.api_client import Client, Request, Filter from modules.auth.auth_client import AuthClient from modules.auth.helper import Helper @@ -72,12 +72,18 @@ def setUpClass(cls): def tearDownClass(cls): """Clean up resources after all tests are done""" print("\nTearing down integration test client...") - if hasattr(cls, 'client'): - cls.client.http_client.close() + try: + if hasattr(cls, 'client'): + cls.client.http_client.close() + except Exception as e: + print(f" - Warning: Failed to close api_client's http_client: {e}") - if hasattr(cls, 'helper') and hasattr(cls.helper, 'stop'): - print("Stopping auth helper thread...") - cls.helper.stop() + try: + if hasattr(cls, 'helper') and hasattr(cls.helper, 'stop'): + print("Stopping auth helper thread...") + cls.helper.stop() + except Exception as e: + print(f" - Warning: Failed to stop auth helper (e.g., token revoke failed): {e}") def test_get_projects_smoke_test(self): """ @@ -165,10 +171,7 @@ def test_head_and_download_batch(self): """ print("Running test_head_and_download_batch...") - # This date gets a recent batch. - # Should you see a 404 due to age, - # you can update this date. - batch_time = datetime(2025, 10, 26, 12) # YYYY, M, D, H + batch_time = datetime.now(timezone.utc) - timedelta(hours=3) req = Request(limit=1) print(f" ... Finding a batch from {batch_time}...") From 13568fce970bd8aa1fb11a25d9f2d95bf9a4bde3 Mon Sep 17 00:00:00 2001 From: jaleman-vdr-wikimedia Date: Thu, 30 Oct 2025 16:04:42 -0600 Subject: [PATCH 2/2] Implement stoppable callbacks, fix examples, and bug fix - Fixed batches, snapshots, and streaming examples to accommodate and showcase stoppable callbacks. - Added a "callback" example, to specifically showcase stoppable callbacks, along its README doc. - Discovered that running examples back to back, causes issues during token revocation cleanup. Due to this, edited auth_client, to handle errors during token revocation. --- example/batches/batches.py | 4 +- example/callback/README.md | 39 ++++++++++ example/callback/callback.py | 129 +++++++++++++++++++++++++++++++++ example/snapshots/snapshots.py | 4 +- example/streaming/streaming.py | 2 + modules/auth/auth_client.py | 14 +++- 6 files changed, 189 insertions(+), 3 deletions(-) create mode 100644 example/callback/README.md create mode 100644 example/callback/callback.py diff --git a/example/batches/batches.py b/example/batches/batches.py index 10cd2a9..9def1dd 100644 --- a/example/batches/batches.py +++ b/example/batches/batches.py @@ -106,11 +106,13 @@ def main(): articles_found = [] - def article_callback(article_json): + def article_callback(article_json) -> bool: """A simple callback to process one article from the batch.""" if 'identifier' in article_json: articles_found.append(article_json['identifier']) + return True + api_client.read_all(buffer, article_callback) logger.info("Successfully processed %s articles from the batch.", len(articles_found)) diff --git a/example/callback/README.md b/example/callback/README.md new file mode 100644 index 0000000..b7f8752 --- /dev/null +++ b/example/callback/README.md @@ -0,0 +1,39 @@ +# Real-Time Stream Callback Demo + +This script demonstrates the ability to programmatically stop a real-time stream from within the processing callback. + +It connects to the `/v2/articles` real-time stream, processes articles one by one, and automatically stops and disconnects after receiving a predefined number of articles (5, by default). + +## Key Features +- Connects to the real-time article stream (`api_client.stream_articles`). + +- Implements a custom callback function (`stream_callback`). + +- Uses the callback's bool return value to control the stream (returning `False` stops the stream). + +- Handles authentication and graceful token revocation using `AuthClient` and `Helper`. + +## How it Works +The core logic of this demo is inside the `stream_callback` function. The `api_client`'s internal streaming loop checks the boolean value returned by this callback after every article it processes. + +- A list, `articles_received_tracker`, is used as a counter (it's a list so it can be mutated from within the callback). + +- The callback receives an `article` (a `dict`) from the stream. + +- It logs the article's details and appends it to the tracker. + +- It checks the count: `if len(articles_received_tracker) >= STOP_AFTER_N_ARTICLES:` + +- If the count is met, it logs a warning and returns `False`. + +- If the count is not met, it returns `True`. + +A `True` value tells the `api_client` to "keep processing." A `False` value tells the `api_client` to "stop immediately," at which point it closes the stream connection and the `api_client.stream_articles()` function returns. + +## Running the Application + +To run the application, use the following command from the project's root: + +```sh +python -m example.callback.callback +``` diff --git a/example/callback/callback.py b/example/callback/callback.py new file mode 100644 index 0000000..e5fa02e --- /dev/null +++ b/example/callback/callback.py @@ -0,0 +1,129 @@ +# pylint: disable=W0718, R0914, R0801, W0612 + +""" +Demonstrates the ability to stop a client callback midway through processing. + +This script connects to the real-time article stream and uses the callback's +boolean return value to stop the stream after receiving 5 articles. +""" + +import logging +import time + +# --- Import custom modules --- +from modules.auth.auth_client import AuthClient +from modules.auth.helper import Helper +from modules.api.api_client import Client, Request +from modules.api.exceptions import APIRequestError, APIStatusError, APIDataError + +# --- Setup Logging --- +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +# --- Configuration --- +STOP_AFTER_N_ARTICLES = 5 + +def main(): + """Runs the callback stop demo""" + helper = None + auth_client = None + + articles_received_tracker = [] + + def stream_callback(article: dict) -> bool: + """ + Callback function to process streamed articles. + + This callback will log the received article and check if it's + time to stop the stream. + + Args: + article (dict): The JSON object for the article. + + Returns: + bool: True to continue processing, False to stop the stream. + """ + try: + article_name = article.get('name', article.get('identifier', 'Unknown')) + event_id = article.get('event', {}).get('identifier', 'unknown_event') + + logger.info( + "[%s] Received article (event: %s): %s", + len(articles_received_tracker) + 1, + event_id, + article_name + ) + + articles_received_tracker.append(article) + + # --- The Stop Logic --- + if len(articles_received_tracker) >= STOP_AFTER_N_ARTICLES: + logger.warning( + "Reached stop limit of %s articles. Returning False to stop stream.", + STOP_AFTER_N_ARTICLES + ) + return False + + except Exception as e: + logger.error("Error within callback function: %s", e) + return False + return True + + try: + # --- Authentication Setup --- + logger.info("Setting up authentication...") + auth_client = AuthClient() + helper = Helper(auth_client) + + api_client = Client(timeout=3600.0) + + token = helper.get_access_token() + api_client.set_access_token(token) + logger.info("Succesfully authenticated!") + + # --- Stream Demonstration --- + logger.info("\nStarting real-time stream callback demo...") + + stream_req = Request( + fields=["name", "abstract", "event.*"] + ) + + logger.info( + "Connecting to real-time article stream (will stop after %s articles)...", + STOP_AFTER_N_ARTICLES + ) + + start_time = time.time() + + api_client.stream_articles(stream_req, stream_callback) + + end_time = time.time() + + logger.info( + "Stream processing finished. Total articles received: %s", + len(articles_received_tracker) + ) + logger.info("Stream was active for %.2f seconds.", end_time - start_time) + logger.info("\n--- Callback stop demo complete ---") + + except (APIRequestError, APIStatusError, APIDataError) as e: + logger.fatal("API Error encountered: %s", e) + if isinstance(e, APIStatusError) and e.response and e.response.status_code == 401: + logger.error("Got 401 Unauthorized. Check your token permissions for the real-time stream.") + except ValueError as e: + logger.fatal("Configuration Error (check .env): %s", e) + except KeyboardInterrupt: + logger.info("\nUser interrupted stream. Shutting down.") + except Exception as e: + logger.fatal("An unexpected error ocurred: %s", e, exc_info=True) + finally: + # --- Graceful Shutdown --- + if helper: + logger.info("Shutting down helper and revoking tokens...") + helper.stop() + elif auth_client: + auth_client.close() + logger.info("Exiting!") + +if __name__ == "__main__": + main() diff --git a/example/snapshots/snapshots.py b/example/snapshots/snapshots.py index fb53438..cd39935 100644 --- a/example/snapshots/snapshots.py +++ b/example/snapshots/snapshots.py @@ -108,13 +108,15 @@ def main(): buffer.seek(0) articles_found_in_snapshot = [] - def snapshot_article_callback(article_json): + def snapshot_article_callback(article_json) -> bool: """Simple callback to collect article names""" if 'name' in article_json: articles_found_in_snapshot.append(article_json['name']) elif 'identifier' in article_json: articles_found_in_snapshot.append(article_json['identifier']) + return True + try: api_client.read_all(buffer, snapshot_article_callback) logger.info("Succesfully processed %s articles from snapshot.", len(articles_found_in_snapshot)) diff --git a/example/streaming/streaming.py b/example/streaming/streaming.py index ba77897..b4c0b57 100644 --- a/example/streaming/streaming.py +++ b/example/streaming/streaming.py @@ -58,6 +58,8 @@ def article_callback(article): logger.info("event.identifiers: %s", article.get('event', {}).get('identifier')) logger.info("-----------END------------\n\n\n") + return True + def main(): """Main execution function to initiate the article stream. diff --git a/modules/auth/auth_client.py b/modules/auth/auth_client.py index 0ddc46c..4ebecc6 100644 --- a/modules/auth/auth_client.py +++ b/modules/auth/auth_client.py @@ -7,6 +7,7 @@ import os import json +import logging from threading import Lock from datetime import datetime, timedelta from dotenv import load_dotenv @@ -15,6 +16,9 @@ # Load environment variables from .env file load_dotenv() +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + class AuthClient: """Manages authentication and token lifecycle for the Wikimedia Enterprise API.""" @@ -77,7 +81,15 @@ def refresh_token(self, refresh_token): def revoke_token(self, refresh_token): """Revokes a refresh token, invalidating the current session.""" data = {"refresh_token": refresh_token} - self._post("/token-revoke", data) + + try: + self._post("/token-revoke", data) + logger.info("Token revoked succesfully!") + except (httpx.ReadTimeout, httpx.NetworkError, httpx.HTTPStatusError) as e: + logger.warning( + "Failed to revoke token: %s." + " If running examples back-to-back, this is expected.", e + ) def get_access_token(self): """