From d6e700d783131d5a27cc57cef9b1e60444fe9abc Mon Sep 17 00:00:00 2001 From: Neukz Date: Mon, 15 Jun 2026 11:22:05 +0200 Subject: [PATCH 1/2] feat: add PyPI Stats API client with retries --- src/repolytics/ingestion/_http.py | 19 ++++ src/repolytics/ingestion/github_client.py | 21 +---- src/repolytics/ingestion/pypi_client.py | 85 ++++++++++++++++++ src/repolytics/ingestion/watermarks.py | 30 +++++++ src/repolytics/ingestion/writer.py | 45 ++++++++++ tests/unit/_clients.py | 27 ++++++ tests/unit/test_github_client.py | 44 ++++----- tests/unit/test_pypi_client.py | 105 ++++++++++++++++++++++ tests/unit/test_watermarks.py | 33 +++++++ tests/unit/test_writer.py | 58 ++++++++++++ 10 files changed, 424 insertions(+), 43 deletions(-) create mode 100644 src/repolytics/ingestion/_http.py create mode 100644 src/repolytics/ingestion/pypi_client.py create mode 100644 src/repolytics/ingestion/watermarks.py create mode 100644 src/repolytics/ingestion/writer.py create mode 100644 tests/unit/_clients.py create mode 100644 tests/unit/test_pypi_client.py create mode 100644 tests/unit/test_watermarks.py create mode 100644 tests/unit/test_writer.py diff --git a/src/repolytics/ingestion/_http.py b/src/repolytics/ingestion/_http.py new file mode 100644 index 0000000..0245c00 --- /dev/null +++ b/src/repolytics/ingestion/_http.py @@ -0,0 +1,19 @@ +"""Shared HTTP helpers for the ingestion API clients.""" + +import random + +import httpx + +# HTTP statuses worth retrying. +RETRYABLE_STATUSES = frozenset({429, 500, 502, 503}) + + +def retry_delay(response: httpx.Response, attempt: int, backoff_base: float) -> float: + """Seconds to wait before a retry: `Retry-After` if present, else backoff.""" + retry_after = response.headers.get("Retry-After") + if retry_after is not None: + try: + return float(retry_after) + except ValueError: + pass + return backoff_base * 2**attempt + random.uniform(0, backoff_base) diff --git a/src/repolytics/ingestion/github_client.py b/src/repolytics/ingestion/github_client.py index 1c25e39..e9956bd 100644 --- a/src/repolytics/ingestion/github_client.py +++ b/src/repolytics/ingestion/github_client.py @@ -1,10 +1,9 @@ """GitHub REST API client. -Methods return raw JSON exactly as received. -Handles authentication, pagination, rate limiting, and retries with exponential backoff. +Methods return raw JSON exactly as received. Handles authentication, +pagination, rate limiting, and retries with exponential backoff. """ -import random import time from collections.abc import Callable from datetime import UTC, datetime @@ -12,9 +11,7 @@ import httpx from repolytics.config import get_settings - -# HTTP statuses worth retrying. -RETRYABLE_STATUSES = frozenset({429, 500, 502, 503}) +from repolytics.ingestion._http import RETRYABLE_STATUSES, retry_delay def _to_iso(dt: datetime) -> str: @@ -167,22 +164,12 @@ def _request( for attempt in range(self._max_retries): if response.status_code not in RETRYABLE_STATUSES: break - self._sleep(self._retry_delay(response, attempt)) + self._sleep(retry_delay(response, attempt, self._backoff_base)) response = self._client.request(method, url, params=params) response.raise_for_status() self._guard_rate_limit(response) return response - def _retry_delay(self, response: httpx.Response, attempt: int) -> float: - """Seconds to wait before a retry: `Retry-After` if present, else backoff.""" - retry_after = response.headers.get("Retry-After") - if retry_after is not None: - try: - return float(retry_after) - except ValueError: - pass - return self._backoff_base * 2**attempt + random.uniform(0, self._backoff_base) - def _guard_rate_limit(self, response: httpx.Response) -> None: """Sleep until the rate-limit window resets when remaining calls run low.""" remaining = response.headers.get("X-RateLimit-Remaining") diff --git a/src/repolytics/ingestion/pypi_client.py b/src/repolytics/ingestion/pypi_client.py new file mode 100644 index 0000000..c65d59d --- /dev/null +++ b/src/repolytics/ingestion/pypi_client.py @@ -0,0 +1,85 @@ +"""PyPI Stats API client. + +Methods return raw JSON exactly as received. Handles retries with exponential +backoff and keeps a courtesy delay between requests to stay within PyPI's +informal rate limits. +""" + +import time +from collections.abc import Callable + +import httpx + +from repolytics.ingestion._http import RETRYABLE_STATUSES, retry_delay + + +class PyPIClient: + """Synchronous client for the PyPI Stats API. + + Wraps `httpx.Client`. Usable as a context manager so the underlying + connection pool is closed on exit. + """ + + def __init__( + self, + *, + base_url: str = "https://pypistats.org/api", + min_interval: float = 1.0, + max_retries: int = 3, + backoff_base: float = 1.0, + timeout: httpx.Timeout | None = None, + transport: httpx.BaseTransport | None = None, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + self._min_interval = min_interval + self._max_retries = max_retries + self._backoff_base = backoff_base + self._sleep = sleep + self._client = httpx.Client( + base_url=base_url, + headers={"Accept": "application/json"}, + timeout=timeout + or httpx.Timeout(connect=30.0, read=60.0, write=60.0, pool=60.0), + transport=transport, + ) + + # --- Lifecycle --- + + def close(self) -> None: + """Close the underlying HTTP connection pool.""" + self._client.close() + + def __enter__(self) -> "PyPIClient": + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + # --- Endpoints --- + + def get_recent_downloads(self, package: str) -> dict: + """`GET /packages/{package}/recent` (last day/week/month totals).""" + return self._request("GET", f"/packages/{package}/recent").json() + + def get_overall_downloads(self, package: str, *, mirrors: bool = False) -> dict: + """`GET /packages/{package}/overall` (daily download time series).""" + params = {"mirrors": str(mirrors).lower()} + return self._request( + "GET", f"/packages/{package}/overall", params=params + ).json() + + # --- Internals --- + + def _request( + self, method: str, url: str, params: dict[str, object] | None = None + ) -> httpx.Response: + """Issue one request, retrying transient failures, then pause politely.""" + response = self._client.request(method, url, params=params) + for attempt in range(self._max_retries): + if response.status_code not in RETRYABLE_STATUSES: + break + self._sleep(retry_delay(response, attempt, self._backoff_base)) + response = self._client.request(method, url, params=params) + response.raise_for_status() + self._sleep(self._min_interval) + return response diff --git a/src/repolytics/ingestion/watermarks.py b/src/repolytics/ingestion/watermarks.py new file mode 100644 index 0000000..50c6294 --- /dev/null +++ b/src/repolytics/ingestion/watermarks.py @@ -0,0 +1,30 @@ +"""Watermark store for incremental extraction. + +Tracks the last-seen timestamp per `{source}/{endpoint}/{repo}` as a flat JSON +map on disk, so each ingestion run only fetches records newer than last time. +""" + +import json +from pathlib import Path + + +def watermark_key(source: str, endpoint: str, repo: str) -> str: + """Build the watermark map key for a source/endpoint/repo target.""" + return f"{source}/{endpoint}/{repo}" + + +def load_watermarks(path: str | Path) -> dict[str, str]: + """Load the watermark map, returning an empty dict when the file is absent.""" + file = Path(path) + if not file.exists(): + return {} + return json.loads(file.read_text(encoding="utf-8")) + + +def save_watermarks(watermarks: dict[str, str], path: str | Path) -> None: + """Atomically write the watermark map to `path`, creating parents as needed.""" + file = Path(path) + file.parent.mkdir(parents=True, exist_ok=True) + tmp = file.with_name(f"{file.name}.tmp") + tmp.write_text(json.dumps(watermarks, indent=2), encoding="utf-8") + tmp.replace(file) diff --git a/src/repolytics/ingestion/writer.py b/src/repolytics/ingestion/writer.py new file mode 100644 index 0000000..d29514f --- /dev/null +++ b/src/repolytics/ingestion/writer.py @@ -0,0 +1,45 @@ +"""Parquet writer for raw API responses. + +Lands each record as a single JSON `data` column plus a `_loaded_at` timestamp. +""" + +import json +from datetime import UTC, datetime +from pathlib import Path + +import polars as pl + +# Raw-landing schema. +_SCHEMA = {"data": pl.Utf8, "_loaded_at": pl.Datetime(time_unit="us", time_zone="UTC")} + + +def write_parquet(records: list[dict], path: str | Path) -> Path: + """Write `records` to `path` as Parquet (JSON `data` column + `_loaded_at`). + + Each record is serialized to a JSON string and stamped with a single batch + timestamp. An empty `records` list still writes a zero-row file with the + correct schema, so the partition stays present and schema-stable. + """ + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + loaded_at = datetime.now(UTC) + frame = pl.DataFrame( + { + "data": [json.dumps(record, ensure_ascii=False) for record in records], + "_loaded_at": [loaded_at] * len(records), + }, + schema=_SCHEMA, + ) + frame.write_parquet(out) + return out + + +def partition_path( + root: str | Path, source: str, table: str, date: datetime | str +) -> Path: + """Build the date-partitioned landing path for a source/table on a date. + + Returns `{root}/{source}/{table}/{YYYY-MM-DD}.parquet`. + """ + day = date.strftime("%Y-%m-%d") if isinstance(date, datetime) else date + return Path(root) / source / table / f"{day}.parquet" diff --git a/tests/unit/_clients.py b/tests/unit/_clients.py new file mode 100644 index 0000000..437ea77 --- /dev/null +++ b/tests/unit/_clients.py @@ -0,0 +1,27 @@ +"""Builders for mock-transport-backed API clients.""" + +from collections.abc import Callable + +import httpx + + +def mock_client( + client_cls: type, + handler: Callable[[httpx.Request], httpx.Response], + *args: object, + **kwargs: object, +) -> tuple[object, list[float]]: + """Build an ingestion client over a `MockTransport` with a spy `sleep`. + + Returns the client and the list recording every sleep duration, so tests can + assert on retry/courtesy waits without blocking. + """ + sleeps: list[float] = [] + client = client_cls( + *args, + transport=httpx.MockTransport(handler), + sleep=sleeps.append, + backoff_base=0.0, + **kwargs, + ) + return client, sleeps diff --git a/tests/unit/test_github_client.py b/tests/unit/test_github_client.py index 493e39a..eca5da0 100644 --- a/tests/unit/test_github_client.py +++ b/tests/unit/test_github_client.py @@ -7,7 +7,9 @@ import httpx import pytest +from repolytics.ingestion._http import RETRYABLE_STATUSES from repolytics.ingestion.github_client import GitHubClient, _to_iso +from tests.unit._clients import mock_client # Absolute next-page URL used to drive pagination in tests. NEXT_URL = "https://api.github.com/x?page=2" @@ -17,25 +19,13 @@ def make_client( handler: Callable[[httpx.Request], httpx.Response], **kwargs: object ) -> tuple[GitHubClient, list[float]]: - """Build a client backed by a `MockTransport` and a spy `sleep`. - - Returns the client and the list that records every sleep duration, so tests - can assert on retry/rate-limit waits without ever blocking. - """ - sleeps: list[float] = [] - client = GitHubClient( - "test-token", - transport=httpx.MockTransport(handler), - sleep=sleeps.append, - backoff_base=0.0, - **kwargs, - ) - return client, sleeps + """Build a `GitHubClient` over a `MockTransport` with a spy `sleep`.""" + return mock_client(GitHubClient, handler, "test-token", **kwargs) -def test_to_iso_appends_z_for_naive_and_aware() -> None: - assert _to_iso(datetime(2024, 1, 1)) == "2024-01-01T00:00:00Z" - assert _to_iso(datetime(2024, 1, 1, tzinfo=UTC)) == "2024-01-01T00:00:00Z" +@pytest.mark.parametrize("dt", [datetime(2024, 1, 1), datetime(2024, 1, 1, tzinfo=UTC)]) +def test_to_iso_appends_z(dt: datetime) -> None: + assert _to_iso(dt) == "2024-01-01T00:00:00Z" def test_paginates_following_next_link() -> None: @@ -77,23 +67,24 @@ def handler(request: httpx.Request) -> httpx.Response: assert sleeps[0] > 0 -def test_retries_then_succeeds() -> None: - statuses = [502, 502, 200] +@pytest.mark.parametrize("status", sorted(RETRYABLE_STATUSES)) +def test_retryable_status_retries_then_succeeds(status: int) -> None: + statuses = [status, 200] calls: list[httpx.Request] = [] def handler(request: httpx.Request) -> httpx.Response: - status = statuses[len(calls)] + current = statuses[len(calls)] calls.append(request) - if status == 200: + if current == 200: return httpx.Response(200, json={"id": 1}) - return httpx.Response(status) + return httpx.Response(current) client, sleeps = make_client(handler) result = client.get_repository("o", "r") assert result == {"id": 1} - assert len(calls) == 3 - assert len(sleeps) == 2 + assert len(calls) == 2 # one failure, one retry + assert len(sleeps) == 1 def test_raises_after_max_retries() -> None: @@ -107,12 +98,13 @@ def handler(request: httpx.Request) -> httpx.Response: assert len(sleeps) == 3 -def test_404_raises_without_retry() -> None: +@pytest.mark.parametrize("status", [400, 401, 403, 404]) +def test_non_retryable_status_raises_without_retry(status: int) -> None: calls: list[httpx.Request] = [] def handler(request: httpx.Request) -> httpx.Response: calls.append(request) - return httpx.Response(404) + return httpx.Response(status) client, sleeps = make_client(handler) with pytest.raises(httpx.HTTPStatusError): diff --git a/tests/unit/test_pypi_client.py b/tests/unit/test_pypi_client.py new file mode 100644 index 0000000..4dceccf --- /dev/null +++ b/tests/unit/test_pypi_client.py @@ -0,0 +1,105 @@ +"""Tests for repolytics.ingestion.pypi_client.PyPIClient.""" + +from collections.abc import Callable + +import httpx +import pytest + +from repolytics.ingestion._http import RETRYABLE_STATUSES +from repolytics.ingestion.pypi_client import PyPIClient +from tests.unit._clients import mock_client + + +def make_client( + handler: Callable[[httpx.Request], httpx.Response], **kwargs: object +) -> tuple[PyPIClient, list[float]]: + """Build a `PyPIClient` over a `MockTransport` with a spy `sleep`.""" + return mock_client(PyPIClient, handler, **kwargs) + + +def test_get_recent_downloads_returns_raw_json() -> None: + payload = { + "data": {"last_day": 1, "last_week": 2, "last_month": 3}, + "package": "polars", + "type": "recent_downloads", + } + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/packages/polars/recent" + return httpx.Response(200, json=payload) + + client, _ = make_client(handler, min_interval=0.0) + assert client.get_recent_downloads("polars") == payload + + +def test_get_overall_sets_mirrors_param() -> None: + captured: dict[str, str | None] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["path"] = request.url.path + captured["mirrors"] = request.url.params.get("mirrors") + return httpx.Response(200, json={"data": []}) + + client, _ = make_client(handler, min_interval=0.0) + client.get_overall_downloads("polars", mirrors=True) + + assert captured["path"] == "/api/packages/polars/overall" + assert captured["mirrors"] == "true" + + +def test_courtesy_delay_applied() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": {}}) + + client, sleeps = make_client(handler, min_interval=1.0) + client.get_recent_downloads("polars") + + assert sleeps == [1.0] + + +@pytest.mark.parametrize("status", sorted(RETRYABLE_STATUSES)) +def test_retryable_status_retries_then_succeeds(status: int) -> None: + statuses = [status, 200] + calls: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + current = statuses[len(calls)] + calls.append(request) + if current == 200: + return httpx.Response(200, json={"data": {}}) + return httpx.Response(current) + + client, sleeps = make_client(handler, min_interval=0.0) + result = client.get_recent_downloads("polars") + + assert result == {"data": {}} + assert len(calls) == 2 + assert len(sleeps) == 2 # one retry backoff + one courtesy delay + + +def test_429_honors_retry_after() -> None: + statuses = [429, 200] + calls: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + status = statuses[len(calls)] + calls.append(request) + if status == 429: + return httpx.Response(429, headers={"Retry-After": "2"}) + return httpx.Response(200, json={"data": {}}) + + client, sleeps = make_client(handler, min_interval=0.0) + result = client.get_recent_downloads("polars") + + assert result == {"data": {}} + assert sleeps == [2.0, 0.0] # Retry-After honored, then courtesy delay + + +def test_context_manager_closes_client() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": {}}) + + client, _ = make_client(handler, min_interval=0.0) + with client as c: + assert c.get_recent_downloads("polars") == {"data": {}} + assert client._client.is_closed diff --git a/tests/unit/test_watermarks.py b/tests/unit/test_watermarks.py new file mode 100644 index 0000000..ea56cbf --- /dev/null +++ b/tests/unit/test_watermarks.py @@ -0,0 +1,33 @@ +"""Tests for repolytics.ingestion.watermarks.""" + +from pathlib import Path + +from repolytics.ingestion.watermarks import ( + load_watermarks, + save_watermarks, + watermark_key, +) + + +def test_watermark_key_format() -> None: + key = watermark_key("github", "commits", "encode/httpx") + assert key == "github/commits/encode/httpx" + + +def test_load_missing_returns_empty(tmp_path: Path) -> None: + assert load_watermarks(tmp_path / "nope.json") == {} + + +def test_save_then_load_roundtrip(tmp_path: Path) -> None: + path = tmp_path / "wm" / ".watermarks.json" + marks = {"github/commits/encode/httpx": "2024-01-01T00:00:00Z"} + save_watermarks(marks, path) + + assert path.exists() + assert load_watermarks(path) == marks + + +def test_save_creates_parent_dirs(tmp_path: Path) -> None: + path = tmp_path / "deep" / "nested" / ".watermarks.json" + save_watermarks({"a": "b"}, path) + assert path.exists() diff --git a/tests/unit/test_writer.py b/tests/unit/test_writer.py new file mode 100644 index 0000000..939491f --- /dev/null +++ b/tests/unit/test_writer.py @@ -0,0 +1,58 @@ +"""Tests for repolytics.ingestion.writer.""" + +import json +from datetime import datetime +from pathlib import Path + +import duckdb +import polars as pl + +from repolytics.ingestion.writer import partition_path, write_parquet + + +def test_write_parquet_roundtrips_via_duckdb( + tmp_path: Path, duckdb_conn: duckdb.DuckDBPyConnection +) -> None: + records = [ + {"id": 1, "owner": {"login": "a"}, "topics": ["x", "y"]}, + {"id": 2, "owner": {"login": "b"}, "license": None}, + ] + out = write_parquet(records, tmp_path / "sub" / "repos.parquet") + + assert out.exists() + rows = duckdb_conn.execute( + "SELECT data, _loaded_at FROM read_parquet(?) ORDER BY data", [str(out)] + ).fetchall() + + assert len(rows) == 2 + loaded = [json.loads(data) for data, _ in rows] + assert {record["id"] for record in loaded} == {1, 2} + assert loaded[0]["owner"]["login"] == "a" # nested structure preserved + assert all(loaded_at is not None for _, loaded_at in rows) + + +def test_write_parquet_columns(tmp_path: Path) -> None: + out = write_parquet([{"id": 1}], tmp_path / "x.parquet") + frame = pl.read_parquet(out) + assert frame.columns == ["data", "_loaded_at"] + + +def test_write_parquet_empty_writes_schema_only(tmp_path: Path) -> None: + out = write_parquet([], tmp_path / "empty.parquet") + frame = pl.read_parquet(out) + assert frame.columns == ["data", "_loaded_at"] + assert frame.height == 0 + + +def test_write_parquet_creates_parent_dirs(tmp_path: Path) -> None: + out = write_parquet([{"id": 1}], tmp_path / "a" / "b" / "c.parquet") + assert out.exists() + + +def test_partition_path_formats_date() -> None: + on_date = datetime(2024, 1, 2) + from_datetime = partition_path("data/raw", "github", "commits", on_date) + assert from_datetime == Path("data/raw/github/commits/2024-01-02.parquet") + + from_string = partition_path("data/raw", "github", "commits", "2024-01-02") + assert from_string == Path("data/raw/github/commits/2024-01-02.parquet") From cae04edf0322a314bb3b61d565593528ccdac660 Mon Sep 17 00:00:00 2001 From: Neukz Date: Mon, 15 Jun 2026 11:35:00 +0200 Subject: [PATCH 2/2] test: assert _loaded_at via SQL to avoid pytz dependency on CI --- tests/unit/test_writer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_writer.py b/tests/unit/test_writer.py index 939491f..47c75a3 100644 --- a/tests/unit/test_writer.py +++ b/tests/unit/test_writer.py @@ -21,14 +21,16 @@ def test_write_parquet_roundtrips_via_duckdb( assert out.exists() rows = duckdb_conn.execute( - "SELECT data, _loaded_at FROM read_parquet(?) ORDER BY data", [str(out)] + "SELECT data, _loaded_at IS NOT NULL AS has_loaded_at " + "FROM read_parquet(?) ORDER BY data", + [str(out)], ).fetchall() assert len(rows) == 2 loaded = [json.loads(data) for data, _ in rows] assert {record["id"] for record in loaded} == {1, 2} assert loaded[0]["owner"]["login"] == "a" # nested structure preserved - assert all(loaded_at is not None for _, loaded_at in rows) + assert all(has_loaded_at for _, has_loaded_at in rows) def test_write_parquet_columns(tmp_path: Path) -> None: