diff --git a/gateway/config/settings/base.py b/gateway/config/settings/base.py index 87fb85ab4..edd935d60 100644 --- a/gateway/config/settings/base.py +++ b/gateway/config/settings/base.py @@ -8,6 +8,7 @@ from urllib.parse import urlparse from celery.schedules import crontab +from celery.schedules import schedule from environs import env from config.settings.logs import ColoredFormatter @@ -543,7 +544,7 @@ def _strip_endpoint_scheme(endpoint_url: str) -> str: }, "monitor-services-health": { "task": "sds_gateway.monitoring.tasks.monitor_services_health", - "schedule": crontab(minute="*/1"), + "schedule": schedule(run_every=60), "options": {"expires": 50}, }, } diff --git a/gateway/docs/changelog.md b/gateway/docs/changelog.md index 4a5336712..fb732c52e 100644 --- a/gateway/docs/changelog.md +++ b/gateway/docs/changelog.md @@ -74,13 +74,13 @@ - Fixes: - [**Added M2M relationships for assets to - datasets**](https://github.com/spectrumx/sds-code/pull/226): Currently the - foreign key (FK) relationships only allow assets to be connected to a single - dataset. This PR expands the schema (with future contraction planned) to correct - the schema issues. + datasets**](https://github.com/spectrumx/sds-code/pull/226): Currently the + foreign key (FK) relationships only allow assets to be connected to a single + dataset. This PR expands the schema (with future contraction planned) to correct + the schema issues. - Requires a data migration for existing datasets (and captures) from the FK - field to M2M field. Run management command `migrate_fk_to_m2m` to transfer - foreign key model relationship to many-to-many field (only adds, does NOT - nullify FK). + field to M2M field. Run management command `migrate_fk_to_m2m` to transfer + foreign key model relationship to many-to-many field (only adds, does NOT + nullify FK). - This is meant to be run on systems with data under the current schema that - require migrations to new schema. + require migrations to new schema. diff --git a/gateway/docs/dev-notes.md b/gateway/docs/dev-notes.md index 2e4520ecb..91823e8d1 100644 --- a/gateway/docs/dev-notes.md +++ b/gateway/docs/dev-notes.md @@ -45,22 +45,22 @@ migrations, or before upgrading Postgres versions. Up to one snapshot per day of: + [x] **PostgreSQL database** - Complete database dump via `pg_dumpall` (all schemas, - tables, and data) + tables, and data) + [x] **OpenSearch indices** - Snapshots of all indices and their data + [x] **Gateway secrets** - Environment files (`.envs/`, `.env.sh`, etc.) and - configuration files + configuration files + [x] **Git repository state** - Branch info, commit SHA, uncommitted changes, staged - changes, and commit archive. + changes, and commit archive. ### What is NOT backed up + ❌ **MinIO objects** - S3-compatible storage data is not included in snapshots (see - warning below). When configuring MinIO, you can configure erasure coding for data - redundancy across multiple drives or nodes. Consider using MinIO's own backup and - replication features for protecting this data instead of this snapshot procedure. + warning below). When configuring MinIO, you can configure erasure coding for data + redundancy across multiple drives or nodes. Consider using MinIO's own backup and + replication features for protecting this data instead of this snapshot procedure. + ❌ **Docker volumes** - Any data stored in Docker volumes outside of the database - container is not included, such as temporary zip files, `uv` cache, virtual - environments, generated static files. + container is not included, such as temporary zip files, `uv` cache, virtual + environments, generated static files. ### Creating Backups @@ -69,17 +69,17 @@ Up to one snapshot per day of: 3. Run `just snapshot` to create a _daily_ backup snapshot of Postgres database. + See `scripts/create-snapshot.sh` for details of what is executed. + Database credentials are automatically loaded from the postgres container, so the - user running the snapshot must have privileges to manage the database container, - and the container must be running with those env variables set. + user running the snapshot must have privileges to manage the database container, + and the container must be running with those env variables set. + The Postgres snapshot is created with `pg_dumpall`. See the docs for details: - [Postgres docs](https://postgresql.org/docs/current/app-pg-dumpall.html). + [Postgres docs](https://postgresql.org/docs/current/app-pg-dumpall.html). + The snapshot creation is non-interactive, so it can be run from automated scripts; but interaction might be required for backup exfiltration. + **Note:** MinIO objects are not included in this snapshot. See [What is NOT backed up](#what-is-not-backed-up) above. 4. Access the created files in `./data/backups/`. + Note you can set up a remote server in `.env.sh` to automatically copy the backups - there after creation using `rsync` over `ssh`. + there after creation using `rsync` over `ssh`. ### Backups Restoration @@ -95,7 +95,7 @@ READ THIS ENTIRE SECTION BEFORE RESTORING A BACKUP. What can be automatically restored: + [x] **PostgreSQL database** - Full restoration of the database from the latest - snapshot. See [Backups Restoration](#backups-restoration) below. + snapshot. See [Backups Restoration](#backups-restoration) below. Everything else (OpenSearch indices, secrets, git state) can be manually restored from the snapshot files, but there is no automated restoration procedure for them yet. @@ -104,11 +104,11 @@ Note that, while the database can be restored, it will contain references to oth sources that might exist or not. Most notably: + OpenSearch indices - The restored database might contain references to documents in - OpenSearch that do not exist if those indices were not manually restored. + OpenSearch that do not exist if those indices were not manually restored. + MinIO objects - The restored database might contain references to objects in MinIO - that do not exist if those objects were not manually restored. This is the main - challenge when restoring a production backup to a staging/QA machine, as most - objects will not exist there due to the commonly large size of object data. + that do not exist if those objects were not manually restored. This is the main + challenge when restoring a production backup to a staging/QA machine, as most + objects will not exist there due to the commonly large size of object data. Run `scripts/restore-snapshot.sh` when you are ready to restore the most recent snapshot. @@ -123,9 +123,9 @@ production environment. ## MinIO Configuration + MinIO config for SDS: [see the production deploy - instructions](./detailed-deploy.md#production-deploy). + instructions](./detailed-deploy.md#production-deploy). + [MinIO reference - document](https://github.com/minio/minio/blob/master/docs/config/README.md). + document](https://github.com/minio/minio/blob/master/docs/config/README.md). ## OpenSearch Cluster Maintenance diff --git a/gateway/sds_gateway/monitoring/tests/test_celery_schedule.py b/gateway/sds_gateway/monitoring/tests/test_celery_schedule.py index 82092a37e..8be4d3a3e 100644 --- a/gateway/sds_gateway/monitoring/tests/test_celery_schedule.py +++ b/gateway/sds_gateway/monitoring/tests/test_celery_schedule.py @@ -1,4 +1,4 @@ -from celery.schedules import crontab +from celery.schedules import schedule from django.conf import settings from django.utils import timezone @@ -12,7 +12,7 @@ def test_monitoring_task_is_in_celery_beat_schedule() -> None: schedule_config["task"] == "sds_gateway.monitoring.tasks.monitor_services_health" ) - assert isinstance(schedule_config["schedule"], crontab) + assert isinstance(schedule_config["schedule"], schedule) _is_due, next_run_seconds = schedule_config["schedule"].is_due(timezone.now()) check_period_sec: int = 60 expected_expires_sec = check_period_sec - 10 diff --git a/sdk/docs/README.md b/sdk/docs/README.md index a8de6dacf..1b9469bc5 100644 --- a/sdk/docs/README.md +++ b/sdk/docs/README.md @@ -338,3 +338,28 @@ If this is needed, SDK users have a few options: One writer (an SDK client that creates, updates, and/or deletes contents) and multiple readers are generally safe. + +## Structured Logging + +The SDK writes structured JSONL logs to `~/.local/state/spectrumx/logs/YYYY-MM-DD.jsonl`. +Each line is a JSON object with fields for timestamp, severity, category, and message. +Filter them with `jq`: + +```bash +# Show only filesystem operations +jq 'select(.cat == "filesystem")' ~/.local/state/spectrumx/logs/*.jsonl + +# Show only download operations +jq 'select(.cat == "download")' ~/.local/state/spectrumx/logs/*.jsonl + +# Show warnings and errors +jq 'select(.lvl == "WARNING" or .lvl == "ERROR")' ~/.local/state/spectrumx/logs/*.jsonl + +# Compact view per line +jq -r '"\(.ts[11:19]) [\(.lvl)] [\(.cat)] \(.msg)"' ~/.local/state/spectrumx/logs/*.jsonl + +# Count messages per category +jq -r '.cat' ~/.local/state/spectrumx/logs/*.jsonl | sort | uniq -c | sort -rn +``` + +See the [full structured logging docs](https://sds.crc.nd.edu/sdk/advanced/structured-logging/) for more examples. diff --git a/sdk/docs/mkdocs.yaml b/sdk/docs/mkdocs.yaml index 9e869041b..de7b4ad2e 100644 --- a/sdk/docs/mkdocs.yaml +++ b/sdk/docs/mkdocs.yaml @@ -25,6 +25,7 @@ nav: - "Changelog": "changelog.md" - "Advanced": - "Concurrent Access": "advanced/concurrent-access.md" + - "Structured Logging": "advanced/structured-logging.md" - "API Reference": - "Client": "api/client.md" - "Errors": "api/errors.md" diff --git a/sdk/docs/mkdocs/advanced/structured-logging.md b/sdk/docs/mkdocs/advanced/structured-logging.md new file mode 100644 index 000000000..e8f48caeb --- /dev/null +++ b/sdk/docs/mkdocs/advanced/structured-logging.md @@ -0,0 +1,101 @@ +# Structured Logging + +The SDK writes structured JSONL logs to `~/.local/state/spectrumx/logs/YYYY-MM-DD.jsonl`. +Each line is a JSON object with fields like `ts`, `pid`, `lvl`, `cat`, and `msg`. + +## Filtering with jq + +Use [`jq`](https://jqlang.org/) to filter and analyze log files on the command line. + +### By category + +```bash +# Filesystem operations only +jq 'select(.cat == "filesystem")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Network operations +jq 'select(.cat == "network")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Upload-specific messages +jq 'select(.cat == "upload")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Download-specific messages +jq 'select(.cat == "download")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Auth messages +jq 'select(.cat == "auth")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Config messages +jq 'select(.cat == "config")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl +``` + +### By severity + +```bash +# Warnings and errors only +jq 'select(.lvl == "WARNING" or .lvl == "ERROR")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Errors only with full exception info +jq 'select(.lvl == "ERROR")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl +``` + +### Compound filters + +```bash +# Filesystem warnings and errors +jq 'select(.cat == "filesystem" and (.lvl == "WARNING" or .lvl == "ERROR"))' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Everything except generic "log" category +jq 'select(.cat != "log")' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Messages with exception details +jq 'select(.exc_info != null)' ~/.local/state/spectrumx/logs/2026-06-29.jsonl +``` + +### Aggregations + +```bash +# Count messages per category +jq -r '.cat' ~/.local/state/spectrumx/logs/2026-06-29.jsonl | sort | uniq -c | sort -rn + +# Count messages per severity level +jq -r '.lvl' ~/.local/state/spectrumx/logs/2026-06-29.jsonl | sort | uniq -c | sort -rn + +# Extract just the messages for a category +jq -r 'select(.cat == "network") | .msg' ~/.local/state/spectrumx/logs/2026-06-29.jsonl +``` + +### Pretty-print and format + +```bash +# Colorized output (jq highlights strings by default) +jq . ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Compact view: timestamp + level + category + message +jq -r '"\(.ts[11:19]) [\(.lvl)] [\(.cat)] \(.msg)"' ~/.local/state/spectrumx/logs/2026-06-29.jsonl + +# Tabular output with header +jq -r '["TS","LVL","CAT","MSG"], ["---","---","---","---"], (.[] | [.ts[11:19], .lvl, .cat, .msg]) | @tsv' ~/.local/state/spectrumx/logs/2026-06-29.jsonl | column -t -s $'\t' +``` + +### Watch in real time + +```bash +# Tail and filter: show only filesystem logs as they arrive +tail -f ~/.local/state/spectrumx/logs/2026-06-29.jsonl | jq 'select(.cat == "filesystem")' + +# Tail with compact format +tail -f ~/.local/state/spectrumx/logs/2026-06-29.jsonl | jq -r '"\(.ts[11:19]) [\(.lvl)] [\(.cat)] \(.msg)"' +``` + +## Log categories reference + +| Category | Description | Examples | +|--------------|---------------------------------|---------------------------------------| +| `log` | Generic/default messages | Dry-run notices, user-facing messages | +| `config` | SDK configuration | Config loading, env variable warnings | +| `auth` | Authentication operations | Token validation, auth success/failure| +| `network` | HTTP/network operations | Response info, connection errors | +| `filesystem` | File and capture CRUD | Upload/download paths, capture ops | +| `download` | File download operations | Download skip/reason, dry-run notices, path resolution | +| `upload` | File upload metadata operations | Upload status, metadata updates | diff --git a/sdk/justfile b/sdk/justfile index ad1c9d794..c351e811b 100644 --- a/sdk/justfile +++ b/sdk/justfile @@ -116,6 +116,7 @@ pyrefly *args: test python=python_version *pytest_args: #!/usr/bin/env bash set -euo pipefail + export COLUMNS=88 echo -e "\n\t\033[34mRunning tests against Python {{ python }}\033[0m\n" # no coverage if -k is used (likely a focused test run) @@ -145,6 +146,7 @@ test python=python_version *pytest_args: test-all *args: #!/usr/bin/env bash set -euo pipefail + COLUMNS=80 echo -e "\n\t\033[34mRunning all local (non-integration) tests\033[0m\n" # run higher bound tests against default python version @@ -169,7 +171,7 @@ test-integration python=python_version workers="auto" *pytest_args: @echo -e "\n\t\033[33mRunning integration tests in parallel\033[0m" @echo -e "\t\033[33m Workers: {{ workers }}\033[0m" @echo -e "\t\033[33m Pass -n 1 or use 'just test-integration-sequential' for sequential execution\033[0m\n" - uv run --resolution highest -p '{{ python }}' pytest \ + COLUMNS=88 uv run --resolution highest -p '{{ python }}' pytest \ -m '{{ test_marker_integration }}' \ -n {{ workers }} \ --dist loadgroup \ @@ -181,7 +183,7 @@ test-integration python=python_version workers="auto" *pytest_args: test-integration-sequential python=python_version *pytest_args: @echo -e "\n\t\033[33mRunning integration tests sequentially\033[0m" @echo -e "\t\033[33m Use 'just test-integration' for parallel execution\033[0m\n" - uv run --resolution highest -p '{{ python }}' pytest \ + COLUMNS=88 uv run --resolution highest -p '{{ python }}' pytest \ -m '{{ test_marker_integration }}' \ {{ pytest_args }} \ tests @@ -190,7 +192,7 @@ test-integration-sequential python=python_version *pytest_args: [group('qa')] test-lowest python=python_version *pytest_args: @echo "Running lowest dep resolution tests for Python '{{ python }}'" - uv run --resolution lowest-direct -p '{{ python }}' pytest \ + COLUMNS=88 uv run --resolution lowest-direct -p '{{ python }}' pytest \ -m '{{ test_marker_default }}' \ {{ pytest_args }} \ tests @@ -198,7 +200,7 @@ test-lowest python=python_version *pytest_args: # runs local tests in verbose mode with stdout capture; useful for debugging test failures [group('qa')] test-verbose python=python_version *pytest_args: - uv run --resolution highest -p '{{ python }}' pytest \ + COLUMNS=88 uv run --resolution highest -p '{{ python }}' pytest \ -vvv \ -m '{{ test_marker_default }}' \ --show-capture=stdout \ @@ -213,7 +215,7 @@ test-verbose python=python_version *pytest_args: test-integration-verbose python=python_version *pytest_args: @echo -e "\n\t\033[33mRunning integration tests: make sure you set up the web application\033[0m" @echo -e "\t\033[33mand configure 'tests/integration/integration.env'\033[0m\n" - uv run --resolution highest -p '{{ python }}' pytest \ + COLUMNS=88 uv run --resolution highest -p '{{ python }}' pytest \ -m '{{ test_marker_integration }}' \ -vvv \ --capture=no \ @@ -226,7 +228,7 @@ test-integration-quick python=python_version workers="auto" *pytest_args: @echo -e "\n\t\033[33mRunning quick integration tests in parallel (skipping heavy tests)\033[0m" @echo -e "\t\033[33m Workers: {{ workers }}\033[0m" @echo -e "\t\033[33m Marker: {{ test_marker_integration }} and not heavy\033[0m\n" - uv run --resolution highest -p '{{ python }}' pytest \ + COLUMNS=88 uv run --resolution highest -p '{{ python }}' pytest \ -m '{{ test_marker_integration }} and not heavy' \ -n {{ workers }} \ --dist loadgroup \ @@ -258,3 +260,8 @@ publish: update: uv sync --upgrade --dev uv run prek autoupdate + +# runs the progress bars demo with the given file count and rate +[group('development')] +demo-progress-bars count="9" rate="10": + uv run scripts/demo_progress_bars.py --file-count {{ count }} --rate {{ rate }} diff --git a/sdk/pyproject.toml b/sdk/pyproject.toml index 3f9ebd238..40a0c037c 100644 --- a/sdk/pyproject.toml +++ b/sdk/pyproject.toml @@ -419,6 +419,9 @@ # Tests commonly use boolean fixtures as positional args — intentional "sdk/tests/**" = ["FBT001", "FBT002"] + # Demo scripts intentionally use print(), lazy imports, private member access for mocking + "sdk/scripts/**" = ["T201", "PLC0415", "SLF001", "S108", "RUF001"] + [tool.ruff.format] indent-style = "space" line-ending = "auto" diff --git a/sdk/scripts/demo_progress_bars.py b/sdk/scripts/demo_progress_bars.py new file mode 100755 index 000000000..58967d39b --- /dev/null +++ b/sdk/scripts/demo_progress_bars.py @@ -0,0 +1,408 @@ +#!/usr/bin/env python3 +"""Demo the real byte-level progress bars from the SDK. + +Creates test files, mocks the gateway HTTP layer to simulate transfers +at a given rate, and runs the real download/upload code paths so you +can see the tqdm bars animate smoothly. +""" + +from __future__ import annotations + +import atexit +import json +import os +import shutil +import time +import uuid +from datetime import UTC +from datetime import datetime +from datetime import timedelta +from pathlib import Path +from textwrap import dedent + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +FILE_COUNT = 4 +FILE_SIZE = 10_000_000 # 10 MB per file +RATE_MBPS = 1 # 1 MB/s default +MOCK_CHUNK = 256 # tiny internal chunks -> smooth tqdm animation +_MIN_SLEEP_THRESHOLD = 0.001 # 1 ms: minimum sleep the OS scheduler can deliver +_BYTE_KB = 1024 # bytes per kilobyte + +# Empirical correction: the time-budget accumulator only accounts for +# sleep duration, not the Python overhead per chunk (byte generation, +# yield, loop, monotonic calls). In practice the actual throughput ends +# up ~5% higher than the nominal rate, so we discount the rate used for +# timing calculations by this factor to match the target. +_RATE_OVERHEAD_FUDGE = 0.954 + +# Set dynamically by main() via globals() update; declare here for linting. +rate_bps: float = RATE_MBPS * _BYTE_KB * _BYTE_KB + + +def _human_bytes(n: int) -> str: + value = float(n) + for unit in ("B", "KB", "MB", "GB"): + if abs(value) < _BYTE_KB: + return f"{value:.1f} {unit}" + value /= _BYTE_KB + return f"{value:.1f} PB" + + +def _human_rate(bps: float) -> str: + return _human_bytes(int(bps)).replace("B", "B/s") + + +# --------------------------------------------------------------------------- +# Mock HTTP layer +# --------------------------------------------------------------------------- + + +class _MockResponse: + """A requests.Response look-alike that the gateway code accepts.""" + + status_code = 200 + ok = True + reason = "OK" + + def __init__(self, content: bytes, file_size: int | None = None): + self._content = content + self._file_size = file_size + self.headers = {} + if file_size is not None: + self.headers["Content-Length"] = str(file_size) + + @property + def content(self): + return self._content + + def json(self): + return json.loads(self._content) + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def iter_content(self, chunk_size: int = 8192): + """Yield tiny chunks at the rate, for smooth tqdm updates. + + Instead of calling ``time.sleep`` per chunk (which is + unreliable for sub-millisecond intervals), we compare cumulative + expected time against wall clock and only sleep when the deficit + exceeds 1 ms. This aggregates many tiny delays into fewer but + longer sleeps that the OS scheduler can deliver accurately. + """ + if self._file_size is not None: + inner = min(chunk_size, MOCK_CHUNK) + offset = 0 + _t0 = time.monotonic() + _expected = 0.0 + while offset < self._file_size: + n = min(inner, self._file_size - offset) + _expected += n / (rate_bps * _RATE_OVERHEAD_FUDGE) + _deficit = _expected - (time.monotonic() - _t0) + if _deficit > _MIN_SLEEP_THRESHOLD: + time.sleep(_deficit) + yield b"\x00" * n + offset += n + elif self._content: + yield self._content + + +def make_mock_response(content: bytes, *, file_size: int | None = None): + """Build a requests.Response look-alike that the gateway code accepts.""" + return _MockResponse(content, file_size=file_size) + + +def make_file_json(i: int, file_uuid: str, size: int, rel: str) -> bytes: + """Return serialised File JSON that the gateway returns for this file.""" + now = datetime.now(tz=UTC).isoformat() + later = (datetime.now(tz=UTC) + timedelta(days=30)).isoformat() + info = { + "uuid": file_uuid, + "name": f"demo_file_{i:02d}.bin", + "media_type": "application/octet-stream", + "size": size, + "directory": rel, + "permissions": "rw-rw-r--", + "created_at": now, + "updated_at": now, + "expiration_date": later, + "sum_blake3": None, + } + return json.dumps(info).encode() + + +# --------------------------------------------------------------------------- +# Demos +# --------------------------------------------------------------------------- + + +def demo_downloads(tmp: Path) -> None: + """Run a full download flow using the real Client code paths. + + The gateway._request is patched so every HTTP call returns a local + mock response. All progress bars are the real ones from gateway.py + and client.py. + """ + from spectrumx.client import Client + from spectrumx.gateway import Endpoints + from spectrumx.gateway import HTTPMethods + from spectrumx.models.files import File + + os.environ["SPECTRUMX_API_KEY"] = "demo_key" + + client = Client(host="localhost", verbose=True) + client.dry_run = False + gateway = client._gateway + + uuids_hex = [uuid.uuid4().hex for _ in range(FILE_COUNT)] + file_json_map: dict[str, bytes] = {} + demo_files: list[File] = [] + for i, uid_hex in enumerate(uuids_hex): + j = make_file_json(i, uid_hex, FILE_SIZE, "/demo") + file_json_map[uid_hex] = j + f = File.model_validate_json(j) + f.size = FILE_SIZE + f.local_path = tmp / f.name + demo_files.append(f) + + total_bytes_h = _human_bytes(FILE_COUNT * FILE_SIZE) + rate_h = _human_rate(rate_bps) + print(f" ↓ Downloading {FILE_COUNT} files ({total_bytes_h} at {rate_h})…\n") + + orig_request = gateway._request + + def _mock_request( + method=None, + endpoint=None, + *, + asset_id=None, + endpoint_args=None, + stream=False, + timeout=None, + verbose=False, + **kw, + ): + nonlocal orig_request, file_json_map + + if endpoint == Endpoints.FILES and method == HTTPMethods.GET and asset_id: + uid = asset_id + if uid in file_json_map: + return make_mock_response(file_json_map[uid]) + + elif endpoint == Endpoints.FILE_DOWNLOAD: + uid = (endpoint_args or {}).get("uuid") or asset_id + if uid in file_json_map: + j = json.loads(file_json_map[uid]) + sz = j.get("size", FILE_SIZE) + return make_mock_response(b"", file_size=sz) + + return make_mock_response(b'{"detail":"mocked ok"}') + + try: + gateway._request = _mock_request + results = client.download( + files_to_download=demo_files, + to_local_path=tmp / "downloads", + verbose=True, + overwrite=True, + ) + ok = sum(1 for r in results if r) + fail = len(results) - ok + print(f" ↓ Done: {ok} ok, {fail} failed\n") + finally: + gateway._request = orig_request + + +def demo_uploads(tmp: Path) -> None: + """Run a real upload flow with the gateway patched to simulate slow I/O. + + The _ProgressFile wrapper in upload_new_file tracks bytes read from the + source file, which updates the tqdm progress bar. We patch _request + so that the FILES endpoint reads from the _ProgressFile (triggering the + bar) at the configured rate before returning a success response. + """ + from spectrumx.client import Client + from spectrumx.gateway import Endpoints + + os.environ["SPECTRUMX_API_KEY"] = "demo_key" + os.environ.pop("PYTEST_CURRENT_TEST", None) + + client = Client(host="localhost", verbose=True) + client.dry_run = False + gateway = client._gateway + + upload_root = tmp / "to_upload" + upload_root.mkdir(parents=True, exist_ok=True) + for i in range(FILE_COUNT): + p = upload_root / f"src_file_{i:02d}.bin" + sz_left = FILE_SIZE + with p.open("wb") as f: + while sz_left: + n = min(65536, sz_left) + f.write(os.urandom(n)) + sz_left -= n + + total_h = _human_bytes(FILE_COUNT * FILE_SIZE) + print(f" ↑ Uploading {FILE_COUNT} files ({total_h} at {_human_rate(rate_bps)})…\n") + + orig_request = gateway._request + + def _mock_request( + method=None, + endpoint=None, + *, + asset_id=None, + endpoint_args=None, + stream=False, + timeout=None, + verbose=False, + **kw, + ): + nonlocal orig_request + + if endpoint == Endpoints.FILE_CONTENTS_CHECK: + return make_mock_response( + json.dumps( + { + "file_exists_in_tree": False, + "file_contents_exist_for_user": False, + "user_mutable_attributes_differ": False, + "asset_id": None, + } + ).encode() + ) + if endpoint == Endpoints.FILES and kw.get("files"): + files = kw["files"] + fp = files.get("file") + if fp is not None: + _t0 = time.monotonic() + _expected = 0.0 + while True: + chunk = fp.read(MOCK_CHUNK) + if not chunk: + break + # discounted by _RATE_OVERHEAD_FUDGE (see Config) + _expected += len(chunk) / (rate_bps * _RATE_OVERHEAD_FUDGE) + _deficit = _expected - (time.monotonic() - _t0) + if _deficit > _MIN_SLEEP_THRESHOLD: + time.sleep(_deficit) + return make_mock_response( + make_file_json( + i=0, + file_uuid=str(uuid.uuid4()), + size=FILE_SIZE, + rel="/demo_uploads", + ) + ) + return make_mock_response(b'{"detail":"mocked"}') + + try: + gateway._request = _mock_request + results = client.upload( + local_path=upload_root, + sds_path="/demo_uploads", + verbose=True, + persist_state=False, + ) + ok = sum(1 for r in results if r) + fail = len(results) - ok + print(f" ↑ Done: {ok} ok, {fail} failed\n") + finally: + gateway._request = orig_request + + +# --------------------------------------------------------------------------- +# Cleanup +# --------------------------------------------------------------------------- + +TMP_DIR = Path("/tmp") / f"sdk_demo_progress_{uuid.uuid4().hex[:8]}" + + +def _cleanup() -> None: + if TMP_DIR.exists(): + shutil.rmtree(TMP_DIR, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + import argparse + + parser = argparse.ArgumentParser( + description=dedent("""\ + Demo real SDK progress bars for byte-level upload/download tracking. + + Creates test files and mocks the gateway HTTP layer so the SDK's + own tqdm progress bars (in gateway.py, client.py, uploads.py) + render with realistic timing — no real server needed. + """) + ) + parser.add_argument( + "--demo", + choices=["download", "upload", "all"], + default="all", + ) + parser.add_argument( + "--rate", + type=float, + default=RATE_MBPS, + help=f"Transfer rate in MB/s (default {RATE_MBPS})", + ) + parser.add_argument( + "--file-count", + type=int, + default=FILE_COUNT, + help=f"Number of test files (default {FILE_COUNT})", + ) + args = parser.parse_args() + + globals().update( + rate_bps=args.rate * 1024 * 1024, + FILE_COUNT=args.file_count, + ) + + # Clean slate (leftover files from crashed runs confuse the upload scanner) + _cleanup() + atexit.register(_cleanup) + TMP_DIR.mkdir(parents=True, exist_ok=True) + + size_h = _human_bytes(FILE_SIZE) + rate_h = _human_rate(rate_bps) + print( + f"\nSDK progress bars demo | {args.file_count} files x {size_h} | {rate_h}\n" + ) + + if args.demo in ("download", "all"): + print("═" * 56) + print(" DOWNLOAD — outer progress bar (client.py)") + print("═" * 56) + print(" Shows file count with byte postfix per completed file") + print() + demo_downloads(TMP_DIR) + + if args.demo in ("upload", "all"): + print("\n") + print("═" * 56) + print(" UPLOAD — workload progress bar (uploads.py)") + print("═" * 56) + print(" Shows per-file completion with total bytes transferred") + print() + demo_uploads(TMP_DIR) + + print("─" * 56) + print(" Demo complete. See the SDK code at:\n") + print(" sdk/src/spectrumx/gateway.py") + print(" sdk/src/spectrumx/client.py") + print(" sdk/src/spectrumx/api/uploads.py\n") + + +if __name__ == "__main__": + main() diff --git a/sdk/src/spectrumx/__init__.py b/sdk/src/spectrumx/__init__.py index 8c5cb173c..9367adb73 100644 --- a/sdk/src/spectrumx/__init__.py +++ b/sdk/src/spectrumx/__init__.py @@ -26,7 +26,13 @@ try: from loguru import logger as log - log.disable(LIB_NAME) + # Remove the default stderr handler that Loguru adds on first import. + # Add a no-op null handler to prevent Loguru from re-adding the default + # handler on the first log call. Individual features + # (enable_logging, enable_structured_logging) add their own sinks + # as needed. + log.remove() + log.add(lambda msg: None, level="TRACE") except ImportError: pass @@ -69,8 +75,9 @@ def enable_logging() -> None: } ] ) - log.enable(LIB_NAME) # pyright: ignore[reportPossiblyUnboundVariable] - log.info(f"Enabled logging for '{LIB_NAME}'") # pyright: ignore[reportPossiblyUnboundVariable] + from spectrumx.utils import LogCategory # noqa: PLC0415 + + log.bind(cat=LogCategory.CONFIG).info(f"Enabled logging for '{LIB_NAME}'") # pyright: ignore[reportPossiblyUnboundVariable] except NameError: import logging # noqa: PLC0415 @@ -79,6 +86,11 @@ def enable_logging() -> None: logger.warning("Install Loguru to enable additional spectrumx logging.") +# ------------------- +# structured logging + +from spectrumx.utils import enable_structured_logging + # ------------------- # package exports @@ -86,6 +98,7 @@ def enable_logging() -> None: "Client", "__version__", "enable_logging", + "enable_structured_logging", "experiments", "models", ] diff --git a/sdk/src/spectrumx/api/captures.py b/sdk/src/spectrumx/api/captures.py index cb2ba6722..5a9367c8b 100644 --- a/sdk/src/spectrumx/api/captures.py +++ b/sdk/src/spectrumx/api/captures.py @@ -19,6 +19,7 @@ from spectrumx.models.captures import CaptureOrigin from spectrumx.models.captures import CaptureType from spectrumx.models.user import User +from spectrumx.utils import LogCategory from spectrumx.utils import log_user_warning if TYPE_CHECKING: @@ -89,23 +90,27 @@ def create( or the user doesn't have permission to create it. """ if index_name: - log.warning( + log.bind(cat=LogCategory.LOG).warning( "The 'index_name' parameter is deprecated and " "will be removed in future versions." ) index_name = index_mapping.get(capture_type, index_name) if not index_name: - log.warning(f"Could not find an index for {capture_type=}") + log.bind(cat=LogCategory.LOG).warning( + f"Could not find an index for {capture_type=}" + ) top_level_dir = PurePosixPath(top_level_dir) if self.verbose: - log.debug( + log.bind(cat=LogCategory.FILESYSTEM).debug( f"Creating capture with {top_level_dir=}, " f"{channel=}, {capture_type=}, {index_name=}, {scan_group=}, {name=}" ) if self.dry_run: - log.debug("Dry run enabled: simulating the capture creation") + log.bind(cat=LogCategory.FILESYSTEM).debug( + "Dry run enabled: simulating the capture creation" + ) return Capture( capture_props={}, capture_type=capture_type, @@ -138,7 +143,9 @@ def create( ) capture = Capture.model_validate_json(capture_raw) if self.verbose: - log.debug(f"Capture created with UUID {capture.uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Capture created with UUID {capture.uuid}" + ) return capture def listing(self, *, capture_type: CaptureType | None = None) -> list[Capture]: @@ -153,9 +160,13 @@ def listing(self, *, capture_type: CaptureType | None = None) -> list[Capture]: A list of the RF captures found owned by the requesting user. """ if self.verbose: - log.debug(f"Listing captures of type {capture_type}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Listing captures of type {capture_type}" + ) if self.dry_run: - log.debug("Dry run enabled: simulating capture listing") + log.bind(cat=LogCategory.FILESYSTEM).debug( + "Dry run enabled: simulating capture listing" + ) num_captures: int = 3 rng = random.Random() # noqa: S311 return [ @@ -167,7 +178,9 @@ def listing(self, *, capture_type: CaptureType | None = None) -> list[Capture]: captures_raw = self.gateway.list_captures(capture_type=capture_type) captures_list_raw, has_more = _extract_page_from_payload(captures_raw) if has_more: - log.warning("Not all capture results may be listed. ") + log.bind(cat=LogCategory.FILESYSTEM).warning( + "Not all capture results may be listed. " + ) # TODO: request more pages if needed captures: list[Capture] = [] for captures_raw in captures_list_raw: @@ -176,10 +189,12 @@ def listing(self, *, capture_type: CaptureType | None = None) -> list[Capture]: captures.append(capture) except ValidationError as err: log_user_warning(f"Validation error loading capture: {captures_raw}") - log.exception(err) + log.bind(cat=LogCategory.FILESYSTEM).exception(err) continue if self.verbose: - log.debug(f"Listing {len(captures)} captures") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Listing {len(captures)} captures" + ) return captures def update( @@ -196,10 +211,14 @@ def update( None """ if self.verbose: - log.debug(f"Updating capture with UUID {capture_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Updating capture with UUID {capture_uuid}" + ) if self.dry_run: - log.debug(f"Dry run enabled: simulating capture update {capture_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Dry run enabled: simulating capture update {capture_uuid}" + ) return capture_raw = self.gateway.update_capture( @@ -208,7 +227,9 @@ def update( ) capture = Capture.model_validate_json(capture_raw) if self.verbose: - log.debug(f"Capture updated with UUID {capture.uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Capture updated with UUID {capture.uuid}" + ) def read( self, @@ -226,12 +247,16 @@ def read( The capture object. """ if self.verbose: - log.debug(f"Reading capture with UUID {capture_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Reading capture with UUID {capture_uuid}" + ) capture_raw = self.gateway.read_capture(capture_uuid=capture_uuid) capture = Capture.model_validate_json(capture_raw) if self.verbose: - log.debug(f"Capture read with UUID {capture.uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Capture read with UUID {capture.uuid}" + ) return capture def delete( @@ -249,17 +274,23 @@ def delete( if it has already been deleted; if this user doesn't own it. """ if self.verbose: - log.debug(f"Deleting capture with UUID {capture_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Deleting capture with UUID {capture_uuid}" + ) if self.dry_run: - log.debug(f"Dry run enabled: would delete capture {capture_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Dry run enabled: would delete capture {capture_uuid}" + ) return True self.gateway.delete_capture( capture_uuid=capture_uuid, ) if self.verbose: - log.debug(f"Capture deleted with UUID {capture_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Capture deleted with UUID {capture_uuid}" + ) return True def revoke_share_permissions(self, capture_uuid: uuid.UUID) -> bool: @@ -268,9 +299,13 @@ def revoke_share_permissions(self, capture_uuid: uuid.UUID) -> bool: Use this (or the web portal) before :meth:`delete` when the capture is shared. """ if self.verbose: - log.debug(f"Revoking share permissions for capture {capture_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Revoking share permissions for capture {capture_uuid}" + ) if self.dry_run: - log.debug("Dry run enabled: would revoke capture share permissions") + log.bind(cat=LogCategory.FILESYSTEM).debug( + "Dry run enabled: would revoke capture share permissions" + ) return True self.gateway.revoke_capture_share_permissions(capture_uuid=capture_uuid) return True @@ -281,9 +316,13 @@ def detach_from_datasets(self, capture_uuid: uuid.UUID) -> bool: Use before :meth:`delete` when the capture is still linked to datasets. """ if self.verbose: - log.debug(f"Detaching capture {capture_uuid} from datasets") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Detaching capture {capture_uuid} from datasets" + ) if self.dry_run: - log.debug("Dry run enabled: would detach capture from datasets") + log.bind(cat=LogCategory.FILESYSTEM).debug( + "Dry run enabled: would detach capture from datasets" + ) return True self.gateway.detach_capture_from_datasets(capture_uuid=capture_uuid) return True @@ -322,7 +361,9 @@ def advanced_search( # TODO: adapt this function to return a Paginator[Capture] object if self.dry_run: - log.debug("Dry run enabled: simulating search results") + log.bind(cat=LogCategory.FILESYSTEM).debug( + "Dry run enabled: simulating search results" + ) rng = random.Random() # noqa: S311 return [ _generate_capture(capture_type=rng.choice(list(CaptureType))) @@ -347,12 +388,14 @@ def advanced_search( log_user_warning( f"Validation error loading search result: {result_raw}" ) - log.exception(err) + log.bind(cat=LogCategory.FILESYSTEM).exception(err) continue else: captures.append(capture) if self.verbose: - log.debug(f"Search returned {len(captures)} captures") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Search returned {len(captures)} captures" + ) return captures diff --git a/sdk/src/spectrumx/api/datasets.py b/sdk/src/spectrumx/api/datasets.py index 9ae8f7b36..8cf0ccd33 100644 --- a/sdk/src/spectrumx/api/datasets.py +++ b/sdk/src/spectrumx/api/datasets.py @@ -11,6 +11,7 @@ from spectrumx.models.datasets import Dataset from spectrumx.models.files import File from spectrumx.ops.pagination import Paginator +from spectrumx.utils import LogCategory from spectrumx.utils import log_user if TYPE_CHECKING: @@ -118,11 +119,15 @@ def get_files( log_user("Dry run enabled: files will be simulated") if artifacts_only and capture_uuids: - log.warning("Capture UUIDs are not allowed when artifacts_only is True.") + log.bind(cat=LogCategory.FILESYSTEM).warning( + "Capture UUIDs are not allowed when artifacts_only is True." + ) capture_uuids = None if artifacts_only and top_level_dirs: - log.info("Top level directories included will ONLY return artifact files.") + log.bind(cat=LogCategory.FILESYSTEM).info( + "Top level directories included will ONLY return artifact files." + ) list_kwargs: dict[str, Any] = {"dataset_uuid": dataset_uuid} if capture_uuids is not None: @@ -155,17 +160,23 @@ def delete( True if the dataset was deleted successfully, or if in dry run mode. """ if self.verbose: - log.debug(f"Deleting dataset with UUID {dataset_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Deleting dataset with UUID {dataset_uuid}" + ) if self.dry_run: - log.debug(f"Dry run enabled: would delete dataset {dataset_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Dry run enabled: would delete dataset {dataset_uuid}" + ) return True self.gateway.delete_dataset( dataset_uuid=dataset_uuid, ) if self.verbose: - log.debug(f"Dataset deleted with UUID {dataset_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Dataset deleted with UUID {dataset_uuid}" + ) return True def revoke_share_permissions(self, dataset_uuid: UUID) -> bool: @@ -174,9 +185,13 @@ def revoke_share_permissions(self, dataset_uuid: UUID) -> bool: Use this (or the web portal) before :meth:`delete` when the dataset is shared. """ if self.verbose: - log.debug(f"Revoking share permissions for dataset {dataset_uuid}") + log.bind(cat=LogCategory.FILESYSTEM).debug( + f"Revoking share permissions for dataset {dataset_uuid}" + ) if self.dry_run: - log.debug("Dry run enabled: would revoke dataset share permissions") + log.bind(cat=LogCategory.FILESYSTEM).debug( + "Dry run enabled: would revoke dataset share permissions" + ) return True self.gateway.revoke_dataset_share_permissions(dataset_uuid=dataset_uuid) return True diff --git a/sdk/src/spectrumx/api/sds_files.py b/sdk/src/spectrumx/api/sds_files.py index eaee3c625..d0fd40dc7 100644 --- a/sdk/src/spectrumx/api/sds_files.py +++ b/sdk/src/spectrumx/api/sds_files.py @@ -3,6 +3,7 @@ # ruff: noqa: SLF001 # pyright: reportPrivateUsage=false +import collections.abc import os import tempfile import uuid @@ -23,6 +24,7 @@ from spectrumx.models.files import File from spectrumx.ops import files from spectrumx.ops.pagination import Paginator +from spectrumx.utils import LogCategory from spectrumx.utils import log_user from spectrumx.utils import log_user_warning @@ -78,6 +80,7 @@ def download_file( to_local_path: Path | str | None = None, skip_contents: bool = False, warn_missing_path: bool = True, + progress_callback: collections.abc.Callable[[int], None] | None = None, ) -> File: """Downloads a file from SDS: metadata and contents (unless skip_contents=True). @@ -109,6 +112,7 @@ def download_file( valid_uuid=valid_uuid, valid_local_path_or_none=valid_local_path_or_none, skip_contents=skip_contents, + progress_callback=progress_callback, ) return file_instance @@ -162,6 +166,7 @@ def upload_file( client: Client, local_file: File | Path | str, sds_path: PurePosixPath | Path | str = "/", + progress_callback: collections.abc.Callable[[int], None] | None = None, ) -> File: """Uploads a file to SDS. @@ -196,7 +201,9 @@ def upload_file( file_instance = files.construct_file(local_file, sds_path=sds_path) del local_file - return __upload_file_mux(client=client, file_instance=file_instance) + return __upload_file_mux( + client=client, file_instance=file_instance, progress_callback=progress_callback + ) def detach_file_from_datasets( @@ -287,19 +294,18 @@ def __download_file_contents_if_applicable( valid_uuid: UUID4, valid_local_path_or_none: Path | None, skip_contents: bool = False, + progress_callback: collections.abc.Callable[[int], None] | None = None, ) -> None: if skip_contents: - msg = ( - "A file instance was provided and skip_contents " - "is True: nothing to download." + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=1).warning( + "Download skipped: skip_contents=True, no file contents to download" ) - log_user_warning(msg) return if client.dry_run: file_instance.local_path = valid_local_path_or_none - log_user( - "Dry run enabled: file contents would be " - f"downloaded as {file_instance.local_path}" + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + f"Download dry-run: file contents would be saved to " + f"{file_instance.local_path}" ) return downloaded_path: Path = __download_file_contents( @@ -307,6 +313,7 @@ def __download_file_contents_if_applicable( file_uuid=valid_uuid, target_path=valid_local_path_or_none, contents_lock=file_instance.contents_lock, # pyright: ignore[reportPrivateUsage] + progress_callback=progress_callback, ) file_instance.local_path = downloaded_path @@ -317,6 +324,7 @@ def __download_file_contents( file_uuid: UUID4 | str, contents_lock: RLock, target_path: Path | None = None, + progress_callback: collections.abc.Callable[[int], None] | None = None, ) -> Path: """Downloads the contents of a file from SDS to a location on disk. @@ -327,8 +335,10 @@ def __download_file_contents( to ensure partial downloads don't leave incomplete files in place. Args: - file_uuid: The UUID of the file to download from SDS. - target_path: The local path to save the downloaded file to. + file_uuid: The UUID of the file to download from SDS. + target_path: The local path to save the downloaded file to. + progress_callback: Optional callable invoked with byte count as + chunks are received from the stream (for progress tracking). Returns: The local path to the downloaded file. """ @@ -345,7 +355,9 @@ def __download_file_contents( ) if client.dry_run: - log_user(f"Dry run enabled: file would be saved as {target_path}") + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + f"Download dry-run: file contents would be saved to {target_path}" + ) return target_path target_path.parent.mkdir(parents=True, exist_ok=True) @@ -370,6 +382,8 @@ def __download_file_contents( with download_path.open(mode="wb") as file_ptr, contents_lock: for chunk in client._gateway.get_file_contents_by_id(uuid=uuid_to_set.hex): file_ptr.write(chunk) + if progress_callback: + progress_callback(len(chunk)) file_ptr.flush() os.fsync(file_ptr.fileno()) @@ -397,11 +411,9 @@ def __extract_download_info_from_file_instance( if to_local_path: file_instance.local_path = Path(to_local_path) if file_instance.local_path is None and warn_missing_path: - msg = ( - "The file instance passed is missing a local path to " - "download to. A temporary file will be created on disk." + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=1).warning( + "Download: no local path set on file instance; using temporary file" ) - log_user_warning(msg) valid_uuid = file_instance.uuid valid_local_path_or_none = file_instance.local_path return file_instance, valid_uuid, valid_local_path_or_none @@ -465,8 +477,9 @@ def __pre_fetch_file_for_download( msg = "Expected a file instance or UUID to download." raise ValueError(msg) if to_local_path is None and warn_missing_path: - msg = "The file will be downloaded as temporary." - log_user_warning(msg) + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=1).warning( + "Download: file will be saved to a temporary location" + ) valid_local_path_or_none: Path | None = ( Path(to_local_path) if to_local_path else None @@ -476,7 +489,9 @@ def __pre_fetch_file_for_download( uuid.UUID(file_uuid) if isinstance(file_uuid, str) else file_uuid ) if client.dry_run: - log_user("Dry run enabled: a sample file is being returned instead") + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + "Download dry-run: returning a sample file instead of downloading" + ) file_instance: File = files.generate_sample_file(valid_uuid) else: file_bytes: bytes = client._gateway.get_file_by_id(uuid=valid_uuid.hex) @@ -484,7 +499,12 @@ def __pre_fetch_file_for_download( return file_instance, valid_uuid, valid_local_path_or_none -def __upload_file_mux(*, client: Client, file_instance: File) -> File: # noqa: C901 +def __upload_file_mux( # noqa: C901 + *, + client: Client, + file_instance: File, + progress_callback: collections.abc.Callable[[int], None] | None = None, +) -> File: """Uploads a file instance to SDS, choosing the right upload mode.""" file_path = file_instance.local_path # check whether sds already has this file for this user @@ -506,12 +526,14 @@ def __upload_file_mux(*, client: Client, file_instance: File) -> File: # noqa: if verbose: log_user(f"Uploading contents and metadata for '{file_path}'") return __upload_contents_and_metadata( - client=client, file_instance=file_instance + client=client, + file_instance=file_instance, + progress_callback=progress_callback, ) case FileUploadMode.UPLOAD_METADATA_ONLY: if verbose: log_user(f"Uploading only metadata for '{file_path}'") - log.debug( + log.bind(cat=LogCategory.UPLOAD).debug( f"Uploading only metadata '{file_path}' with sibling '{asset_id}'" ) if asset_id is None: @@ -523,7 +545,9 @@ def __upload_file_mux(*, client: Client, file_instance: File) -> File: # noqa: case FileUploadMode.UPDATE_METADATA_ONLY: # pragma: no cover if verbose: log_user(f"Updating metadata for '{file_path}'") - log.debug(f"Updating metadata '{file_path}' with asset '{asset_id}'") + log.bind(cat=LogCategory.UPLOAD).debug( + f"Updating metadata '{file_path}' with asset '{asset_id}'" + ) assert asset_id is not None, "Expected an asset ID when updating metadata" return __update_existing_file_metadata_only( client=client, file_instance=file_instance, asset_id=asset_id @@ -537,6 +561,7 @@ def __upload_contents_and_metadata( *, client: Client, file_instance: File, + progress_callback: collections.abc.Callable[[int], None] | None = None, ) -> File: """UPLOADS a new file instance to SDS with contents and metadata. @@ -551,7 +576,9 @@ def __upload_contents_and_metadata( return file_instance assert not client.dry_run, "Internal error: expected dry run to be disabled." - file_response = client._gateway.upload_new_file(file_instance=file_instance) + file_response = client._gateway.upload_new_file( + file_instance=file_instance, progress_callback=progress_callback + ) uploaded_file = File.model_validate_json(file_response) uploaded_file.local_path = file_instance.local_path return uploaded_file diff --git a/sdk/src/spectrumx/api/uploads.py b/sdk/src/spectrumx/api/uploads.py index f2a936742..720f19be2 100644 --- a/sdk/src/spectrumx/api/uploads.py +++ b/sdk/src/spectrumx/api/uploads.py @@ -3,9 +3,11 @@ from __future__ import annotations import asyncio +import contextlib import hashlib import json import sys +import threading from datetime import UTC from datetime import datetime from pathlib import Path @@ -14,6 +16,7 @@ from typing import NoReturn from anyio import Path as AsyncPath +from loguru import logger as log from pydantic import BaseModel from pydantic import Field from tqdm import tqdm @@ -30,8 +33,11 @@ # pydantic complains if not defined out of type checking block ) from spectrumx.ops import files as file_ops +from spectrumx.utils import LogCategory +from spectrumx.utils import credit_unstreamed_file_bytes from spectrumx.utils import get_prog_bar from spectrumx.utils import is_test_env +from spectrumx.utils import log_context from spectrumx.utils import log_user from spectrumx.utils import log_user_warning from spectrumx.vendor.xdg_base_dirs import xdg_state_home @@ -275,6 +281,7 @@ class UploadWorkload(BaseModel): sds_path: PurePosixPath = Field(default_factory=lambda: PurePosixPath("/")) max_concurrent_uploads: int = 5 persist_state: bool = True + progress_log_period_secs: int = 30 # file buffers fq_discovered: list[File] = Field(default_factory=list) @@ -295,6 +302,7 @@ class UploadWorkload(BaseModel): tqdm[NoReturn] | None # pyrefly: ignore[bad-specialization] ) = None _persistence_manager: UploadPersistenceManager | None = None + _progress_log_task: asyncio.Task | None = None model_config = { "populate_by_name": True, @@ -304,6 +312,7 @@ class UploadWorkload(BaseModel): def model_post_init(self, context) -> None: """Initialize the workload with an async lock.""" self._state_lock = asyncio.Lock() + self._prog_bar_lock = threading.Lock() self._workload_id = self._compute_workload_id() self._persistence_manager = UploadPersistenceManager( local_root=self.local_root, @@ -316,7 +325,6 @@ def model_post_init(self, context) -> None: **upload_prog_bar_kwargs, # pyrefly: ignore[bad-argument-type] leave=True, ) - self._prog_uploaded_bytes.disable = True self._prog_uploaded_bytes.clear() super().model_post_init(context) @@ -476,16 +484,32 @@ async def _mark_completed_result(self, successful_result: Result[File]) -> None: self.fq_completed.append(uploaded_file) assert self._persistence_manager is not None await self._persistence_manager.save_persisted_upload(uploaded_file) - await self._update_prog_bars(num_bytes=uploaded_file.size) + await self._update_prog_bars() + log.bind(cat=LogCategory.UPLOAD).info( + f"Uploaded: {uploaded_file.name}", + file_name=uploaded_file.name, + file_size=uploaded_file.size, + ) + + def _update_prog_bar_bytes(self, num_bytes: int) -> None: + """Thread-safe byte counter update for the shared upload progress bar.""" + bar = self._prog_uploaded_bytes + if bar is None: + return + with self._prog_bar_lock: + bar.update(num_bytes) async def _update_prog_bars(self, num_bytes: int | None = None) -> None: """Update progress bars based on current upload state.""" if self.verbose and isinstance(self._prog_uploaded_bytes, tqdm): async with self._state_lock: - self._prog_uploaded_bytes.disable = False - self._prog_uploaded_bytes.set_description(self._get_progress_string()) - if num_bytes: - self._prog_uploaded_bytes.update(num_bytes) + with self._prog_bar_lock: + self._prog_uploaded_bytes.disable = False + self._prog_uploaded_bytes.set_description( + self._get_progress_string() + ) + if num_bytes: + self._prog_uploaded_bytes.update(num_bytes) def _get_progress_string(self) -> str: """Returns a progress string for the upload progress bar.""" @@ -529,7 +553,8 @@ async def _reset_progress(self) -> None: self.fq_completed.clear() self.fq_failed.clear() if isinstance(self._prog_uploaded_bytes, tqdm): - self._prog_uploaded_bytes.clear() + with self._prog_bar_lock: + self._prog_uploaded_bytes.clear() @property def total_files(self) -> int: @@ -560,14 +585,44 @@ async def _upload_next_file(self) -> Result[File]: return Result(exception=StopAsyncIteration("No more files to upload")) try: + # Throttled wrapper: accumulate bytes and only push to tqdm + # every ~100 KB, to avoid excessive refresh overhead (mirrors + # the same pattern used in _download_files). + _acc: list[int] = [0] + _streamed: list[int] = [0] + _throttle = 100_000 + bar = self._prog_uploaded_bytes + + def _throttled_update(n: int) -> None: + if bar is None: + return + _acc[0] += n + _streamed[0] += n + if _acc[0] >= _throttle: + self._update_prog_bar_bytes(_acc[0]) + _acc[0] = 0 + result = Result( value=await asyncio.to_thread( self.client._sds_files.upload_file, # noqa: SLF001 # pyright: ignore[reportPrivateUsage] client=self.client, local_file=next_file, sds_path=self.sds_path, + progress_callback=_throttled_update if bar is not None else None, ) ) + # Flush any remaining bytes that didn't reach the threshold + if bar is not None and _acc[0] > 0: + self._update_prog_bar_bytes(_acc[0]) + _acc[0] = 0 + if bar is not None: + credited = credit_unstreamed_file_bytes( + file_size=next_file.size, + bytes_streamed=_streamed[0], + prog_bar=None, + ) + if credited: + self._update_prog_bar_bytes(credited) await self._mark_completed_result(successful_result=result) except SDSError as err: await self._mark_failed_file(sds_file=next_file, reason=str(err)) @@ -603,16 +658,50 @@ async def _has_pending_files(self) -> bool: async with self._state_lock: return bool(self.fq_pending) + async def _periodic_progress_logger(self) -> None: + """Emit consolidated progress log every N seconds.""" + while True: + await asyncio.sleep(self.progress_log_period_secs) + async with self._state_lock: + completed = len(self.fq_completed) + total = len(self.fq_discovered) + total_bytes = self.total_bytes + uploaded_bytes = sum(f.size for f in self.fq_completed) + failures = len(self.fq_failed) + in_progress = len(self.fq_in_progress) + skipped = len(self.fq_skipped) + log.bind(cat=LogCategory.UPLOAD).info( + "Upload progress", + completed=completed, + total=total, + bytes_uploaded=uploaded_bytes, + bytes_total=total_bytes, + failures=failures, + in_progress=in_progress, + skipped=skipped, + ) + async def _execute_uploads(self) -> None: """Execute uploads concurrently using worker coroutines.""" await self._reset_progress() + # Start periodic progress logger + self._progress_log_task = asyncio.create_task(self._periodic_progress_logger()) + workers = [self._upload_worker(i) for i in range(self.max_concurrent_uploads)] - await asyncio.gather(*workers) + try: + await asyncio.gather(*workers) + finally: + # Cancel progress logger + if self._progress_log_task: + self._progress_log_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._progress_log_task # close progress bar if isinstance(self._prog_uploaded_bytes, tqdm): - self._prog_uploaded_bytes.close() + with self._prog_bar_lock: + self._prog_uploaded_bytes.close() async def _reset_state(self) -> None: """Reset all state buffers.""" @@ -632,7 +721,8 @@ async def _register_discovered_file(self, file_model: File) -> None: self.fq_pending.append(file_model) self.total_bytes += file_model.size if isinstance(self._prog_uploaded_bytes, tqdm): - self._prog_uploaded_bytes.total = self.total_bytes + with self._prog_bar_lock: + self._prog_uploaded_bytes.total = self.total_bytes async def _add_skipped(self, skipped: SkippedUpload) -> None: """Add a skipped file entry.""" @@ -649,26 +739,60 @@ async def run(self, client: Client) -> list[Result[File]]: List of results for each uploaded file. """ self.client = client - await self._discover_files() - await self._execute_uploads() + with log_context( + upload_id=self._workload_id, + upload_dir=str(self.local_root), + ): + await self._discover_files() + await self._execute_uploads() + + # Completion summary + elapsed = datetime.now(UTC) - ( + self.discovery_started_at or datetime.now(UTC) + ) + elapsed_sec = elapsed.total_seconds() + uploaded_bytes = sum(f.size for f in self.fq_completed) + avg_speed_bps: float | None = ( + uploaded_bytes / elapsed_sec if elapsed_sec > 0 else None + ) + status = "clean" if len(self.fq_failed) == 0 else "interrupted" + log.bind( + cat=LogCategory.UPLOAD, + total_files=len(self.fq_discovered), + total_bytes=self.total_bytes, + uploaded=len(self.fq_completed), + failed=len(self.fq_failed), + skipped=len(self.fq_skipped), + elapsed_seconds=elapsed_sec, + avg_speed_bps=avg_speed_bps, + status=status, + ).info( + "Upload complete", + ) - results = [Result(value=file_obj) for file_obj in self.fq_completed] - results.extend( - Result( - exception=error, - error_info={ - "sds_file": error.sds_file, - "reason": error.reason, - "error": str(error), - }, + results = [Result(value=file_obj) for file_obj in self.fq_completed] + results.extend( + Result( + exception=error, + error_info={ + "sds_file": error.sds_file, + "reason": error.reason, + "error": str(error), + }, + ) + for error in self.fq_failed ) - for error in self.fq_failed - ) - return results + return results @staticmethod def _remove_from_buffer(sds_file: File, buffer: list[File]) -> None: + """Remove from buffer, matching by local_path (persists through upload).""" + if sds_file.local_path: + for i, f in enumerate(buffer): + if f.local_path and f.local_path == sds_file.local_path: + buffer.pop(i) + return try: buffer.remove(sds_file) except ValueError: @@ -719,6 +843,7 @@ def upload_resumable( verbose=verbose, warn_skipped=warn_skipped, persist_state=persist_state, + progress_log_period_secs=client._config.progress_log_period_secs, # noqa: SLF001 ) return asyncio.run(upload_workload.run(client=client)) diff --git a/sdk/src/spectrumx/client.py b/sdk/src/spectrumx/client.py index 7ff6b4450..158416c78 100644 --- a/sdk/src/spectrumx/client.py +++ b/sdk/src/spectrumx/client.py @@ -1,7 +1,10 @@ """Client for the SpectrumX Data System.""" +import collections.abc +import time from collections.abc import Collection from collections.abc import Mapping +from datetime import UTC from datetime import datetime from pathlib import Path from pathlib import PurePosixPath @@ -23,15 +26,20 @@ from spectrumx.ops.pagination import Paginator from . import __version__ +from . import utils from .config import SDSConfig from .gateway import GatewayClient from .models.files import File from .ops import files +from .utils import LogCategory from .utils import clean_local_path +from .utils import credit_unstreamed_file_bytes +from .utils import enable_structured_logging from .utils import get_prog_bar from .utils import log_user from .utils import log_user_error from .utils import log_user_warning +from .utils import set_persistent_log_context def _normalize_top_level_dir_prefix( @@ -96,6 +104,7 @@ def __init__( env_file: Path | None = None, env_config: Mapping[str, Any] | None = None, verbose: bool = False, + log_file: Path | None = None, ) -> None: # avoids circular imports from spectrumx.api import sds_files as _sds_files_api # noqa: PLC0415 @@ -112,6 +121,22 @@ def __init__( env_config=env_config, sds_host=host, verbose=self.verbose, + log_file=log_file, + ) + + # Enable structured logging + _log_path = self._config.log_file + enable_structured_logging(log_path=_log_path) + log_user(f"Structured log: {utils._current_log_path}") # noqa: SLF001 + + # Bind persistent context + api_key_prefix = "" + if self._config.api_key: + api_key_prefix = self._config.api_key.split(".")[0] + + set_persistent_log_context( + api_key_prefix=api_key_prefix, + timeout=self._config.timeout, ) # initialize the gateway @@ -139,11 +164,7 @@ def __str__(self) -> str: def _issue_user_alerts(self) -> None: """Logs important messages on initialization.""" - if __version__.startswith("0.1."): - log_user_warning( - "This version of the SDK is in early development. " - "Expect breaking changes in the future." - ) + log_user(f"SDK version: {__version__}") if self.dry_run: log_user("Dry run enabled: no SDS requests will be made or files written") @@ -203,9 +224,11 @@ def authenticate(self) -> None: log_user("Dry run is enabled: assuming successful authentication") log_user("To authenticate against SDS, set Client.dry_run to False") else: - log.warning(f"Dry run DISABLED: authenticating against '{self.host}'") + log.bind(cat=LogCategory.AUTH).warning( + f"Dry run DISABLED: authenticating against '{self.host}'" + ) self._gateway.authenticate() - log.info("Authenticated successfully") + log.bind(cat=LogCategory.AUTH).info("Authenticated successfully") self.is_authenticated = True # ======= FILE METHODS @@ -315,7 +338,9 @@ def _prepare_download_directory(self, to_local_path: Path) -> None: """Prepare the download directory.""" if not to_local_path.exists(): if self.dry_run: - log_user(f"Dry run: would create the directory '{to_local_path}'") + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + f"Download dry-run: directory would be created at '{to_local_path}'" + ) else: to_local_path.mkdir(parents=True) @@ -330,12 +355,13 @@ def _get_files_to_download( ) -> DownloadFileSource: """Get the list of files to download.""" if self.dry_run: - log_user( - "Called download() in dry run mode: " - "files are fabricated and not written to disk." + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + "Download dry-run: files are simulated (not written to disk)" ) files_to_download = files.generate_random_files(num_files=10) - log_user(f"Dry run: discovered {len(files_to_download)} files (samples)") + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + f"Download dry-run: discovered {len(files_to_download)} sample files" + ) else: if from_sds_path is not None: files_to_download = self.list_files( @@ -350,7 +376,10 @@ def _get_files_to_download( ) raise ValueError(error_msg) if verbose: - log_user(f"Discovered {len(files_to_download)} files") + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + f"Download: discovered {len(files_to_download)} files " + f"in SDS path '{from_sds_path}'" + ) return files_to_download @@ -365,10 +394,183 @@ def _download_files( ) -> list[Result[File]]: """Download the files and return results.""" prefix = "Dry-run: simulating download:" if self.dry_run else "Downloading:" + + total_files = len(files_to_download) + + # Mutable container so per-chunk closure can update byte count + bytes_downloaded_shared: list[int] = [0] + _download_start = datetime.now(UTC) + + # Only compute total bytes for lists; Paginator would be fully consumed. + if isinstance(files_to_download, list): + total_bytes_total = sum(f.size for f in files_to_download) + results = self._download_files_with_byte_progress( + files_to_download=files_to_download, + total_bytes_total=total_bytes_total, + to_local_path=to_local_path, + skip_contents=skip_contents, + overwrite=overwrite, + verbose=verbose, + prefix=prefix, + total_files=total_files, + period=self._config.progress_log_period_secs, + bytes_downloaded_shared=bytes_downloaded_shared, + ) + else: + total_bytes_total = None + results = self._download_files_fallback( + files_to_download=files_to_download, + to_local_path=to_local_path, + skip_contents=skip_contents, + overwrite=overwrite, + verbose=verbose, + prefix=prefix, + ) + + # Completion summary (shared by both branches) + completed = sum(1 for r in results if r) + failed = len(results) - completed + elapsed = datetime.now(UTC) - _download_start + elapsed_sec = elapsed.total_seconds() + transferred_bytes = ( + bytes_downloaded_shared[0] if total_bytes_total is not None else None + ) + avg_speed_bps: float | None = ( + transferred_bytes / elapsed_sec + if transferred_bytes is not None and elapsed_sec > 0 + else None + ) + status = "clean" if failed == 0 else "interrupted" + log.bind( + cat=LogCategory.DOWNLOAD, + total_files=total_files, + total_bytes=total_bytes_total, + downloaded=completed, + failed=failed, + elapsed_seconds=elapsed_sec, + avg_speed_bps=avg_speed_bps, + status=status, + ).info( + "Download complete", + ) + + return results + + def _download_files_with_byte_progress( # noqa: PLR0913 + self, + *, + files_to_download: list[File], + total_bytes_total: int, + to_local_path: Path, + skip_contents: bool, + overwrite: bool, + verbose: bool, + prefix: str, + total_files: int, + period: float, + bytes_downloaded_shared: list[int], + ) -> list[Result[File]]: + """Byte-level progress bar with throttled per-chunk callback.""" prog_bar = get_prog_bar( - files_to_download, desc="Downloading", disable=not verbose + total=total_bytes_total, + desc="Downloading", + unit="B", + unit_scale=True, + unit_divisor=1024, + disable=not verbose or utils.is_test_env(), ) + # Throttle callback updates: ~100 KB between tqdm refreshes + _acc: list[int] = [0] + _file_streamed: list[int] = [0] + _throttle = 100_000 + + def _on_download_bytes(n: int) -> None: + """Per-chunk callback: accumulate bytes, update bar every ~100 KB.""" + _acc[0] += n + _file_streamed[0] += n + bytes_downloaded_shared[0] += n + if _acc[0] >= _throttle: + prog_bar.update(_acc[0]) + _acc[0] = 0 + + _last_progress_log = 0.0 + results: list[Result[File]] = [] + for file_info in files_to_download: + _file_streamed[0] = 0 + prog_bar.set_description(f"{prefix} '{file_info.name}'") + result = self.download_single_file( + file_info=file_info, + to_local_path=to_local_path, + skip_contents=skip_contents, + overwrite=overwrite, + progress_callback=_on_download_bytes, + ) + results.append(result) + + if _acc[0] > 0: + prog_bar.update(_acc[0]) + _acc[0] = 0 + if result: + credit_unstreamed_file_bytes( + file_size=file_info.size, + bytes_streamed=_file_streamed[0], + prog_bar=prog_bar, + bytes_accounted=bytes_downloaded_shared, + ) + + # Per-file completion log + if result: + log.bind(cat=LogCategory.DOWNLOAD).info( + f"Downloaded: {file_info.name}", + file_name=file_info.name, + file_size=file_info.size, + ) + else: + log.bind(cat=LogCategory.DOWNLOAD).warning( + f"Download failed: {file_info.name}", + file_name=file_info.name, + error=str(result.exception_or(Exception("Unknown"))), + ) + + # Periodic progress log + now = time.monotonic() + if now - _last_progress_log >= period: + _last_progress_log = now + completed = sum(1 for r in results if r) + failed = len(results) - completed + log.bind(cat=LogCategory.DOWNLOAD).info( + "Download progress", + completed=completed, + total=total_files, + bytes_downloaded=bytes_downloaded_shared[0], + bytes_total=total_bytes_total, + failures=failed, + ) + + # Flush any remaining accumulated bytes + if _acc[0] > 0: + prog_bar.update(_acc[0]) + prog_bar.refresh() + + return results + + def _download_files_fallback( + self, + *, + files_to_download: DownloadFileSource, + to_local_path: Path, + skip_contents: bool, + overwrite: bool, + verbose: bool, + prefix: str, + ) -> list[Result[File]]: + """Fallback: file-level progress bar (total bytes unknown, e.g. Paginator).""" + prog_bar = get_prog_bar( + files_to_download, + desc="Downloading", + disable=not verbose or utils.is_test_env(), + ) results: list[Result[File]] = [] for file_info in prog_bar: prog_bar.set_description(f"{prefix} '{file_info.name}'") @@ -380,6 +582,20 @@ def _download_files( ) results.append(result) + # Per-file completion log + if result: + log.bind(cat=LogCategory.DOWNLOAD).info( + f"Downloaded: {file_info.name}", + file_name=file_info.name, + file_size=file_info.size, + ) + else: + log.bind(cat=LogCategory.DOWNLOAD).warning( + f"Download failed: {file_info.name}", + file_name=file_info.name, + error=str(result.exception_or(Exception("Unknown"))), + ) + return results def download_single_file( @@ -389,6 +605,7 @@ def download_single_file( to_local_path: Path, skip_contents: bool, overwrite: bool, + progress_callback: collections.abc.Callable[[int], None] | None = None, ) -> Result[File]: """Download a single file and return the result.""" # Handle local files without UUID @@ -417,22 +634,31 @@ def download_single_file( # skip file download when local exists and not overwriting if not skip_contents and local_file_path.exists() and not overwrite: - log_user(f"Skipping local file: '{local_file_path}'") + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + f"Download skipped: local file exists " + f"(set overwrite=True to replace): '{local_file_path}'" + ) return Result(value=file_info) # never re-download identical files if not skip_contents and local_file_path.exists() and overwrite: if file_info.sum_blake3 == file_info.compute_sum_blake3(): - log_user(f"Skipping identical file: '{local_file_path}'") + log.bind(cat=LogCategory.DOWNLOAD).opt(depth=2).info( + f"Download skipped: local copy is identical to SDS version " + f"(BLAKE3 match): '{local_file_path}'" + ) return Result(value=file_info) # download the file try: - log.debug(f"Dw: {local_file_path}") + log.bind(cat=LogCategory.DOWNLOAD).debug( + f"Downloading file contents to: {local_file_path}" + ) downloaded_file = self.download_file( file_uuid=file_info.uuid, to_local_path=local_file_path, skip_contents=skip_contents, + progress_callback=progress_callback, ) return Result(value=downloaded_file) except SDSError as err: @@ -454,8 +680,10 @@ def _is_path_relative_to_target( """Check if the local file path is relative to the target directory.""" local_file_path = local_file_path.resolve() to_local_path = to_local_path.resolve() - log.debug(f"Resolved path: {local_file_path}") - log.debug(f"Target path: {to_local_path}") + log.bind(cat=LogCategory.DOWNLOAD).debug( + f"Download path resolved: {local_file_path}" + ) + log.bind(cat=LogCategory.DOWNLOAD).debug(f"Download target: {to_local_path}") return local_file_path.is_relative_to(to_local_path) def list_files( @@ -493,6 +721,7 @@ def download_file( to_local_path: Path | str | None = None, skip_contents: bool = False, warn_missing_path: bool = True, + progress_callback: collections.abc.Callable[[int], None] | None = None, ) -> File: """Downloads a file from SDS: metadata and maybe contents. @@ -503,12 +732,15 @@ def download_file( otherwise the download will create a temporary file on disk. Args: - file_instance: The file instance to download. - file_uuid: The UUID of the file to download. - to_local_path: The local path to save the downloaded file to. - skip_contents: When True, only the metadata is downloaded - and no files are created on disk. - warn_missing_path: Show a warning when the download location is undefined. + file_instance: The file instance to download. + file_uuid: The UUID of the file to download. + to_local_path: The local path to save the downloaded file to. + skip_contents: When True, only the metadata is downloaded + and no files are created on disk. + warn_missing_path: Show a warning when the download location + is undefined. + progress_callback: Optional callable invoked with byte count as + data is received from the stream (for progress tracking). Returns: The file instance with updated attributes, or a sample when in dry run. """ @@ -519,6 +751,7 @@ def download_file( to_local_path=to_local_path, skip_contents=skip_contents, warn_missing_path=warn_missing_path, + progress_callback=progress_callback, ) def download_dataset( diff --git a/sdk/src/spectrumx/config.py b/sdk/src/spectrumx/config.py index bb3e1e0b1..0ca1126de 100644 --- a/sdk/src/spectrumx/config.py +++ b/sdk/src/spectrumx/config.py @@ -14,6 +14,7 @@ from spectrumx.errors import Unset from .models import SDSModel +from .utils import LogCategory from .utils import into_human_bool from .utils import log_user from .utils import log_user_warning @@ -40,8 +41,10 @@ class Attr: CFG_NAME_LOOKUP: dict[str, Attr] = { "dry_run": Attr(attr_name="dry_run", cast_fn=into_human_bool), "http_timeout": Attr(attr_name="timeout", cast_fn=int), + "log_file": Attr(attr_name="log_file"), "sds_host": Attr(attr_name="sds_host"), "sds_secret_token": Attr(attr_name="api_key"), + "progress_log_period_secs": Attr(attr_name="progress_log_period_secs", cast_fn=int), } @@ -74,6 +77,8 @@ class SDSConfig: api_key: str = "" dry_run: bool = True # safer default timeout: int = DEFAULT_HTTP_TIMEOUT + log_file: Path | None = None + progress_log_period_secs: int = 30 _active_config: list[Attr] _env_file: Path | None = None @@ -85,6 +90,7 @@ def __init__( env_config: Mapping[str, Any] | None = None, sds_host: None | str = None, verbose: bool = False, + log_file: Path | None = None, ) -> None: """Initialize the configuration. Args: @@ -92,8 +98,10 @@ def __init__( env_config: Overrides for the environment file. sds_host: The host to connect to (the config file has priority over this). verbose: Show which config files are loaded and which attributes are set. + log_file: Path for the structured JSONL log file. """ self.sds_host = sds_host + self.log_file = log_file self.init_config(env_file=env_file, env_config=env_config, verbose=verbose) def init_config( @@ -126,7 +134,7 @@ def init_config( def show_config(self, log_fn: Callable[[str], None] = print) -> None: """Show the active configuration.""" header = "SDS_Config: active configuration:" - log.debug(header) # for sdk developers + log.bind(cat=LogCategory.CONFIG).debug(header) # for sdk developers log_fn(header) # for users for attr in self._active_config: _log_redacted( @@ -138,16 +146,15 @@ def show_config(self, log_fn: Callable[[str], None] = print) -> None: def _set_config(self, clean_config: list[Attr]) -> None: """Sets the instance attributes.""" for attr in clean_config: - setattr( - self, - attr.attr_name, - attr.attr_value, - ) - _log_redacted(key=attr.attr_name, value=attr.attr_value) + value = attr.attr_value + if attr.attr_name == "log_file" and isinstance(value, str): + value = Path(value) + setattr(self, attr.attr_name, value) + _log_redacted(key=attr.attr_name, value=value) # validate attributes / show user warnings if not self.api_key: - log.error( + log.bind(cat=LogCategory.CONFIG).error( "SDS_Config: API key not set. Check your environment" " file has an SDS_SECRET_TOKEN set." ) @@ -177,7 +184,9 @@ def __load_config( if secret := os.environ.get("SDS_SECRET_TOKEN"): env_vars["SDS_SECRET_TOKEN"] = secret env_vars = {k: v for k, v in env_vars.items() if v is not None} - log.debug(f"SDS_Config: from local env: {list(env_vars.keys())}") + log.bind(cat=LogCategory.CONFIG).debug( + f"SDS_Config: from local env: {list(env_vars.keys())}" + ) # merge file, cli, and env vars configs if ( @@ -225,14 +234,14 @@ def _clean_config( # skip if no value if sensitive_value is None: - log.warning(f"SDS_Config: {key} has no value") + log.bind(cat=LogCategory.CONFIG).warning(f"SDS_Config: {key} has no value") continue attr: Attr | None = name_lookup.get(normalized_key) # warn of invalid config if attr is None: msg = f"SDS_Config: {key} not recognized" - log.warning(msg) + log.bind(cat=LogCategory.CONFIG).warning(msg) logger.warning(msg) continue @@ -250,7 +259,7 @@ def _clean_config( def _log_redacted( *, key: str, - value: str | float | bool | None, + value: Any, log_fn: Callable[[str], None] = logger.info, depth: int = 1, ) -> None: @@ -269,7 +278,7 @@ def _log_redacted( ) del value msg = f"\tSDS_Config: set {key}={safe_value}" - log.opt(depth=depth).debug(msg) # for sdk developers + log.bind(cat=LogCategory.CONFIG).opt(depth=depth).debug(msg) # for sdk developers log_fn(msg) diff --git a/sdk/src/spectrumx/gateway.py b/sdk/src/spectrumx/gateway.py index a4d644755..511e436a8 100644 --- a/sdk/src/spectrumx/gateway.py +++ b/sdk/src/spectrumx/gateway.py @@ -2,6 +2,7 @@ import json import uuid +from collections.abc import Callable from collections.abc import Collection from collections.abc import Iterator from enum import StrEnum @@ -10,6 +11,7 @@ from pathlib import PurePosixPath from typing import Annotated from typing import Any +from typing import BinaryIO import requests from loguru import logger as log @@ -27,6 +29,8 @@ from .models.files import FileUpload from .models.files import PermissionRepresentation from .ops import network +from .utils import LogCategory +from .utils import LogContext from .utils import is_test_env from .utils import log_user_warning @@ -72,6 +76,33 @@ class FileContentsCheck(BaseModel): asset_id: Annotated[uuid.UUID | None, Field()] = None +class _ProgressFileReader: + """Wraps a binary file to report bytes read via a callback.""" + + def __init__( + self, + file: BinaryIO, + callback: Callable[[int], None], + ) -> None: + self._file = file + self._callback = callback + + def read(self, n: int = -1) -> bytes: + data = self._file.read(n) + if data: + self._callback(len(data)) + return data + + def __getattr__(self, name: str) -> Any: + return getattr(self._file, name) + + def __enter__(self): + return self + + def __exit__(self, *args: object) -> None: + self._file.close() + + class GatewayClient: """Communicates with the SDS API.""" @@ -177,15 +208,25 @@ def _request( debug_str += f"\nparams={kwargs['params']}" if "asset_id" in kwargs: debug_str += f"\nasset_id={kwargs['asset_id']}" - log.opt(depth=1).debug(debug_str) - - return requests.request( - timeout=self.timeout if timeout is None else timeout, - method=method, - stream=stream, - **payload, - **kwargs, - ) + log.bind(cat=LogCategory.NETWORK).opt(depth=1).debug(debug_str) + + url = payload["url"] + with LogContext(url=url): + try: + response = requests.request( + timeout=self.timeout if timeout is None else timeout, + method=method, + stream=stream, + **payload, + **kwargs, + ) + with LogContext(http_status=response.status_code): + return response + except requests.exceptions.RequestException as err: + log.bind(cat=LogCategory.NETWORK).opt(depth=1).error( + f"Network error: {method} {url}: {err}" + ) + raise @property def base_url(self) -> str: @@ -208,7 +249,7 @@ def authenticate(self, *, verbose: bool = False) -> None: msg = "No response code received from authentication request." raise AuthError(msg) status = HTTPStatus(code) - log.debug(f"Authentication response: {status}") + log.bind(cat=LogCategory.NETWORK).debug(f"Authentication response: {status}") if status.is_success: return msg = f"Authentication failed: {response.text}" @@ -328,13 +369,21 @@ def check_file_contents_exist( content: bytes | Any = response.content return FileContentsCheck.model_validate_json(content) - def upload_new_file(self, file_instance: File, *, verbose: bool = False) -> bytes: + def upload_new_file( + self, + file_instance: File, + *, + verbose: bool = False, + progress_callback: Callable[[int], None] | None = None, + ) -> bytes: """Uploads a local file to the SDS API. Uploads file contents and metadata. Args: - file_instance: The file to upload, as a models.File instance. + file_instance: The file to upload, as a models.File instance. + progress_callback: Optional callable invoked with byte count as + data is read from the source file (for progress tracking). """ if file_instance.local_path is None: msg = "Attempting to upload a remote file. Download it first." @@ -344,8 +393,13 @@ def upload_new_file(self, file_instance: File, *, verbose: bool = False) -> byte context={"mode": PermissionRepresentation.STRING} ) all_chunks: bytes = b"" + + file_ptr: BinaryIO | _ProgressFileReader = file_instance.local_path.open("rb") + if progress_callback is not None: + file_ptr = _ProgressFileReader(file_ptr, progress_callback) + with ( - file_instance.local_path.open("rb") as file_ptr, + file_ptr, self._request( method=HTTPMethods.POST, endpoint=Endpoints.FILES, @@ -643,7 +697,9 @@ def delete_capture( """ endpoint = f"{Endpoints.CAPTURES}/{capture_uuid.hex}" if self.verbose: - log.debug(f"Sending DELETE request to {endpoint}") + log.bind(cat=LogCategory.NETWORK).debug( + f"Sending DELETE request to {endpoint}" + ) response = self._request( method=HTTPMethods.DELETE, @@ -652,7 +708,9 @@ def delete_capture( ) network.success_or_raise(response, ContextException=CaptureError) if self.verbose: - log.debug(f"Capture with UUID {capture_uuid} deleted successfully") + log.bind(cat=LogCategory.NETWORK).debug( + f"Capture with UUID {capture_uuid} deleted successfully" + ) def revoke_capture_share_permissions( self, diff --git a/sdk/src/spectrumx/ops/files.py b/sdk/src/spectrumx/ops/files.py index 365640359..dd00097a2 100644 --- a/sdk/src/spectrumx/ops/files.py +++ b/sdk/src/spectrumx/ops/files.py @@ -11,6 +11,7 @@ from loguru import logger as log from spectrumx.models.files import File +from spectrumx.utils import LogCategory from spectrumx.utils import log_user from spectrumx.utils import log_user_warning @@ -38,7 +39,7 @@ def _load_undesired_globs(ignore_file: Path | None = None) -> list[str]: ] undesired_basenames = sorted([line for line in undesired_basenames if line]) else: - log.info("No .sds-ignore file found") + log.bind(cat=LogCategory.FILESYSTEM).info("No .sds-ignore file found") return undesired_basenames diff --git a/sdk/src/spectrumx/ops/network.py b/sdk/src/spectrumx/ops/network.py index 352c355ae..aa2ef3313 100644 --- a/sdk/src/spectrumx/ops/network.py +++ b/sdk/src/spectrumx/ops/network.py @@ -3,6 +3,7 @@ import sys from http import HTTPStatus +from spectrumx.utils import LogCategory from spectrumx.utils import is_test_env # python 3.11 backport @@ -33,7 +34,7 @@ def success_or_raise( if status.is_success: return - log.info(response) + log.bind(cat=LogCategory.NETWORK).info(response) try: error_json = response.json() error_details = error_json.get("detail", error_json) @@ -45,7 +46,7 @@ def success_or_raise( ) error_details = str(error_details) - log.opt(depth=1).exception(error_details) + log.bind(cat=LogCategory.NETWORK).opt(depth=1).exception(error_details) if status in {HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}: raise errors.AuthError(message=error_details) if status.is_client_error: @@ -62,7 +63,9 @@ def extract_error_details_from_html(response: requests.Response) -> str: BeautifulSoup, # pyright: ignore[reportMissingModuleSource] ) except ImportError: - log.warning("Install 'BeautifulSoup' to have a more detailed error message.") + log.bind(cat=LogCategory.NETWORK).warning( + "Install 'BeautifulSoup' to have a more detailed error message." + ) reason = response.reason if reason is None: reason = "Unknown reason" @@ -78,4 +81,11 @@ def extract_error_details_from_html(response: requests.Response) -> str: combined_text = "\n".join(texts) no_blanks = "\n".join(filter(None, combined_text.splitlines())) - return no_blanks.strip() + extracted = no_blanks.strip() + if extracted: + return extracted + + reason = response.reason + if reason is None: + reason = "Unknown reason" + return reason diff --git a/sdk/src/spectrumx/ops/pagination.py b/sdk/src/spectrumx/ops/pagination.py index c96eab586..c4309a344 100644 --- a/sdk/src/spectrumx/ops/pagination.py +++ b/sdk/src/spectrumx/ops/pagination.py @@ -18,6 +18,7 @@ from spectrumx.gateway import GatewayClient from spectrumx.models import SDSModel from spectrumx.ops import files +from spectrumx.utils import LogCategory from spectrumx.utils import log_user_warning if TYPE_CHECKING: @@ -192,7 +193,7 @@ def _total_pages(self) -> int: def _debug(self, message: str, depth: int = 1) -> None: """Logs a debug message if verbose mode is enabled.""" if self._verbose: # pragma: no cover - log.opt(depth=depth).debug(message) + log.bind(cat=LogCategory.FILESYSTEM).opt(depth=depth).debug(message) @property def _has_next_page(self) -> bool: @@ -204,7 +205,9 @@ def _has_next_page(self) -> bool: def _fetch_next_page(self) -> None: """Fetches the next page of results.""" if self._is_fetching: # pragma: no cover - log.warning("Already fetching the next page.") + log.bind(cat=LogCategory.FILESYSTEM).warning( + "Already fetching the next page." + ) return self._is_fetching = True # try-finally to unset self._is_fetching when done @@ -234,7 +237,7 @@ def _fetch_next_page(self) -> None: # log an unexpected FileError if it happens if "invalid page" not in str(err).lower(): msg = "Unexpected error while fetching the next page:" - log.exception(msg) + log.bind(cat=LogCategory.FILESYSTEM).exception(msg) msg = "No more pages available." raise StopIteration(msg) from err self._next_page += 1 @@ -301,7 +304,7 @@ def _process_file_fake(my_file: files.File) -> None: # pragma: no cover def main() -> None: # pragma: no cover """Usage example for paginator.""" - log.info("Running the main script.") + log.bind(cat=LogCategory.FILESYSTEM).info("Running the main script.") file_paginator = Paginator[files.File]( Entry=files.File, gateway=GatewayClient( @@ -318,23 +321,27 @@ def main() -> None: # pragma: no cover dry_run=True, # in dry-run this should always generate 2.5 pages verbose=True, ) - log.info(f"Total files matched: {len(file_paginator)}") + log.bind(cat=LogCategory.FILESYSTEM).info( + f"Total files matched: {len(file_paginator)}" + ) processed_count: int = 0 for my_file in file_paginator: - log.info(f"Processing file: {my_file.name}") + log.bind(cat=LogCategory.FILESYSTEM).info(f"Processing file: {my_file.name}") _process_file_fake(my_file) processed_count += 1 # new pages are fetched automatically - log.info(f"Processed {processed_count} / {len(file_paginator)} files.") + log.bind(cat=LogCategory.FILESYSTEM).info( + f"Processed {processed_count} / {len(file_paginator)} files." + ) - log.info("Trying another loop:") + log.bind(cat=LogCategory.FILESYSTEM).info("Trying another loop:") for _my_file in file_paginator: msg = "This will not run, as the paginator was consumed." raise AssertionError(msg) - log.info("No more files to process.") + log.bind(cat=LogCategory.FILESYSTEM).info("No more files to process.") - log.info("Paginator demo finished.") + log.bind(cat=LogCategory.FILESYSTEM).info("Paginator demo finished.") if __name__ == "__main__": diff --git a/sdk/src/spectrumx/utils.py b/sdk/src/spectrumx/utils.py index a2d8faa05..85ed421cc 100644 --- a/sdk/src/spectrumx/utils.py +++ b/sdk/src/spectrumx/utils.py @@ -2,13 +2,23 @@ from __future__ import annotations +import contextlib +import contextvars +import importlib.metadata +import json import logging import os +import platform import random import re import string +import tempfile +from datetime import UTC +from datetime import datetime +from enum import StrEnum from pathlib import Path from typing import TYPE_CHECKING +from typing import Any from typing import NoReturn from typing import TypeVar from typing import cast @@ -18,6 +28,8 @@ from tqdm import auto as auto_tqdm from tqdm import tqdm +from spectrumx.vendor.xdg_base_dirs import xdg_state_home + if TYPE_CHECKING: from collections.abc import Iterable @@ -168,3 +180,193 @@ def get_prog_bar( "tqdm[T]", # pyrefly: ignore[bad-specialization] auto_tqdm.tqdm(iterable, *args, **kwargs), ) + + +def credit_unstreamed_file_bytes( + *, + file_size: int, + bytes_streamed: int, + prog_bar: tqdm[NoReturn] | None, + bytes_accounted: list[int] | None = None, +) -> int: + """Credit progress for file bytes that were not transferred over the wire. + + Byte-level progress tracks streamed chunks. When a file completes without + transferring its full contents (skip, metadata-only, local exists, etc.), + credit the remaining bytes so totals stay accurate. + + Returns: + The number of bytes credited. + """ + remaining = max(file_size - bytes_streamed, 0) + if remaining == 0: + return 0 + if bytes_accounted is not None: + bytes_accounted[0] += remaining + if prog_bar is not None: + prog_bar.update(remaining) + return remaining + + +# --- Structured Logging --- + + +class LogCategory(StrEnum): + """Categories for structured log messages.""" + + LOG = "log" + CONFIG = "config" + AUTH = "auth" + NETWORK = "network" + FILESYSTEM = "filesystem" + DOWNLOAD = "download" + UPLOAD = "upload" + + +LOG_CATEGORY_LOG = LogCategory.LOG + + +_log_context: contextvars.ContextVar[dict[str, Any] | None] = contextvars.ContextVar( + "_log_context", default=None +) + +_persistent_context: dict[str, Any] = {} +_structured_sink_id: int | None = None +_current_log_path: Path | None = None + + +class LogContext: + """Context manager for scoped structured logging context binding.""" + + def __init__(self, **kwargs: Any) -> None: + self._kwargs = kwargs + self._token: contextvars.Token[dict[str, Any] | None] | None = None + + def __enter__(self) -> None: + current = (_log_context.get() or {}).copy() + current.update(self._kwargs) + self._token = _log_context.set(current) + + def __exit__(self, *args: object) -> None: + if self._token is not None: + _log_context.reset(self._token) + + +# Alias for LogContext - used by uploads.py +log_context = LogContext + + +def set_persistent_log_context(**kwargs: Any) -> None: + """Set persistent context fields that appear on all log lines.""" + _persistent_context.update(kwargs) + + +def _get_default_log_path() -> Path: + """Resolve default structured log path: XDG state home > temp dir.""" + try: + base = xdg_state_home() / "spectrumx" / "logs" + base.mkdir(parents=True, exist_ok=True) + return base / f"{datetime.now(UTC).strftime('%Y-%m-%d')}.jsonl" + except (ImportError, OSError): + return Path(tempfile.gettempdir()) / "spectrumx_sdk.jsonl" + + +def _emit_boot_message(log_path: Path) -> None: + """Write system info boot message to the log file.""" + try: + sdk_version = importlib.metadata.version("spectrumx") + except importlib.metadata.PackageNotFoundError: + sdk_version = "unknown" + + boot_entry = { + "ts": datetime.now(UTC).isoformat(), + "pid": os.getpid(), + "lvl": "INFO", + "cat": LogCategory.LOG, + "msg": "system info", + "sdk_version": sdk_version, + "os": platform.platform(), + "python": platform.python_version(), + "log_file": str(log_path), + } + try: + with log_path.open("a", encoding="utf-8") as f: + f.write(json.dumps(boot_entry, default=str) + "\n") + except OSError: + pass + + +def _structured_log_sink(message: Any) -> None: + """Loguru function sink: reads ContextVar, writes JSON line to file.""" + if _current_log_path is None: + return + + rec = message.record + msg_str = rec["message"].rstrip() + + entry: dict[str, Any] = { + "ts": rec["time"].isoformat(), + "pid": rec["process"].id, + "lvl": rec["level"].name, + "cat": rec.get("extra", {}).get("cat", LogCategory.LOG), + "msg": msg_str, + } + entry.update(rec.get("extra", {})) + + # Merge persistent context + entry.update({k: v for k, v in _persistent_context.items() if v is not None}) + + # Merge scoped context + ctx = _log_context.get() or {} + entry.update({k: v for k, v in ctx.items() if v is not None}) + + # Extract exc_info from loguru record (for error log lines) + exc_info = rec.get("exception") + if exc_info is not None and exc_info.type is not None: + entry["exc_info"] = f"{exc_info.type.__name__}: {exc_info.value}" + + try: + with _current_log_path.open("a", encoding="utf-8") as f: + f.write(json.dumps(entry, default=str) + "\n") + except OSError: + pass + + +def enable_structured_logging(log_path: Path | str | None = None) -> None: + """Enable or reconfigure the structured JSONL logging sink.""" + global _structured_sink_id, _current_log_path # noqa: PLW0603 + + if _structured_sink_id is not None: + with contextlib.suppress(Exception): + log.remove(_structured_sink_id) + + # Resolve path + resolved = Path(log_path) if log_path is not None else _get_default_log_path() + + _current_log_path = resolved + parent = resolved.parent + parent.mkdir(parents=True, exist_ok=True) + + # Emit system info boot message every time the logger is initialized + _emit_boot_message(resolved) + + try: + _structured_sink_id = log.add( + _structured_log_sink, + format="{message}", + filter=lambda record: True, + ) + except Exception: # noqa: BLE001 + _structured_sink_id = None + + +def reset_structured_logging() -> None: + """Reset structured logging state. Useful for tests.""" + global _structured_sink_id, _current_log_path # noqa: PLW0603 + if _structured_sink_id is not None: + with contextlib.suppress(Exception): + log.remove(_structured_sink_id) + _structured_sink_id = None + _current_log_path = None + _persistent_context.clear() + _log_context.set(None) diff --git a/sdk/tests/ops/test_datasets.py b/sdk/tests/ops/test_datasets.py index b3e23fc68..380a2d678 100644 --- a/sdk/tests/ops/test_datasets.py +++ b/sdk/tests/ops/test_datasets.py @@ -56,8 +56,8 @@ def test_get_files_artifacts_only_clears_capture_uuids_and_warns() -> None: ): paginator = api.get_files(ds, capture_uuids=(cap,), artifacts_only=True) assert len(paginator) == 0 - log_mock.warning.assert_called_once() - assert "capture" in log_mock.warning.call_args[0][0].lower() + log_mock.bind.return_value.warning.assert_called_once() + assert "capture" in log_mock.bind.return_value.warning.call_args[0][0].lower() kw = mock_get.call_args.kwargs assert kw.get("capture_uuids") in (None, ()) assert kw["artifacts_only"] is True @@ -78,7 +78,7 @@ def test_get_files_artifacts_only_with_top_level_dirs_logs_info() -> None: artifacts_only=True, ) assert len(paginator) == 0 - log_mock.info.assert_called_once() + log_mock.bind.return_value.info.assert_called_once() kw = mock_get.call_args.kwargs assert kw["top_level_dirs"] == ("/pytest/root",) assert kw["artifacts_only"] is True @@ -154,3 +154,105 @@ def test_revoke_dataset_share_permissions( assert len(responses.calls) == 1 assert responses.calls[0].request.method == "PUT" assert responses.calls[0].request.url == revoke_url + + +def test_get_dry_run(client: Client) -> None: + """Dry run for get() returns a shell Dataset without HTTP calls.""" + client.dry_run = True + dataset_uuid = uuidlib.uuid4() + result = client.datasets.get(dataset_uuid) + assert result.uuid == dataset_uuid + + +@responses.activate +def test_get_gateway(client: Client, responses: responses.RequestsMock) -> None: + """get() calls gateway and returns a parsed Dataset.""" + client.dry_run = False + dataset_uuid = uuidlib.uuid4() + mock_data = { + "uuid": dataset_uuid.hex, + "name": "test-dataset-from-gateway", + } + responses.add( + method=responses.GET, + url=get_datasets_endpoint(client, dataset_id=dataset_uuid.hex), + json=mock_data, + status=200, + ) + result = client.datasets.get(dataset_uuid) + assert result.uuid is not None + assert result.uuid.hex == dataset_uuid.hex + assert result.name == "test-dataset-from-gateway" + assert len(responses.calls) == 1 + + +def test_list_captures_dry_run(client: Client) -> None: + """Dry run for list_captures() returns an empty list.""" + client.dry_run = True + dataset_uuid = uuidlib.uuid4() + result = client.datasets.list_captures(dataset_uuid) + assert result == [] + + +@responses.activate +def test_list_captures_gateway( + client: Client, responses: responses.RequestsMock +) -> None: + """list_captures() calls gateway and returns parsed capture payloads.""" + client.dry_run = False + dataset_uuid = uuidlib.uuid4() + mock_data = { + "captures": [ + {"uuid": str(uuidlib.uuid4()), "name": "capture-1"}, + ], + } + responses.add( + method=responses.GET, + url=get_datasets_endpoint(client, dataset_id=dataset_uuid.hex), + json=mock_data, + status=200, + ) + result = client.datasets.list_captures(dataset_uuid) + assert len(result) == 1 + assert result[0]["name"] == "capture-1" + assert len(responses.calls) == 1 + + +def test_list_artifact_files_dry_run(client: Client) -> None: + """Dry run for list_artifact_files() returns an empty list.""" + client.dry_run = True + dataset_uuid = uuidlib.uuid4() + result = client.datasets.list_artifact_files(dataset_uuid) + assert result == [] + + +@responses.activate +def test_list_artifact_files_gateway( + client: Client, responses: responses.RequestsMock +) -> None: + """list_artifact_files() calls gateway and returns parsed file payloads.""" + client.dry_run = False + dataset_uuid = uuidlib.uuid4() + mock_data = { + "files": [ + {"uuid": str(uuidlib.uuid4()), "name": "file-1.txt"}, + ], + } + responses.add( + method=responses.GET, + url=get_datasets_endpoint(client, dataset_id=dataset_uuid.hex), + json=mock_data, + status=200, + ) + result = client.datasets.list_artifact_files(dataset_uuid) + assert len(result) == 1 + assert result[0]["name"] == "file-1.txt" + assert len(responses.calls) == 1 + + +def test_revoke_share_permissions_dry_run(client: Client) -> None: + """Dry run for revoke_share_permissions() returns True.""" + client.dry_run = True + dataset_uuid = uuidlib.uuid4() + result = client.datasets.revoke_share_permissions(dataset_uuid) + assert result is True diff --git a/sdk/tests/ops/test_files.py b/sdk/tests/ops/test_files.py index ed163629f..6c0ea832a 100644 --- a/sdk/tests/ops/test_files.py +++ b/sdk/tests/ops/test_files.py @@ -17,10 +17,14 @@ import responses from spectrumx import Client from spectrumx.api.sds_files import delete_file +from spectrumx.api.sds_files import download_file from spectrumx.api.sds_files import file_list_time_query_param from spectrumx.api.sds_files import list_files +from spectrumx.api.uploads import UploadPersistenceManager from spectrumx.errors import FileError +from spectrumx.errors import SDSError from spectrumx.gateway import API_TARGET_VERSION +from spectrumx.models.files import File from spectrumx.ops.files import ( _load_undesired_globs, # pyright: ignore[reportPrivateUsage] ) @@ -870,3 +874,466 @@ def test_large_valid_file(self, tmp_path: Path) -> None: f"Large valid file should pass validation. Reasons: {reasons}" ) assert reasons == [], "No reasons should be given for valid files" + + +# ============================================================================= +# sds_files.py coverage tests (phase 2) +# ============================================================================= + + +def test_upload_file_invalid_type(client: Client) -> None: + """upload_file raises TypeError for non-File/Path/str input.""" + with pytest.raises(TypeError, match="file_path must be a Path, str, or File"): + client.upload_file(local_file=123) # type: ignore[arg-type] + + +def test_upload_file_instance_directory_composition( + client: Client, tmp_path: Path +) -> None: + """Uploading a File instance with directory set prepends sds_path.""" + client.dry_run = True + sub_path = PurePosixPath("sub/dir") + file_path = tmp_path / "test_compose.txt" + file_path.write_text("content") + + file_instance = File( + uuid=uuidlib.uuid4(), + name=file_path.name, + media_type="text/plain", + size=file_path.stat().st_size, + directory=sub_path, + permissions="rw-r--r--", + created_at=datetime(2024, 12, 1, 12, 0, 0, tzinfo=UTC), + updated_at=datetime(2024, 12, 1, 12, 0, 0, tzinfo=UTC), + expiration_date=datetime(2026, 12, 1, 12, 0, 0, tzinfo=UTC), + local_path=file_path, + ) + + result = client.upload_file( + local_file=file_instance, + sds_path=PurePosixPath("/base"), + ) + + assert result.directory == PurePosixPath("/base/sub/dir"), ( + f"Expected '/base/sub/dir', got '{result.directory}'" + ) + + +def test_upload_file_skip_mode( + client: Client, + responses: responses.RequestsMock, + temp_file_with_text_contents: Path, +) -> None: + """When file already exists in tree, upload_file returns input without UUID.""" + client.dry_run = False + asset_id = uuidlib.uuid4() + + responses.add( + method=responses.POST, + url=get_content_check_endpoint(client), + status=201, + json={ + "file_contents_exist_for_user": True, + "file_exists_in_tree": True, + "user_mutable_attributes_differ": False, + "asset_id": asset_id.hex, + }, + ) + + result = client.upload_file( + local_file=temp_file_with_text_contents, + sds_path=PurePosixPath("/test/skip"), + ) + + assert result.uuid is None, "SKIP mode should return the local file without UUID" + assert result.local_path == temp_file_with_text_contents + assert len(responses.calls) == 1, "Only content-check request expected" + + +def test_upload_file_metadata_only_mode( + client: Client, + responses: responses.RequestsMock, + temp_file_with_text_contents: Path, +) -> None: + """When contents exist for user, upload metadata only.""" + client.dry_run = False + asset_id = uuidlib.uuid4() + file_id = uuidlib.uuid4() + file_size = temp_file_with_text_contents.stat().st_size + + # content-check response: contents exist for user but not in tree + responses.add( + method=responses.POST, + url=get_content_check_endpoint(client), + status=201, + json={ + "file_contents_exist_for_user": True, + "file_exists_in_tree": False, + "user_mutable_attributes_differ": True, + "asset_id": asset_id.hex, + }, + ) + + # metadata-only upload response + responses.add( + method=responses.POST, + url=get_files_endpoint(client), + status=201, + json={ + "uuid": file_id.hex, + "name": temp_file_with_text_contents.name, + "media_type": "text/plain", + "size": file_size, + "directory": "/test/metadata-only", + "permissions": "rw-r--r--", + "created_at": "2024-12-01T12:00:00Z", + "updated_at": "2024-12-01T12:00:00Z", + "expiration_date": "2026-12-01T12:00:00Z", + }, + ) + + result = client.upload_file( + local_file=temp_file_with_text_contents, + sds_path=PurePosixPath("/test/metadata-only"), + ) + + assert result.uuid == file_id, "Metadata-only upload should return a server File" + assert result.directory == PurePosixPath("/test/metadata-only") + assert len(responses.calls) == 2, ( # noqa: PLR2004 + "Expected content-check + metadata-upload calls" + ) + + +def test_download_file_with_file_instance( + client: Client, responses: responses.RequestsMock +) -> None: + """download_file accepts a File instance directly.""" + client.dry_run = False + file_id = uuidlib.uuid4() + file_contents = "downloaded via file-instance" + + now_ts = datetime(2024, 12, 1, 12, 0, 0, tzinfo=UTC) + + file_instance = File( + uuid=file_id, + name="test.txt", + media_type="text/plain", + size=len(file_contents), + directory=PurePosixPath("/test/"), + permissions="rw-r--r--", + created_at=now_ts, + updated_at=now_ts, + expiration_date=datetime(2026, 12, 1, 12, 0, 0, tzinfo=UTC), + ) + + # only the download endpoint is needed (file info is in the instance) + responses.add( + method=responses.GET, + url=_download_file_endpoint(client, file_id=file_id.hex), + status=200, + body=file_contents, + ) + + result = download_file(client=client, file_instance=file_instance) + + assert result.uuid == file_id + assert result.local_path is not None + assert result.local_path.exists() + assert result.local_path.read_text() == file_contents + + # cleanup + result.local_path.unlink(missing_ok=True) + + +def test_download_file_skip_contents( + client: Client, responses: responses.RequestsMock +) -> None: + """skip_contents=True prevents downloading file contents.""" + client.dry_run = False + file_id = uuidlib.uuid4() + now_ts = datetime(2024, 12, 1, 12, 0, 0, tzinfo=UTC) + + file_instance = File( + uuid=file_id, + name="test.txt", + media_type="text/plain", + size=100, + directory=PurePosixPath("/test/"), + permissions="rw-r--r--", + created_at=now_ts, + updated_at=now_ts, + expiration_date=datetime(2026, 12, 1, 12, 0, 0, tzinfo=UTC), + ) + + # Do NOT register a download endpoint — if skip_contents works, + # the download endpoint should NOT be called. + result = download_file( + client=client, + file_instance=file_instance, + skip_contents=True, + ) + + assert result.uuid == file_id + assert result.local_path is None, "No local path should be set when skipping" + assert len(responses.calls) == 0, "No HTTP calls should be made" + + +def test_download_file_missing_uuid(client: Client) -> None: + """Downloading a File instance without a UUID raises ValueError.""" + now_ts = datetime(2024, 12, 1, 12, 0, 0, tzinfo=UTC) + + file_instance = File( + name="test.txt", + media_type="text/plain", + size=100, + directory=PurePosixPath("/test/"), + permissions="rw-r--r--", + created_at=now_ts, + updated_at=now_ts, + expiration_date=datetime(2026, 12, 1, 12, 0, 0, tzinfo=UTC), + ) + + with pytest.raises(ValueError, match="local reference"): + download_file(client=client, file_instance=file_instance) + + +def test_download_file_contents_error_cleanup( + client: Client, tmp_path: Path, responses: responses.RequestsMock +) -> None: + """When download fails, temp files are cleaned up and exception propagates.""" + client.dry_run = False + file_id = uuidlib.uuid4() + target_path = tmp_path / "should-not-exist.txt" + now_ts = datetime(2024, 12, 1, 12, 0, 0, tzinfo=UTC) + + file_instance = File( + uuid=file_id, + name="test.txt", + media_type="text/plain", + size=100, + directory=PurePosixPath("/test/"), + permissions="rw-r--r--", + created_at=now_ts, + updated_at=now_ts, + expiration_date=datetime(2026, 12, 1, 12, 0, 0, tzinfo=UTC), + ) + + responses.add( + method=responses.GET, + url=_download_file_endpoint(client, file_id=file_id.hex), + status=500, + json={"detail": "Simulated download failure"}, + ) + + with pytest.raises(SDSError, match="Simulated download failure"): + download_file( + client=client, + file_instance=file_instance, + to_local_path=target_path, + ) + + assert not target_path.exists(), ( + "Target path should not exist after failed download" + ) + # verify no stray .downloading temp files remain + temp_files = list(tmp_path.glob("*.downloading")) + assert len(temp_files) == 0, f"Temp files not cleaned up: {temp_files}" + + +def test_delete_file_get_file_fails( + client: Client, responses: responses.RequestsMock +) -> None: + """When get_file fails before deletion, a warning is logged but delete proceeds.""" + test_uuid = uuidlib.uuid4() + test_uuid_hex = test_uuid.hex + client.dry_run = False + + # GET request for file info will receive a response that causes parse failure + responses.add( + method=responses.GET, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=200, + body=b"not-valid-json", + ) + + # DELETE request succeeds + responses.add( + method=responses.DELETE, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=204, + ) + + result = delete_file(client=client, file_uuid=test_uuid) + + assert result is True, "Delete should succeed even if get_file fails" + assert len(responses.calls) == 2, "Expected GET + DELETE calls" # noqa: PLR2004 + + +def test_delete_file_cleanup_called( + client: Client, responses: responses.RequestsMock +) -> None: + """When sum_blake3 is set and delete succeeds, cleanup is triggered.""" + test_uuid = uuidlib.uuid4() + test_uuid_hex = test_uuid.hex + test_checksum = "abc123def456" + client.dry_run = False + + # GET returns file info WITH sum_blake3 + responses.add( + method=responses.GET, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=200, + json={ + "uuid": test_uuid_hex, + "name": "test.txt", + "media_type": "text/plain", + "size": 100, + "directory": "/test/", + "permissions": "rw-r--r--", + "sum_blake3": test_checksum, + "created_at": "2024-12-01T12:00:00Z", + "updated_at": "2024-12-01T12:00:00Z", + "expiration_date": "2026-12-01T12:00:00Z", + }, + ) + + # DELETE request succeeds + responses.add( + method=responses.DELETE, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=204, + ) + + with patch.object( + UploadPersistenceManager, + "remove_persisted_uploads_by_checksum", + ) as mock_cleanup: + result = delete_file(client=client, file_uuid=test_uuid) + + assert result is True + mock_cleanup.assert_called_once_with(checksum=test_checksum) + + assert len(responses.calls) == 2, "Expected GET + DELETE calls" # noqa: PLR2004 + + +def test_delete_file_gateway_returns_false( + client: Client, responses: responses.RequestsMock +) -> None: + """When the gateway rejects the DELETE, the error propagates as FileError.""" + test_uuid = uuidlib.uuid4() + test_uuid_hex = test_uuid.hex + client.dry_run = False + + # GET returns file info + responses.add( + method=responses.GET, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=200, + json={ + "uuid": test_uuid_hex, + "name": "test.txt", + "media_type": "text/plain", + "size": 100, + "directory": "/test/", + "permissions": "rw-r--r--", + "created_at": "2024-12-01T12:00:00Z", + "updated_at": "2024-12-01T12:00:00Z", + "expiration_date": "2026-12-01T12:00:00Z", + }, + ) + + responses.add( + method=responses.DELETE, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=400, + json={"detail": "Cannot delete"}, + ) + + with pytest.raises(FileError, match="Cannot delete"): + delete_file(client=client, file_uuid=test_uuid) + + +def test_list_files_empty_result( + client: Client, responses: responses.RequestsMock +) -> None: + """list_files returns no files when the gateway returns an empty page.""" + client.dry_run = False + + responses.add( + method=responses.GET, + url=get_files_endpoint(client), + status=200, + json={"count": 0, "results": []}, + ) + + paginator = list_files( + client=client, + sds_path=PurePosixPath("/test/empty"), + ) + results = list(paginator) + + assert results == [], "Expected no files from empty listing" + + +def test_list_files_dry_run(client: Client) -> None: + """list_files in dry run returns a paginator without network calls.""" + client.dry_run = True + + paginator = list_files( + client=client, + sds_path=PurePosixPath("/test/dry"), + ) + + # In dry run, the paginator yields sample files + results = list(paginator) + assert len(results) > 0, "Expected sample files in dry run" + for f in results: + assert f.uuid is not None + assert isinstance(f.name, str) + assert len(f.name) > 0 + + +def test_cleanup_persisted_uploads_error( + client: Client, responses: responses.RequestsMock +) -> None: + """When UploadPersistenceManager raises, the error is logged but not re-raised.""" + test_uuid = uuidlib.uuid4() + test_uuid_hex = test_uuid.hex + client.dry_run = False + + # GET returns file info WITH sum_blake3 + responses.add( + method=responses.GET, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=200, + json={ + "uuid": test_uuid_hex, + "name": "test.txt", + "media_type": "text/plain", + "size": 100, + "directory": "/test/", + "permissions": "rw-r--r--", + "sum_blake3": "somechecksum", + "created_at": "2024-12-01T12:00:00Z", + "updated_at": "2024-12-01T12:00:00Z", + "expiration_date": "2026-12-01T12:00:00Z", + }, + ) + + # DELETE succeeds + responses.add( + method=responses.DELETE, + url=f"{get_files_endpoint(client)}{test_uuid_hex}/", + status=204, + ) + + with patch.object( + UploadPersistenceManager, + "remove_persisted_uploads_by_checksum", + side_effect=OSError("Disk error"), + ): + # Should NOT raise — error is caught and logged + result = delete_file(client=client, file_uuid=test_uuid) + + assert result is True, "Delete should still succeed even if cleanup fails" + assert len(responses.calls) == 2, "Expected GET + DELETE calls" # noqa: PLR2004 diff --git a/sdk/tests/ops/test_network.py b/sdk/tests/ops/test_network.py index f09b84347..1f06b5f9f 100644 --- a/sdk/tests/ops/test_network.py +++ b/sdk/tests/ops/test_network.py @@ -6,9 +6,12 @@ # ruff: noqa: SLF001 # pyright: ignore[reportPrivateUsage] +from unittest.mock import patch + import pytest import requests from spectrumx import errors +from spectrumx.ops.network import extract_error_details_from_html from spectrumx.ops.network import success_or_raise @@ -74,3 +77,64 @@ class CustomException(errors.SDSError): with pytest.raises(CustomException, match="Bad Request"): success_or_raise(response, ContextException=CustomException) + + +def test_success_or_raise_no_status_code() -> None: + """No status code should raise SDSError.""" + response = requests.Response() + object.__setattr__(response, "status_code", None) + with pytest.raises(errors.SDSError, match="No status code"): + success_or_raise(response) + + +def test_success_or_raise_json_decode_error_no_bs4() -> None: + """Non-JSON error body falls back to reason when not in test env.""" + response = requests.Response() + response.status_code = 500 + response._content = b"Not JSON" + response.reason = "Internal Server Error" + with ( + patch("spectrumx.ops.network.is_test_env", return_value=False), + pytest.raises(errors.ServiceError, match="Internal Server Error"), + ): + success_or_raise(response) + + +def test_success_or_raise_json_decode_error_with_bs4() -> None: + """Non-JSON error body extracts error from HTML when in test env.""" + response = requests.Response() + response.status_code = 500 + response._content = b'
  • Error details
    • ' + with ( + patch("spectrumx.ops.network.is_test_env", return_value=True), + pytest.raises(errors.ServiceError, match="Error details"), + ): + success_or_raise(response) + + +def test_success_or_raise_catchall_fallback() -> None: + """3xx status falls through to generic SDSError.""" + response = requests.Response() + response.status_code = 300 + response._content = b'{"detail": "redirected"}' + with pytest.raises(errors.SDSError, match="redirected"): + success_or_raise(response) + + +def test_extract_error_details_from_html_no_matching_element() -> None: + """When HTML lacks #summary/#pastebinTraceback, falls back to reason.""" + response = requests.Response() + response.status_code = 500 + response._content = b"No useful error element" + response.reason = "Internal Server Error" + result = extract_error_details_from_html(response) + assert result == "Internal Server Error" + + +def test_extract_error_details_from_html_with_bs4() -> None: + """Extracts text from HTML id=summary element.""" + response = requests.Response() + response.status_code = 500 + response._content = b'
  • Extracted error
    • ' + result = extract_error_details_from_html(response) + assert result == "Extracted error" diff --git a/sdk/tests/test_client.py b/sdk/tests/test_client.py index c53c383d3..146434b3d 100644 --- a/sdk/tests/test_client.py +++ b/sdk/tests/test_client.py @@ -6,17 +6,33 @@ from pathlib import Path from pathlib import PurePosixPath from typing import Any +from unittest.mock import MagicMock from unittest.mock import patch +from uuid import UUID import pytest +import responses from loguru import logger as log from spectrumx.client import Client +from spectrumx.client import _normalize_top_level_dir_prefix from spectrumx.client import resolve_dataset_capture_filter_params from spectrumx.config import CFG_NAME_LOOKUP from spectrumx.config import SDSConfig +from spectrumx.errors import CaptureError +from spectrumx.errors import Result +from spectrumx.errors import SDSError +from spectrumx.gateway import API_TARGET_VERSION +from spectrumx.models.captures import CaptureType from spectrumx.models.files import File from spectrumx.ops import files +from tests.conftest import get_captures_endpoint +from tests.conftest import get_content_check_endpoint +from tests.conftest import get_datasets_endpoint +from tests.conftest import get_files_endpoint + +_DRY_RUN_FILE_COUNT = 10 + class LogLevels(IntEnum): """Log levels for testing.""" @@ -231,7 +247,6 @@ def test_download_dry_run_happy_path( def test_download_fails_for_invalid_files( caplog: pytest.LogCaptureFixture, - capsys: pytest.CaptureFixture[str], client: Client, ) -> None: """Ensures download() fails when an invalid local path is provided.""" @@ -264,16 +279,9 @@ def _get_problematic_list_of_files(num_files: int = 10) -> list[File]: attribute="generate_random_files", side_effect=_get_problematic_list_of_files, ): - _ = capsys.readouterr() results = client.download(from_sds_path=sds_path, to_local_path=local_path) assert len(results) == target_num_files, "Three files should be generated" - captured = capsys.readouterr() - log.debug(captured.err) - assert "simulating download" in captured.err, ( - "Expected 'simulating download' in stderr output" - ) - # assert "Skipping local file" in captured.err successful_files = [result() for result in results if result] error_infos = [result.error_info for result in results if not result] log.error( @@ -445,3 +453,513 @@ def test_resolve_dataset_capture_filter_normalizes_top_level_dirs() -> None: assert active assert uuids is None assert dirs == ["/foo/bar", "/baz"] + + +# ====================================================================== +# _normalize_top_level_dir_prefix +# ====================================================================== + + +def test_normalize_top_level_dir_prefix() -> None: + """Covers _normalize_top_level_dir_prefix edge cases.""" + assert _normalize_top_level_dir_prefix("foo/bar") == "/foo/bar" + assert _normalize_top_level_dir_prefix("/baz/") == "/baz" + assert _normalize_top_level_dir_prefix("/") == "/" + assert _normalize_top_level_dir_prefix("a\\b") == "/a/b" + + +# ====================================================================== +# resolve_dataset_capture_filter_params edge cases +# ====================================================================== + + +def test_resolve_dataset_capture_filter_none_params() -> None: + """When both captures and dirs are None, filter is not active (line 61).""" + active, uuids, dirs = resolve_dataset_capture_filter_params( + capture_uuids=None, + top_level_dirs=None, + dry_run=False, + ) + assert not active + assert uuids is None + assert dirs is None + + +def test_resolve_dataset_capture_filter_empty_collections() -> None: + """Empty collections produce active filter with empty sets (lines 71-75).""" + active, uuids, dirs = resolve_dataset_capture_filter_params( + capture_uuids=[], + top_level_dirs=[], + dry_run=False, + ) + assert active + assert uuids is not None + assert len(uuids) == 0 + assert dirs is not None + assert len(dirs) == 0 + + +def test_resolve_dataset_capture_filter_uuid_str_conversion() -> None: + """String UUIDs are properly converted to UUID objects (line 71).""" + uuid_str = "12345678-1234-5678-1234-567812345678" + active, uuids, _dirs = resolve_dataset_capture_filter_params( + capture_uuids=[uuid_str], + top_level_dirs=None, + dry_run=False, + ) + assert active + assert uuids is not None + assert len(uuids) == 1 + u = next(iter(uuids)) + assert isinstance(u, UUID) + assert str(u) == uuid_str + + +# ====================================================================== +# Client miscellaneous properties and methods +# ====================================================================== + + +def test_client_str() -> None: + """Client.__str__ returns a descriptive string (line 159).""" + client = Client(host="sds-dev.crc.nd.edu") + assert "sds-dev.crc.nd.edu" in str(client) + + +def test_verbose_enables_output( + client: Client, caplog: pytest.LogCaptureFixture +) -> None: + """Verbose mode produces log output during operations.""" + caplog.set_level(LogLevels.DEBUG) + client.dry_run = True + client.verbose = True + client.get_file(file_uuid=uuid.uuid4()) + assert len(caplog.records) > 0 + + +def test_non_verbose_suppresses_output( + client: Client, caplog: pytest.LogCaptureFixture +) -> None: + """Non-verbose mode suppresses verbose log output.""" + caplog.set_level(LogLevels.DEBUG) + client.dry_run = True + client.verbose = False + client.get_file(file_uuid=uuid.uuid4()) + verbose_messages = [r for r in caplog.records if "verbose" in r.message.lower()] + assert len(verbose_messages) == 0 + + +def test_client_base_url_properties(client: Client) -> None: + """base_url and base_url_no_port properties (lines 214, 219).""" + assert client.base_url.startswith("http") + assert isinstance(client.base_url_no_port, str) + + +# ====================================================================== +# Download method tests +# ====================================================================== + + +def test_download_mutual_exclusion( + tmp_path: Path, + client: Client, +) -> None: + """from_sds_path and files_to_download are mutually exclusive (lines 308-314).""" + with pytest.raises(ValueError, match="Both a path in the SDS"): + client.download( + from_sds_path="/some/path", + to_local_path=tmp_path, + files_to_download=[], + ) + + +def _get_file_download_endpoint(client: Client, file_id: str) -> str: + """Returns the endpoint for downloading a file's contents.""" + return ( + client.base_url + f"/api/{API_TARGET_VERSION}/assets/files/{file_id}/download/" + ) + + +@responses.activate +def test_download_single_file_sdserror( + tmp_path: Path, + client: Client, + caplog: pytest.LogCaptureFixture, + responses: responses.RequestsMock, +) -> None: + """download_single_file catches SDSError (lines 461-463).""" + client.dry_run = False + caplog.set_level(LogLevels.ERROR) + file_info = files.generate_sample_file(uuid.uuid4()) + file_info.directory = PurePosixPath("remote/dir") + assert file_info.uuid is not None + file_id_hex = file_info.uuid.hex + + # File info GET succeeds + responses.add( + method=responses.GET, + url=get_files_endpoint(client) + f"{file_id_hex}/", + status=200, + json={ + "uuid": file_id_hex, + "name": file_info.name, + "media_type": "text/plain", + "size": 100, + "directory": "/remote/dir", + "permissions": "rw-r--r--", + "created_at": "2024-12-01T12:00:00Z", + "updated_at": "2024-12-01T12:00:00Z", + "expiration_date": "2026-12-01T12:00:00Z", + }, + ) + # Download GET fails with server error + responses.add( + method=responses.GET, + url=_get_file_download_endpoint(client, file_id_hex), + status=500, + json={"detail": "Internal Server Error"}, + ) + + result = client.download_single_file( + file_info=file_info, + to_local_path=tmp_path, + skip_contents=False, + overwrite=False, + ) + assert not result + exc = result.exception_or(None) + assert exc is not None + assert ( + "500" in str(exc) + or "Server Error" in str(exc) + or "Internal Server Error" in str(exc) + ) + + +def test_download_no_path_no_files( + tmp_path: Path, + client: Client, +) -> None: + """Missing both path and files raises ValueError (lines 369-374).""" + client.dry_run = False + with pytest.raises(ValueError, match="Either a path in the SDS"): + client.download(to_local_path=tmp_path) + + +# ====================================================================== +# download_dataset endpoint tests +# ====================================================================== + + +@pytest.mark.parametrize("uuid_arg", [uuid.uuid4(), str(uuid.uuid4())]) +def test_download_dataset_dry_run( + tmp_path: Path, + client: Client, + uuid_arg: uuid.UUID | str, +) -> None: + """download_dataset works in dry run mode (lines 588-652).""" + results = client.download_dataset( + dataset_uuid=uuid_arg, + to_local_path=tmp_path, + verbose=False, + ) + assert len(results) == _DRY_RUN_FILE_COUNT + + +def test_download_dataset_empty_filters_active( + tmp_path: Path, + client: Client, +) -> None: + """download_dataset with empty filters produces empty file list (lines 604-606). + + When filter_active=True but both uuid_set and dir_prefixes_norm are empty, + files_to_download is set to an empty list directly. + """ + client.dry_run = False + results = client.download_dataset( + dataset_uuid=uuid.uuid4(), + to_local_path=tmp_path, + capture_uuids=[], + top_level_dirs=[], + verbose=False, + ) + assert len(results) == 0 + + +@responses.activate +@pytest.mark.parametrize( + ("capture_uuids", "top_level_dirs"), + [ + ([uuid.uuid4()], ["/some/path"]), + ([uuid.uuid4()], None), + (None, ["/some/path"]), + ], +) +def test_download_dataset_filter_branches( + tmp_path: Path, + client: Client, + responses: responses.RequestsMock, + capture_uuids, + top_level_dirs, +) -> None: + """download_dataset filter_active paths with HTTP-level mocking. + + Covers all three branches: both filters active, only uuids, and only dirs. + """ + client.dry_run = False + dataset_uuid = uuid.uuid4() + # Mock the dataset files endpoint to return empty results + responses.add( + method=responses.GET, + url=get_datasets_endpoint(client, dataset_id=dataset_uuid.hex) + "files/", + status=200, + json={"count": 0, "results": []}, + ) + results = client.download_dataset( + dataset_uuid=dataset_uuid, + to_local_path=str(tmp_path), + capture_uuids=capture_uuids, + top_level_dirs=top_level_dirs, + verbose=False, + ) + assert len(results) == 0 + + +# ====================================================================== +# Dataset read methods with string UUID (lines 662-678) +# ====================================================================== + + +def test_get_dataset_string_uuid(client: Client) -> None: + """get_dataset accepts a string UUID (lines 662-664).""" + ds = client.get_dataset(dataset_uuid=str(uuid.uuid4())) + assert ds.uuid is not None + + +def test_list_dataset_captures_string_uuid(client: Client) -> None: + """list_dataset_captures accepts a string UUID (lines 668-670).""" + result = client.list_dataset_captures(dataset_uuid=str(uuid.uuid4())) + assert result == [] + + +def test_list_dataset_artifact_files_string_uuid(client: Client) -> None: + """list_dataset_artifact_files accepts a string UUID (lines 676-678).""" + result = client.list_dataset_artifact_files(dataset_uuid=str(uuid.uuid4())) + assert result == [] + + +# ====================================================================== +# Upload method tests +# ====================================================================== + + +@responses.activate +def test_upload_capture_sdserror( + client: Client, + tmp_path: Path, + responses: responses.RequestsMock, +) -> None: + """upload_capture: SDSError during capture creation returns None (lines 872-877).""" + client.dry_run = False + local_file = tmp_path / "test_upload.txt" + local_file.write_text("test content") + + file_id = uuid.uuid4() + + # The resumable upload discovers files in tmp_path (test file + log file). + # Each file triggers: content-check POST then file-upload POST. + for _ in range(2): + # Content check endpoint — file does not exist in tree + responses.add( + method=responses.POST, + url=get_content_check_endpoint(client), + status=200, + json={ + "file_contents_exist_for_user": False, + "file_exists_in_tree": False, + "user_mutable_attributes_differ": True, + }, + ) + # File upload endpoint succeeds + responses.add( + method=responses.POST, + url=get_files_endpoint(client), + status=201, + json={ + "uuid": file_id.hex, + "name": "test_upload.txt", + "media_type": "text/plain", + "size": local_file.stat().st_size, + "directory": "/", + "permissions": "rw-r--r--", + "created_at": "2024-12-01T12:00:00Z", + "updated_at": "2024-12-01T12:00:00Z", + "expiration_date": "2026-12-01T12:00:00Z", + }, + ) + # Capture creation endpoint fails + responses.add( + method=responses.POST, + url=get_captures_endpoint(client), + status=400, + json={"detail": "Capture failed"}, + ) + + result = client.upload_capture( + local_path=tmp_path, + sds_path="/", + capture_type=CaptureType.DigitalRF, + verbose=False, + raise_on_error=False, + ) + assert result is None + + +def test_upload_multichannel_drf_capture_empty_channels( + client: Client, + tmp_path: Path, +) -> None: + """upload_multichannel with empty channels returns [] (lines 954-956).""" + # Create a file so the resumable upload discovers something + local_file = tmp_path / "test_upload.txt" + local_file.write_text("test content") + + client.dry_run = True + results = client.upload_multichannel_drf_capture( + local_path=tmp_path, + sds_path="/", + channels=[], + verbose=False, + ) + assert results is not None + assert len(results) == 0 + + +def test_upload_multichannel_drf_capture_dry_run( + client: Client, + tmp_path: Path, +) -> None: + """upload_multichannel works end-to-end in dry run mode (lines 937-987).""" + channels = ["chan1", "chan2"] + # Create a file so the resumable upload discovers something + local_file = tmp_path / "test_upload.txt" + local_file.write_text("test content") + + client.dry_run = True + results = client.upload_multichannel_drf_capture( + local_path=tmp_path, + sds_path="/", + channels=channels, + verbose=False, + ) + assert results is not None + assert len(results) == len(channels) + for capture in results: + assert capture.capture_type == CaptureType.DigitalRF + + +# ====================================================================== +# _handle_existing_capture_error tests (lines 889-902) +# +# These test a private utility method directly because triggering the +# specific error-handling branches through the public API would require +# complex HTTP mocking with no additional behavioral coverage benefit. +# ====================================================================== + + +def test_handle_existing_capture_error_non_capture_error(client: Client) -> None: + """Non-CaptureError returns (False, None) (lines 889-890).""" + err = SDSError("Some other error") + handled, capture = client._handle_existing_capture_error(err) # noqa: SLF001 + assert not handled + assert capture is None + + +def test_handle_existing_capture_error_no_existing_uuid(client: Client) -> None: + """CaptureError without existing UUID returns (False, None) (lines 892-894).""" + err = CaptureError("Generic capture error without UUID info") + handled, capture = client._handle_existing_capture_error(err) # noqa: SLF001 + assert not handled + assert capture is None + + +def test_handle_existing_capture_error_read_fails(client: Client) -> None: + """CaptureError with existing UUID, but read fails (lines 899-900).""" + err = CaptureError( + "drf_unique_channel_and_tld another capture: " + "12345678-1234-5678-1234-567812345678" + ) + with patch.object(client.captures, "read", side_effect=SDSError("Read failed")): + handled, capture = client._handle_existing_capture_error(err) # noqa: SLF001 + assert not handled + assert capture is None + + +def test_download_byte_progress_credits_skipped_content( + client: Client, tmp_path: Path +) -> None: + """Progress should reach file totals when content transfer is skipped.""" + file_info = files.generate_sample_file(uuid.uuid4()) + file_info.size = 1000 + bytes_downloaded_shared: list[int] = [0] + prog_bar = MagicMock() + + with ( + patch.object( + Client, + "download_single_file", + return_value=Result(value=file_info), + ), + patch("spectrumx.client.get_prog_bar", return_value=prog_bar), + ): + results = client._download_files_with_byte_progress( # noqa: SLF001 + files_to_download=[file_info], + total_bytes_total=file_info.size, + to_local_path=tmp_path, + skip_contents=True, + overwrite=False, + verbose=True, + prefix="Downloading", + total_files=1, + period=30.0, + bytes_downloaded_shared=bytes_downloaded_shared, + ) + + assert len(results) == 1 + assert results[0] + assert bytes_downloaded_shared[0] == file_info.size + prog_bar.update.assert_called_once_with(file_info.size) + + +def test_download_byte_progress_credits_partial_stream( + client: Client, tmp_path: Path +) -> None: + """Progress should include both streamed and unstreamed bytes per file.""" + file_info = files.generate_sample_file(uuid.uuid4()) + file_info.size = 1000 + bytes_downloaded_shared: list[int] = [0] + + def fake_download_single_file(*, progress_callback=None, **kwargs): + if progress_callback is not None: + progress_callback(400) + return Result(value=file_info) + + with patch.object( + Client, + "download_single_file", + side_effect=fake_download_single_file, + ): + client._download_files_with_byte_progress( # noqa: SLF001 + files_to_download=[file_info], + total_bytes_total=file_info.size, + to_local_path=tmp_path, + skip_contents=False, + overwrite=True, + verbose=True, + prefix="Downloading", + total_files=1, + period=30.0, + bytes_downloaded_shared=bytes_downloaded_shared, + ) + + assert bytes_downloaded_shared[0] == file_info.size diff --git a/sdk/tests/test_structured_logging.py b/sdk/tests/test_structured_logging.py new file mode 100644 index 000000000..674a3875b --- /dev/null +++ b/sdk/tests/test_structured_logging.py @@ -0,0 +1,239 @@ +"""Tests for structured logging functionality.""" + +import json +from pathlib import Path + +import pytest +from loguru import logger as log +from spectrumx import utils +from spectrumx.utils import LOG_CATEGORY_LOG +from spectrumx.utils import LogContext +from spectrumx.utils import enable_structured_logging +from spectrumx.utils import reset_structured_logging +from spectrumx.utils import set_persistent_log_context + +MIN_LOG_LINES = 2 + + +@pytest.fixture(autouse=True) +def reset_logging_state(): + """Reset structured logging state before and after each test.""" + reset_structured_logging() + yield + reset_structured_logging() + + +def test_jsonl_file_creation_default_xdg_path(tmp_path: Path, monkeypatch) -> None: + """Test JSONL file creation at default XDG path.""" + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "xdg_state")) + + enable_structured_logging() + log.info("test message for default path") + + current_log_path = utils._current_log_path # noqa: SLF001 + assert current_log_path is not None + assert current_log_path.exists() + assert current_log_path.suffix == ".jsonl" + assert "spectrumx" in current_log_path.parts + assert "logs" in current_log_path.parts + + +def test_jsonl_file_creation_custom_path(tmp_path: Path) -> None: + """Test JSONL file creation at custom path.""" + custom_path = tmp_path / "custom" / "my_log.jsonl" + + enable_structured_logging(log_path=custom_path) + log.info("test message for custom path") + + current_log_path = utils._current_log_path # noqa: SLF001 + assert current_log_path == custom_path + assert custom_path.exists() + lines = custom_path.read_text().strip().split("\n") + assert len(lines) >= 1 + entry = json.loads(lines[-1]) + assert entry["msg"] == "test message for custom path" + + +def test_boot_message_emission(tmp_path: Path) -> None: + """Test boot message is emitted every time structured logging is enabled.""" + log_file = tmp_path / "boot_test.jsonl" + + enable_structured_logging(log_path=log_file) + log.info("trigger boot message") + + lines = log_file.read_text().strip().split("\n") + assert len(lines) >= MIN_LOG_LINES # boot message + at least one log line + + boot_entry = json.loads(lines[0]) + assert boot_entry["msg"] == "system info" + assert boot_entry["lvl"] == "INFO" + assert boot_entry["cat"] == LOG_CATEGORY_LOG + assert "sdk_version" in boot_entry + assert "os" in boot_entry + assert "python" in boot_entry + assert "log_file" in boot_entry + + +def test_core_fields_present_on_every_log_line(tmp_path: Path) -> None: + """Test core fields (ts, pid, lvl, cat, msg) present on every log line.""" + log_file = tmp_path / "core_fields.jsonl" + + enable_structured_logging(log_path=log_file) + log.info("core fields test") + log.warning("another core fields test") + + lines = log_file.read_text().strip().split("\n") + # First line is boot message, rest are log lines + for line in lines[1:]: + entry = json.loads(line) + assert "ts" in entry + assert "pid" in entry + assert "lvl" in entry + assert "cat" in entry + assert "msg" in entry + + +def test_contextual_fields_only_when_set(tmp_path: Path) -> None: + """Test contextual fields only present when explicitly set.""" + log_file = tmp_path / "contextual_fields.jsonl" + + enable_structured_logging(log_path=log_file) + + # Without any context set + log.info("no context") + + # With persistent context set + set_persistent_log_context(api_key_prefix="test_key", timeout=120) + log.info("with persistent context") + + # With scoped context + with LogContext(upload_id="abc123", upload_dir="/some/dir"): + log.info("with scoped context") + + lines = log_file.read_text().strip().split("\n") + # lines[0] = boot, lines[1] = no context, lines[2] = persistent, lines[3] = scoped + no_ctx_entry = json.loads(lines[1]) + persistent_entry = json.loads(lines[2]) + scoped_entry = json.loads(lines[3]) + + # No context line: no upload_id, no upload_dir, no api_key_prefix + assert "upload_id" not in no_ctx_entry + assert "upload_dir" not in no_ctx_entry + assert "api_key_prefix" not in no_ctx_entry + + # Persistent context line: has api_key_prefix and timeout + assert persistent_entry.get("api_key_prefix") == "test_key" + assert persistent_entry.get("timeout") == 120 # noqa: PLR2004 + + # Scoped context line: has upload_id and upload_dir + assert scoped_entry.get("upload_id") == "abc123" + assert scoped_entry.get("upload_dir") == "/some/dir" + # Also has persistent context + assert scoped_entry.get("api_key_prefix") == "test_key" + + +def test_category_binding_via_log_bind(tmp_path: Path) -> None: + """Test category binding via log.bind().""" + log_file = tmp_path / "category_bind.jsonl" + + enable_structured_logging(log_path=log_file) + log.bind(cat="network").info("network message") + log.bind(cat="auth").warning("auth message") + log.info("default category message") # no bind -> default "log" + + lines = log_file.read_text().strip().split("\n") + network_entry = json.loads(lines[1]) + auth_entry = json.loads(lines[2]) + default_entry = json.loads(lines[3]) + + assert network_entry["cat"] == "network" + assert auth_entry["cat"] == "auth" + assert default_entry["cat"] == LOG_CATEGORY_LOG + + +def test_scoped_context_cleared_on_exit(tmp_path: Path) -> None: + """Test scoped context is cleared on exit (no leakage between operations).""" + log_file = tmp_path / "scoped_clear.jsonl" + + enable_structured_logging(log_path=log_file) + + with LogContext(upload_id="scoped-id-1"): + log.info("inside scope 1") + + # After scope exit, upload_id should not leak + log.info("outside scope") + + lines = log_file.read_text().strip().split("\n") + inside_entry = json.loads(lines[1]) + outside_entry = json.loads(lines[2]) + + assert inside_entry.get("upload_id") == "scoped-id-1" + assert "upload_id" not in outside_entry + + +def test_persistent_context_persists(tmp_path: Path) -> None: + """Test persistent context persists across operations.""" + log_file = tmp_path / "persistent_ctx.jsonl" + + enable_structured_logging(log_path=log_file) + set_persistent_log_context(api_key_prefix="persist_key", timeout=999) + + log.info("msg 1") + log.warning("msg 2") + + lines = log_file.read_text().strip().split("\n") + msg1 = json.loads(lines[1]) + msg2 = json.loads(lines[2]) + + assert msg1.get("api_key_prefix") == "persist_key" + assert msg1.get("timeout") == 999 # noqa: PLR2004 + assert msg2.get("api_key_prefix") == "persist_key" + assert msg2.get("timeout") == 999 # noqa: PLR2004 + + +def test_enable_structured_logging_override_changes_path(tmp_path: Path) -> None: + """Test calling enable_structured_logging again changes the log path.""" + path_a = tmp_path / "log_a.jsonl" + path_b = tmp_path / "log_b.jsonl" + + enable_structured_logging(log_path=path_a) + log.info("written to path A") + + assert utils._current_log_path == path_a # noqa: SLF001 + assert path_a.exists() + + # Re-enable with new path — boot message emitted again on new file + enable_structured_logging(log_path=path_b) + log.info("written to path B") + + assert utils._current_log_path == path_b # noqa: SLF001 + assert path_b.exists() + + # Path A should not have the second message + lines_a = path_a.read_text().strip().split("\n") + log_entries_a = [json.loads(line) for line in lines_a] + msgs_a = [e["msg"] for e in log_entries_a if e["msg"] != "system info"] + assert "written to path B" not in msgs_a + + # Path B should have the second message and its own boot message + lines_b = path_b.read_text().strip().split("\n") + boot_b = json.loads(lines_b[0]) + assert boot_b["msg"] == "system info" + msgs_b = [json.loads(line)["msg"] for line in lines_b] + assert "written to path B" in msgs_b + + +def test_reset_structured_logging_clears_state() -> None: + """Test reset_structured_logging clears all state.""" + enable_structured_logging() + set_persistent_log_context(api_key_prefix="should_be_cleared") + + with LogContext(upload_id="should_be_cleared"): + pass + + reset_structured_logging() + + assert utils._structured_sink_id is None # noqa: SLF001 + assert utils._current_log_path is None # noqa: SLF001 + assert utils._persistent_context == {} # noqa: SLF001 + assert utils._log_context.get() is None # noqa: SLF001 diff --git a/sdk/tests/test_uploads_persistence.py b/sdk/tests/test_uploads_persistence.py index fbf4f2810..7ed6b7fe8 100644 --- a/sdk/tests/test_uploads_persistence.py +++ b/sdk/tests/test_uploads_persistence.py @@ -268,6 +268,7 @@ async def test_mark_completed_persists_file( file_path.write_text("uploaded content", encoding="utf-8") file_model = MagicMock(spec=File) + file_model.name = "uploaded_file.txt" file_model.local_path = file_path file_model.size = 100 file_model.compute_sum_blake3.return_value = "persisted_checksum" diff --git a/sdk/tests/test_uploads_workflow.py b/sdk/tests/test_uploads_workflow.py index a65292f6b..ad1232b58 100644 --- a/sdk/tests/test_uploads_workflow.py +++ b/sdk/tests/test_uploads_workflow.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import contextlib import uuid from datetime import UTC from datetime import datetime @@ -22,8 +23,11 @@ from spectrumx.errors import UploadError from spectrumx.models.files.file import File from spectrumx.ops import files as file_ops +from tqdm import tqdm if TYPE_CHECKING: + from collections.abc import Callable + from spectrumx.client import Client # ruff: noqa: SLF001 @@ -823,7 +827,37 @@ def is_valid_side_effect(file_path: Path, **kwargs: bool) -> tuple[bool, list[st assert len(discovered) == 1 assert len(workload.fq_skipped) == 1 - assert workload.fq_skipped[0].path == root / "invalid.bin" + + +@pytest.mark.anyio +async def test_periodic_progress_logger_emits_logs( + upload_workload: UploadWorkload, client: Client +) -> None: + """Test that the periodic progress logger emits logs every N seconds.""" + # Set a very short period to avoid waiting + upload_workload.progress_log_period_secs = 0.01 + + with patch("spectrumx.api.uploads.log") as mock_log: + # log.bind(...) returns a logger instance, so we mock that return value + mock_log.bind.return_value = mock_log + + # Start the logger in a task + task = asyncio.create_task(upload_workload._periodic_progress_logger()) + + # Wait for a few log intervals + await asyncio.sleep(0.05) + + # Cancel and cleanup + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + # Verify that the logger's info method was called + assert mock_log.info.called, "Expected progress logger to emit at least one log" + + # Verify the log content contains "Upload progress" + first_call_args = mock_log.info.call_args[0] + assert "Upload progress" in first_call_args[0] def test_remove_from_buffer_handles_missing_file( @@ -842,3 +876,105 @@ def test_remove_from_buffer_removes_file( UploadWorkload._remove_from_buffer(mock_file, upload_workload.fq_pending) assert mock_file not in upload_workload.fq_pending + + +@pytest.mark.anyio +async def test_upload_next_file_credits_unstreamed_bytes( + upload_workload: UploadWorkload, mock_file: File +) -> None: + """Upload progress should reach file size when content is not fully streamed.""" + mock_file.size = 1000 + await upload_workload._register_discovered_file(mock_file) + prog_bar = MagicMock() + upload_workload._prog_uploaded_bytes = prog_bar + + def fake_upload_file(*, progress_callback=None, **kwargs): + if progress_callback is not None: + progress_callback(250) + return mock_file + + with patch.object( + upload_workload.client._sds_files, + "upload_file", + side_effect=fake_upload_file, + ): + result = await upload_workload._upload_next_file() + + assert result + prog_bar.update.assert_any_call(250) + prog_bar.update.assert_any_call(750) + + +@pytest.mark.anyio +async def test_upload_next_file_credits_metadata_only_upload( + upload_workload: UploadWorkload, mock_file: File +) -> None: + """Upload progress should reach file size for metadata-only uploads.""" + mock_file.size = 1000 + await upload_workload._register_discovered_file(mock_file) + prog_bar = MagicMock() + upload_workload._prog_uploaded_bytes = prog_bar + + with patch.object( + upload_workload.client._sds_files, + "upload_file", + return_value=mock_file, + ): + result = await upload_workload._upload_next_file() + + assert result + prog_bar.update.assert_called_once_with(1000) + + +@pytest.mark.anyio +async def test_concurrent_upload_progress_bar_byte_count( + tmp_path: Path, client: Client +) -> None: + """Concurrent upload workers must not lose byte-count updates on the shared bar.""" + root = tmp_path / "upload_root" + root.mkdir() + + num_files = 4 + chunk_size = 50_000 + chunks_per_file = 4 + file_size = chunk_size * chunks_per_file + files = [ + _create_mock_file(name=f"file_{i}.bin", size=file_size) + for i in range(num_files) + ] + + workload = UploadWorkload( + client=client, + local_root=root, + sds_path=PurePosixPath("/"), + max_concurrent_uploads=num_files, + verbose=True, + ) + for file_obj in files: + await workload._register_discovered_file(file_obj) + + def mock_upload_file( + *, + client: Client, + local_file: File, + sds_path: PurePosixPath, + progress_callback: Callable[[int], None] | None = None, + ) -> File: + del client, sds_path + if progress_callback is not None: + for _ in range(chunks_per_file): + progress_callback(chunk_size) + return local_file + + with patch.object( + workload.client._sds_files, + "upload_file", + side_effect=mock_upload_file, + ): + await workload._execute_uploads() + + expected_bytes = file_size * num_files + assert isinstance(workload._prog_uploaded_bytes, tqdm) + assert workload._prog_uploaded_bytes.n == expected_bytes, ( + "Expected progress bar byte count to match total uploaded bytes" + ) diff --git a/sdk/tests/test_utils.py b/sdk/tests/test_utils.py index 43ba73990..0b6e262a7 100644 --- a/sdk/tests/test_utils.py +++ b/sdk/tests/test_utils.py @@ -3,6 +3,7 @@ import warnings from contextlib import contextmanager from pathlib import Path +from unittest.mock import MagicMock import pytest import urllib3 @@ -151,3 +152,41 @@ def test_into_bool_falsey() -> None: ] for falsey in falsey_values: assert utils.into_human_bool(falsey) is False, f"{falsey} failed" + + +def test_credit_unstreamed_file_bytes_credits_remaining() -> None: + """Unstreamed bytes should be credited to the progress bar and counter.""" + file_size = 1000 + bytes_streamed = 400 + credited_bytes = file_size - bytes_streamed + prog_bar = MagicMock() + bytes_accounted = [bytes_streamed] + + credited = utils.credit_unstreamed_file_bytes( + file_size=file_size, + bytes_streamed=bytes_streamed, + prog_bar=prog_bar, + bytes_accounted=bytes_accounted, + ) + + assert credited == credited_bytes + assert bytes_accounted[0] == file_size + prog_bar.update.assert_called_once_with(credited_bytes) + + +def test_credit_unstreamed_file_bytes_noop_when_fully_streamed() -> None: + """No credit is applied when all bytes were already streamed.""" + file_size = 1000 + prog_bar = MagicMock() + bytes_accounted = [file_size] + + credited = utils.credit_unstreamed_file_bytes( + file_size=file_size, + bytes_streamed=file_size, + prog_bar=prog_bar, + bytes_accounted=bytes_accounted, + ) + + assert credited == 0 + assert bytes_accounted[0] == file_size + prog_bar.update.assert_not_called()