diff --git a/.python-version b/.python-version new file mode 100644 index 000000000..3e388a4ac --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.13.2 diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..16dbd28c2 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,7 @@ +# Agent Workflow Notes + +## Mandatory checklist discipline +- Before starting implementation work, review `TODO.md` and select the items being addressed. +- After every code change, immediately update `TODO.md` by checking completed items and adjusting status text when needed. +- Do not finish a coding task without reconciling `TODO.md` to match the actual repository state. + diff --git a/CHANGELOG.md b/CHANGELOG.md index 19be64c73..9e40fce2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,18 @@ ## Unreleased +## 1.33.0 (2026-03-03) + +- feat: add resilient stateful engine mode with SQLite task/checkpoint persistence and resumable leases (`--state-db`) +- feat: add unified retry/backoff configuration for metadata and downloads (`--max-retries`, backoff controls, `--respect-retry-after`) +- feat: add bounded adaptive download concurrency and chunked streaming controls (`--download-workers`, `--download-chunk-bytes`) +- feat: add download integrity verification controls (`--verify-size`, `--verify-checksum`) and URL-refresh retry path for expired download URLs +- feat: add machine-readable run outputs and observability options (`--log-format json`, `--metrics-json`) +- feat: add state DB maintenance controls (`--state-db-prune-completed-days`, `--state-db-vacuum`) +- feat: add explicit mode contract, exit semantics, graceful cancellation/requeue behavior, and repeated-throttling alerting +- docs: add architecture note, migration/concurrency/troubleshooting guides, and benchmark docs for workers/chunk sizing +- test: extend integration coverage for checkpoint resume, range edge cases, shutdown behavior, mode parity, and URL refresh flows + ## 1.32.2 (2025-09-01) - fix: HTTP response content not captured for authentication and non-streaming requests [#1240](https://github.com/icloud-photos-downloader/icloud_photos_downloader/issues/1240) diff --git a/Improving iCloud Photos Downloader resilience and efficiency at scale.pdf b/Improving iCloud Photos Downloader resilience and efficiency at scale.pdf new file mode 100644 index 000000000..c0f8d29c1 Binary files /dev/null and b/Improving iCloud Photos Downloader resilience and efficiency at scale.pdf differ diff --git a/README.md b/README.md index c4d0c20fe..d7cc98b43 100644 --- a/README.md +++ b/README.md @@ -1,16 +1,48 @@ -# !!!! [Looking for MAINTAINER for this project](https://github.com/icloud-photos-downloader/icloud_photos_downloader/issues/1305) !!!! - # iCloud Photos Downloader [![Quality Checks](https://github.com/icloud-photos-downloader/icloud_photos_downloader/workflows/Quality%20Checks/badge.svg)](https://github.com/icloud-photos-downloader/icloud_photos_downloader/actions/workflows/quality-checks.yml) [![Build and Package](https://github.com/icloud-photos-downloader/icloud_photos_downloader/workflows/Produce%20Artifacts/badge.svg)](https://github.com/icloud-photos-downloader/icloud_photos_downloader/actions/workflows/produce-artifacts.yml) [![MIT License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE) -- A command-line tool to download all your iCloud photos. +A command-line tool to download all your iCloud photos. + - Works on Linux, Windows, and macOS; laptop, desktop, and NAS - Available as an executable for direct downloading and through package managers/ecosystems ([Docker](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#docker), [PyPI](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#pypi), [AUR](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#aur), [npm](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#npm)) -- Developed and maintained by volunteers (we are always looking for [help](CONTRIBUTING.md)). +- Developed and maintained by volunteers (we are always looking for [help](CONTRIBUTING.md)) + +See [Documentation](https://icloud-photos-downloader.github.io/icloud_photos_downloader/) for more details. Also, check [Issues](https://github.com/icloud-photos-downloader/icloud_photos_downloader/issues). -See [Documentation](https://icloud-photos-downloader.github.io/icloud_photos_downloader/) for more details. Also, check [Issues](https://github.com/icloud-photos-downloader/icloud_photos_downloader/issues) +> [!IMPORTANT] +> The project is currently looking for a maintainer. See [issue #1305](https://github.com/icloud-photos-downloader/icloud_photos_downloader/issues/1305). We aim to release new versions once a week (Friday), if there is something worth delivering. +## New in 1.33.0 (Resilient Engine) + +- Unified retry + exponential backoff across metadata and downloads (`429`/`503` aware) +- Adaptive throttling cooldown and bounded download workers (`--download-workers`) +- Optional persistent SQLite task/checkpoint state (`--state-db`) for resumable long runs +- Configurable streaming chunk size (`--download-chunk-bytes`) +- Optional integrity verification (`--verify-size`, `--verify-checksum`) +- JSON logging and run metrics export (`--log-format json`, `--metrics-json`) +- State DB maintenance controls (`--state-db-prune-completed-days`, `--state-db-vacuum`) + +## Engine Modes + +| Mode | Best For | +| --- | --- | +| Classic (stateless) | Small libraries and simple one-off runs | +| Stateful engine | Large libraries and long-running/resumable jobs | + +Stateful mode is enabled with `--state-db`. It stores per-asset tasks/checkpoints, supports deterministic resume after interruption, and requeues stale in-progress tasks safely on restart. + +Backward compatibility: if `--state-db` is not used, `icloudpd` behaves like previous stateless versions (filesystem skip-based behavior). + +### Key Performance and Resilience Flags + +- `--state-db [PATH]` +- `--download-workers N` +- `--max-retries N` +- `--download-chunk-bytes N` +- `--verify-size` / `--verify-checksum` +- `--no-remote-count` + ## iCloud Prerequisites To make iCloud Photo Downloader work, ensure the iCloud account is configured with the following settings, otherwise Apple Servers will return an ACCESS_DENIED error: @@ -22,7 +54,7 @@ To make iCloud Photo Downloader work, ensure the iCloud account is configured wi ## Install and Run There are three ways to run `icloudpd`: -1. Download executable for your platform from the GitHub [Release](https://github.com/icloud-photos-downloader/icloud_photos_downloader/releases/tag/v1.32.2) and run it +1. Download executable for your platform from the GitHub [Release](https://github.com/icloud-photos-downloader/icloud_photos_downloader/releases/tag/v1.33.0) and run it 1. Use package manager to install, update, and, in some cases, run ([Docker](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#docker), [PyPI](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#pypi), [AUR](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#aur), [npm](https://icloud-photos-downloader.github.io/icloud_photos_downloader/install.html#npm)) 1. Build and run from the source diff --git a/TODO.md b/TODO.md new file mode 100644 index 000000000..af9bf3c64 --- /dev/null +++ b/TODO.md @@ -0,0 +1,180 @@ +# iCloud Photos Downloader Improvement Checklist + +Last updated: 2026-03-03 + +Use this as the source of truth for implementation progress. Every code change should update this file. + +## 0. Project hygiene and tracking +- [x] Add/update architecture note describing current pipeline and target pipeline. +- [x] Keep this checklist aligned with actual implemented code and tests. +- [ ] For each completed task, reference the related PR/commit in this file. +- [x] Keep changelog/release notes in sync when user-facing flags/behavior change. +- [x] Ensure local development/testing uses Python 3.13 in `.venv` to match project constraints. + +## 1. Unified retry and backoff (metadata + downloads) +### 1.1 Policy and configuration +- [x] Define one retry policy module shared by metadata calls and file downloads. +- [x] Add CLI option: `--max-retries` (default target: 6). +- [x] Add CLI option: `--backoff-base-seconds`. +- [x] Add CLI option: `--backoff-max-seconds`. +- [x] Add CLI option: `--respect-retry-after/--no-respect-retry-after`. +- [x] Add CLI option: `--throttle-cooldown-seconds`. +- [x] Ensure defaults preserve safe behavior for existing users. + +### 1.2 Error classification +- [x] Classify fatal auth/config errors as no-retry (invalid creds, MFA unavailable, ADP/web-disabled). +- [x] Classify session-invalid errors as re-auth-then-retry. +- [x] Classify transient errors as retryable (429, 503, timeouts, connection resets, throttling-like denials). +- [x] Centralize retry decision logging (attempt, reason, next delay). + +### 1.3 Integration points +- [x] Apply shared retry policy to album/asset enumeration calls. +- [x] Apply shared retry policy to download calls. +- [x] Remove/replace duplicated ad-hoc retry loops in existing code paths. +- [x] Add jitter to exponential backoff. +- [x] Honor `Retry-After` when present on retryable responses. + +### 1.4 Verification +- [x] Unit tests for retry classifier. +- [x] Unit tests for backoff math and jitter bounds. +- [x] Unit tests for `Retry-After` handling. +- [x] Integration tests: metadata retry behavior under simulated 429/503. +- [x] Integration tests: download retry behavior under simulated 429/503/reset. + +## 2. Persistent state DB and resumable task queue +### 2.1 Data model +- [x] Add `--state-db` option (or equivalent path option) with sensible default. +- [x] Create DB initialization/migration path. +- [x] Create `assets` table. +- [x] Create `tasks` table with status/attempt/error fields. +- [x] Create `checkpoints` table for pagination progress. +- [x] Add indexes for task leasing and status filtering. + +### 2.2 Enumeration persistence +- [x] Persist enumerated assets in batches. +- [x] Persist tasks per asset version. +- [x] Save checkpoint every page (or configurable page interval). +- [x] Resume enumeration from checkpoint after restart. + +### 2.3 Worker/task lifecycle +- [x] Add task states: `pending`, `in_progress`, `done`, `failed`. +- [x] Add lease timestamp/owner for `in_progress`. +- [x] Requeue stale leased tasks on startup. +- [x] Track per-task attempts and last error. + +### 2.4 Verification +- [x] Unit tests for DB schema creation and migrations. +- [x] Unit tests for lease/requeue behavior. +- [x] Integration test: crash mid-run and resume without redoing completed tasks. +- [x] Integration test: checkpoint resume after partial enumeration. + +### 2.5 URL freshness +- [x] Detect expired/invalid persisted download URLs and refresh asset version metadata. +- [x] Add task/state marker for URL refresh path (e.g., `needs_url_refresh`) and retry flow. + +## 3. Bounded adaptive concurrency +### 3.1 CLI and defaults +- [x] Add `--download-workers` option (default target: 4). +- [x] Keep metadata enumeration single-threaded by default. +- [x] Document deprecation relationship with `--threads-num`. + +### 3.2 Limiting and adaptation +- [x] Implement shared account-level limiter for download workers. +- [x] Separate metadata and download request budgets (if needed by code design). +- [x] Implement AIMD or equivalent adaptive reduction on throttling events. +- [x] Add global cool-down behavior when repeated throttle signals occur. + +### 3.3 Session/cookie safety +- [x] Audit all session/cookie writes under concurrent access. +- [x] Add locking or redesign to avoid concurrent write races. +- [x] Ensure no cookie/session corruption under multithreaded runs. + +### 3.4 Verification +- [x] Unit tests for limiter/token bucket behavior. +- [x] Concurrency tests for session persistence safety. +- [x] Integration tests for worker pool drain/stop/restart behavior. +- [x] Benchmark runs at workers = 1, 2, 4, 8 and record throughput + error rate. + +## 4. Download efficiency and integrity +### 4.1 Throughput improvements +- [x] Add `--download-chunk-bytes` option (default target: 262144). +- [x] Replace fixed 1 KiB streaming chunk with configurable larger chunk. +- [x] Verify memory usage remains bounded by worker count and chunk size. +- [x] Benchmark chunk-size/verification combinations for throughput vs CPU tradeoff. + +### 4.2 Integrity checks +- [x] Add `--verify-size/--no-verify-size` option. +- [x] Add `--verify-checksum/--no-verify-checksum` option. +- [x] Validate downloaded file size against expected metadata. +- [x] Implement optional checksum validation strategy. +- [x] Store local checksum/result in state DB when enabled. + +### 4.3 Range resume hardening +- [x] Keep `.part` resume behavior with `Range` requests. +- [x] Detect non-`206` response when resuming and safely restart partial file. +- [x] Add corruption-safe handling for mismatched range behavior. + +### 4.4 Verification +- [x] Unit tests for chunk-size configuration and defaults. +- [x] Unit tests for size verification success/failure. +- [x] Unit tests for checksum verification success/failure. +- [x] Integration tests for resume with partial files and range edge cases. + +## 5. Request volume and enumeration efficiency +- [x] Add `--album-page-size` option (target range: 50-500). +- [x] Add `--no-remote-count` option to skip expensive album count calls. +- [x] Reduce redundant metadata queries where possible. +- [x] Add/align chunked date-based run options (`since/until added date` behavior). +- [x] Document clear behavior differences between added-date and created-date usage. +- [x] Add tests for new pagination and remote-count toggles. + +## 6. Observability and operations +### 6.1 Logging +- [x] Add structured JSON log mode. +- [x] Include stable fields (`run_id`, `asset_id`, `attempt`, `http_status`, etc.). +- [x] Ensure sensitive data redaction remains enforced. + +### 6.2 Metrics and health +- [x] Add metrics endpoint or export path (if compatible with current stack). +- [x] Track throughput, retries, throttle events, queue depth, success gap. +- [x] Add low-disk-space warning/error classification. +- [x] Provide JSON stats snapshot output suitable for GUI wrappers (`--metrics-json`). + +### 6.3 Alerts and notifications +- [x] Add alert condition for repeated throttling. +- [x] Keep MFA expiry notification path working with new engine. +- [x] Add docs for recommended operational thresholds. + +## 7. Documentation and migration +- [x] Update CLI reference docs for all new options. +- [x] Add migration guide: stateless mode vs stateful mode. +- [x] Document compatibility and unchanged default behavior. +- [x] Document concurrency limitations and safe defaults. +- [x] Add troubleshooting guide for throttling/session issues. + +## 9. Runtime Semantics and Operability Hardening +### 9.1 Mode contract +- [x] Define explicit legacy/stateless mode contract (no DB required, filesystem skip semantics). +- [x] Define explicit stateful engine mode contract (resume guarantees, task-state semantics). +- [x] Add integration tests asserting mode-specific behavior and parity expectations. + +### 9.2 Exit and summary semantics +- [x] Define process exit code contract (success, partial success, fatal auth/config, cancelled, stalled). +- [x] Emit machine-readable end-of-run summary with totals/failures/error location hints. + +### 9.3 Cancellation and shutdown +- [x] Handle SIGINT/SIGTERM with graceful stop (drain or safe requeue of in-flight work). +- [x] Ensure clean shutdown is distinguishable from crash and restart behavior is deterministic. + +### 9.4 State DB growth and retention +- [x] Add DB retention/pruning policy (completed task cleanup / capped error history). +- [x] Document and/or automate WAL checkpointing and vacuum guidance. + +## 8. Final validation before release +- [ ] Full test suite passes. +- [x] New tests added for each new subsystem. +- [x] Lint/type checks pass. +- [ ] Manual end-to-end dry run on small sample library. +- [ ] Manual end-to-end run with injected transient failures. +- [x] Confirm no regressions in naming/dedup/folder behavior. +- [x] Confirm watch mode behavior is unchanged unless explicitly modified. diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 000000000..efe1cda93 --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,42 @@ +# Engine Architecture + +This note describes the current downloader execution pipeline and the target resilient pipeline now implemented. + +## Modes + +- `legacy_stateless`: + - No state DB required. + - Filesystem existence checks drive skip/retry behavior. + - Preserves legacy CLI expectations. +- `stateful_engine`: + - Uses SQLite state DB (`--state-db`) for assets, tasks, and checkpoints. + - Supports deterministic resume with leased task recovery. + +## Pipeline Stages + +1. Authenticate and initialize per-user run context (`run_id`, retry/limiter/metrics). +2. Enumerate remote assets (single-threaded) and persist checkpoints/tasks in stateful mode. +3. Download via bounded worker pool with adaptive limiter. +4. Apply unified retry/backoff policy for metadata and downloads (with jitter and `Retry-After`). +5. Verify integrity (size and optional checksum). +6. Persist task outcomes and emit end-of-run summary (machine-readable + human logs). + +## Resilience Guarantees + +- Shared retry classifier for transient vs fatal errors. +- Re-auth on session-invalid failures. +- URL freshness path for expired download URLs (`401`/`403`/`410`) with one metadata refresh retry. +- Graceful cancellation (`SIGINT`/`SIGTERM`) with safe requeue semantics. +- Restart safety for stale leases and pagination checkpoints. + +## Throughput and Safety Controls + +- Configurable chunked streaming (`--download-chunk-bytes`) with bounded memory behavior. +- Adaptive concurrency (`--download-workers`) with throttle backoff/cooldown. +- Optional remote-count skip and page-size tuning for lower API pressure. + +## Operability + +- Structured JSON logs (`--log-format json`). +- JSON metrics snapshot (`--metrics-json`) for wrappers/GUI integration. +- State DB maintenance options (`--state-db-prune-completed-days`, `--state-db-vacuum`). diff --git a/docs/benchmark_download_chunks.md b/docs/benchmark_download_chunks.md new file mode 100644 index 000000000..a21e5a77b --- /dev/null +++ b/docs/benchmark_download_chunks.md @@ -0,0 +1,42 @@ +# Download Chunk Benchmark (Synthetic) + +Date: 2026-03-03 + +Command used: + +```bash +.venv/bin/python scripts/benchmark_download_chunks.py \ + --workers 4 \ + --size-mib 8 \ + --iterations 2 \ + --chunk-bytes 65536 262144 1048576 +``` + +## Throughput vs CPU + +| Chunk bytes | Verify checksum | Avg Throughput (MiB/s) | Avg CPU seconds | +|---|---|---:|---:| +| 65536 | no | 277.90 | 0.0505 | +| 65536 | yes | 137.00 | 0.1496 | +| 262144 | no | 663.89 | 0.0297 | +| 262144 | yes | 207.23 | 0.1291 | +| 1048576 | no | 704.66 | 0.0293 | +| 1048576 | yes | 181.65 | 0.1437 | + +Notes: +- Larger chunks improve throughput significantly vs 64 KiB in this synthetic stream test. +- Enabling checksum verification increases CPU cost and reduces throughput, as expected. +- `262144` and `1048576` are close on CPU cost when checksum is disabled; `262144` remains a good default. + +## Memory boundedness verification + +A dedicated integration test verifies streaming memory remains bounded during large transfers: + +```bash +.venv/bin/python -m pytest \ + tests/test_download_config.py::DownloadConfigTestCase::test_download_response_streaming_memory_is_bounded -q +``` + +Test behavior: +- Streams a 64 MiB response to disk using `--download-chunk-bytes=65536`. +- Asserts peak traced memory stays below 8 MiB (well below transferred bytes), confirming bounded streaming behavior. diff --git a/docs/benchmark_download_workers.md b/docs/benchmark_download_workers.md new file mode 100644 index 000000000..5f8f6c4ca --- /dev/null +++ b/docs/benchmark_download_workers.md @@ -0,0 +1,73 @@ +# Download Worker Benchmark (Synthetic) + +Date: 2026-03-03 + +Command used: + +```bash +.venv/bin/python scripts/benchmark_download_workers.py \ + --workers 1 2 4 8 \ + --tasks 400 \ + --latency-seconds 0.01 \ + --throttle-probability 0.03 \ + --seed 1337 +``` + +Results: + +| Workers | Throughput (tasks/s) | Error Rate | +|---|---:|---:| +| 1 | 92.92 | 0.02 | +| 2 | 157.01 | 0.03 | +| 4 | 289.67 | 0.02 | +| 8 | 364.21 | 0.03 | + +Raw JSON: + +```json +{ + "timestamp_epoch": 1772548722.9070098, + "seed": 1337, + "tasks": 400, + "latency_seconds": 0.01, + "throttle_probability": 0.03, + "results": [ + { + "workers": 1, + "tasks": 400, + "successes": 392, + "errors": 8, + "error_rate": 0.02, + "throughput_tasks_per_sec": 92.91940860752013, + "elapsed_seconds": 4.304805701998703 + }, + { + "workers": 2, + "tasks": 400, + "successes": 388, + "errors": 12, + "error_rate": 0.03, + "throughput_tasks_per_sec": 157.00543983061695, + "elapsed_seconds": 2.5476824270008365 + }, + { + "workers": 4, + "tasks": 400, + "successes": 392, + "errors": 8, + "error_rate": 0.02, + "throughput_tasks_per_sec": 289.6725751729342, + "elapsed_seconds": 1.3808694170002127 + }, + { + "workers": 8, + "tasks": 400, + "successes": 388, + "errors": 12, + "error_rate": 0.03, + "throughput_tasks_per_sec": 364.2110641947837, + "elapsed_seconds": 1.0982642739982111 + } + ] +} +``` diff --git a/docs/concurrency.md b/docs/concurrency.md new file mode 100644 index 000000000..f86758f86 --- /dev/null +++ b/docs/concurrency.md @@ -0,0 +1,26 @@ +# Concurrency and Safe Defaults + +## Current Concurrency Model + +- Metadata enumeration is single-threaded. +- Downloads use bounded worker concurrency via `--download-workers`. +- Account-level adaptive limiting reduces effective concurrency under throttling. + +## Safe Defaults + +- Default `--download-workers` is `4`. +- Start at `2` for constrained networks/NAS devices. +- Increase gradually only after verifying stable error rates. + +## Practical Limits + +- High worker counts can trigger more throttling and retries. +- Very short watch intervals plus high worker counts increase request pressure. +- One process per account/cookie directory remains the safest operational pattern. + +## Tuning Order + +1. Set `--download-workers`. +2. Validate throughput vs retries/throttle events. +3. Adjust `--watch-with-interval` for recurring runs. +4. Keep `--no-remote-count` enabled when operating near throttle limits. diff --git a/docs/conf.py b/docs/conf.py index 6ed1e365d..c4b8c12cf 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -9,7 +9,7 @@ project = "icloudpd" copyright = "2024, Contributors" author = "Contributors" -release = "1.32.2" +release = "1.33.0" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration @@ -19,7 +19,7 @@ templates_path = ["_templates"] exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] -html_title = "icloudpd 1.32.2" +html_title = "icloudpd 1.33.0" language = "en" # language = 'Python' diff --git a/docs/index.md b/docs/index.md index e116c4829..8341ff9fe 100644 --- a/docs/index.md +++ b/docs/index.md @@ -12,7 +12,13 @@ A command-line tool to download your iCloud photos to local storage install authentication naming +architecture mode +migration +concurrency +troubleshooting +benchmark_download_workers +benchmark_download_chunks size raw webui diff --git a/docs/install.md b/docs/install.md index cd3b38124..4945ce3b9 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1,7 +1,7 @@ # Install and Run There are three ways to run `icloudpd`: -1. Download executable for your platform from the GitHub [Release](https://github.com/icloud-photos-downloader/icloud_photos_downloader/releases/tag/v1.32.2) and run it, e.g.: +1. Download executable for your platform from the GitHub [Release](https://github.com/icloud-photos-downloader/icloud_photos_downloader/releases/tag/v1.33.0) and run it, e.g.: ```sh icloudpd --username your@email.address --directory photos --watch-with-interval 3600 @@ -107,13 +107,13 @@ npx --yes icloudpd --directory /data --username my@email.address --watch-with-in Here are the steps to make it work: - Download the binary from GitHub [Releases](https://github.com/icloud-photos-downloader/icloud_photos_downloader/releases) into the desired local folder -- Add the executable flag by running `chmod +x icloudpd-1.32.2-macos-amd64` -- Start it from the terminal: `icloudpd-1.32.2-macos-amd64` +- Add the executable flag by running `chmod +x icloudpd-1.33.0-macos-amd64` +- Start it from the terminal: `icloudpd-1.33.0-macos-amd64` - Apple will tell you that it cannot check for malicious software and refuse to run the app; click "OK" -- Open "System Settings"/"Privacy & Security" and find `icloudpd-1.32.2-macos-amd64` as a blocked app; click "Allow" -- Start `icloudpd-1.32.2-macos-amd64` from the terminal again +- Open "System Settings"/"Privacy & Security" and find `icloudpd-1.33.0-macos-amd64` as a blocked app; click "Allow" +- Start `icloudpd-1.33.0-macos-amd64` from the terminal again - Apple will show another warning; click "Open" -- After that, you can run `icloudpd-1.32.2-macos-amd64 --help` or any other supported command/option +- After that, you can run `icloudpd-1.33.0-macos-amd64 --help` or any other supported command/option ## Error on the First Run diff --git a/docs/migration.md b/docs/migration.md new file mode 100644 index 000000000..6cf53d1a0 --- /dev/null +++ b/docs/migration.md @@ -0,0 +1,37 @@ +# Migration Guide + +This guide describes migration from the legacy stateless behavior to the stateful engine behavior. + +## Stateless vs Stateful + +Stateless mode (default) +: No `--state-db` configured. + +: Behavior: + - No SQLite DB is required. + - Existing local files are skipped using filesystem checks. + - Restarts do not persist task/checkpoint state. + +Stateful mode +: Enabled by `--state-db` (or `--state-db PATH`). + +: Behavior: + - Persists task/checkpoint state in SQLite. + - Supports deterministic resume semantics for long runs. + - Requeues stale in-progress leases on restart. + +## Compatibility and Defaults + +- Default behavior remains stateless unless `--state-db` is explicitly set. +- Existing naming, folder structure, and dedup behavior remain unchanged. +- `--threads-num` remains accepted for compatibility, but download concurrency is controlled by `--download-workers`. +- Watch mode behavior is unchanged unless engine options are explicitly enabled. + +## Recommended Migration Steps + +1. Start with your current command and add `--state-db` only. +2. Verify first run output and resulting local files match expected parity. +3. Optionally add: + - `--state-db-prune-completed-days` for retention + - `--state-db-vacuum` for periodic space reclamation +4. If throttling appears, lower `--download-workers` before other changes. diff --git a/docs/mode.md b/docs/mode.md index 669a866d3..95f5874fe 100644 --- a/docs/mode.md +++ b/docs/mode.md @@ -20,3 +20,23 @@ Move : Download assets from iCloud that are not in the local storage (same as Copy). Then delete assets in iCloud that are in local storage, optionally leaving recent ones in iCloud This mode is selected with [`--keep-icloud-recent-days`](keep-icloud-recent-days-parameter) parameter + +## Engine Modes (Execution Contract) + +`icloudpd` also has two execution-engine modes that control resume/task behavior: + +Legacy / Stateless Engine (`run_mode=legacy_stateless`) +: Default when `--state-db` is not used. + +: Contract: + - No SQLite state DB is required. + - Existing local files are skipped using filesystem checks (`already exists` semantics). + - Restart behavior is best-effort and stateless (no task/checkpoint persistence). + +Stateful Engine (`run_mode=stateful_engine`) +: Enabled when `--state-db` is used. + +: Contract: + - Uses persistent SQLite task/checkpoint state. + - Supports deterministic resume via checkpoint/task state. + - In-progress/stale leases are safely requeued on restart/cancellation. diff --git a/docs/reference.md b/docs/reference.md index 681e71736..3f2318acd 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -219,6 +219,120 @@ This is a list of all options available for the command line interface (CLI) of : If specified, assets downloaded locally will be deleted in iCloud (actually moved to the Recently Deleted album). Deprecated, use [`--keep-icloud-recent-days`](keep-icloud-recent-days-parameter) instead. +(skip-created-before-parameter)= +`--skip-created-before X` / `--skip-created-after X` + +: Filters assets using the original capture time (asset created date). + +(skip-added-before-parameter)= +`--skip-added-before X` / `--skip-added-after X` + +: Filters assets using the iCloud added date (when asset was added/synced to iCloud). + +Date filter behavior: + +- `created` filters are about when the media was taken. +- `added` filters are about when it appeared in iCloud. +- `--recent` and `--until-found` traversal order uses added-date ordering. + +(state-db-parameter)= +`--state-db [PATH]` + +: Enables persistent state DB for resumable runs. If PATH is omitted, the DB is created at `/icloudpd.sqlite`. + +`--state-db-prune-completed-days DAYS` + +: Optional run-end retention policy. Removes `done`/`failed` task rows older than `DAYS`. + +`--state-db-vacuum` + +: Optional run-end full `VACUUM` to reclaim SQLite file space (slower on large DBs). + +State DB maintenance behavior: + +- WAL checkpoint (`PRAGMA wal_checkpoint(PASSIVE)`) is executed automatically at run end when `--state-db` is enabled. +- Use `--state-db-vacuum` periodically (not every run) for large databases. +- If a download URL is expired/denied (`401`/`403`/`410`), `icloudpd` attempts one metadata URL refresh and retry. +- Unrecoverable URL-expiry cases are marked in task state via `needs_url_refresh=1`. + +## Resilience and Engine Options + +```{versionadded} 1.33.0 +``` + +`--log-format {text,json}` + +: Selects text logs or structured JSON logs. + +`--metrics-json PATH` + +: Writes machine-readable end-of-run metrics and summary JSON. + +`--max-retries N` + +: Maximum transient retry attempts across metadata/download paths. + +`--backoff-base-seconds X` + +: Base delay for exponential retry backoff. + +`--backoff-max-seconds X` + +: Maximum delay cap for retry backoff. + +`--respect-retry-after` / `--no-respect-retry-after` + +: Enables/disables honoring server `Retry-After` headers. + +`--throttle-cooldown-seconds X` + +: Minimum cooldown delay used when throttling is detected. + +`--download-workers N` + +: Download worker concurrency limit. Metadata enumeration remains single-threaded. + +`--download-chunk-bytes N` + +: Streaming chunk size for downloads. + +`--verify-size` / `--no-verify-size` + +: Enables/disables downloaded-size validation against metadata. + +`--verify-checksum` / `--no-verify-checksum` + +: Enables/disables checksum validation when checksum metadata is available. + +`--album-page-size N` + +: Album API page size used for enumeration. + +`--no-remote-count` + +: Skips remote count queries to reduce metadata call volume. + +## Operational Thresholds + +- Repeated throttling alert: emitted when a user run records 5 or more throttle events. +- Recommended first responses when alerts appear: + - lower `--download-workers` (for example, from `4` to `2`) + - increase `--watch-with-interval` + - keep `--no-remote-count` enabled to reduce metadata calls + +## Exit Codes + +- `0`: run completed (may be `partial_success` if some downloads failed; inspect run summary/metrics). +- `1`: runtime/authentication/processing error. +- `2`: CLI usage/validation error. +- `130`: cancelled by `SIGINT`/`SIGTERM` with graceful shutdown. + +When `--metrics-json` is enabled, the output includes: + +- `exit_code` +- `status` (`success`, `partial_success`, `cancelled`, `runtime_error`, `cli_error`) +- per-user counters and throughput/error metrics (including `run_mode`: `legacy_stateless` or `stateful_engine`). + ```{note} If remote assets were not downloaded, e.g., because they were already in local storage, they will NOT be deleted in iCloud. ``` diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md new file mode 100644 index 000000000..87a2d80f3 --- /dev/null +++ b/docs/troubleshooting.md @@ -0,0 +1,38 @@ +# Troubleshooting + +## Throttling (429/503, slow progress, repeated retries) + +Symptoms: +- frequent retry logs +- repeated throttle warnings +- reduced throughput over time + +Actions: +1. Lower `--download-workers` (for example, `4` to `2`). +2. Increase `--watch-with-interval` for watch mode runs. +3. Use `--no-remote-count` to reduce metadata request load. +4. Keep retry defaults unless you have a measured reason to tune them. + +## Session/Cookie Issues + +Symptoms: +- repeated re-authentication prompts +- intermittent authentication failures across runs + +Actions: +1. Keep one active process per account/cookie directory. +2. Use separate `--cookie-directory` values for different accounts. +3. If using stateful mode, keep `--state-db` in the same account-scoped directory. + +## Resume Expectations + +- Stateless mode resumes only via filesystem skip checks. +- Stateful mode resumes via persisted task/checkpoint state. +- On clean cancellation (`SIGINT`/`SIGTERM`), in-progress state is safely requeued. +- Expired URL failures in stateful mode are marked with `needs_url_refresh=1` for affected tasks. + +## State DB Size Growth + +If the state DB grows large: +1. Add `--state-db-prune-completed-days` (for example, `30`). +2. Run periodic `--state-db-vacuum` (not on every run). diff --git a/manual-e2e-cookie/christorpgmailcom b/manual-e2e-cookie/christorpgmailcom new file mode 100644 index 000000000..9d64a7a7a --- /dev/null +++ b/manual-e2e-cookie/christorpgmailcom @@ -0,0 +1,4 @@ +#LWP-Cookies-2.0 +Set-Cookie3: dslang="US-EN"; path="/"; domain=".apple.com"; path_spec; secure; discard; HttpOnly=None; version=0 +Set-Cookie3: site=USA; path="/"; domain=".apple.com"; path_spec; secure; discard; HttpOnly=None; version=0 +Set-Cookie3: aasp=3D9348A6598A974AFF2E0B99DDAF957ED72485918522EDDFA7B4835B5155B825511C2BFA51C0A03E94012493855935388B9A26056B4E594B7492F97AC3B7FB9A396CA36ABCC82E6283AE263AEB51424E762D8D41A1A40D94E0C4F8A1375D5D5CDA6A69238EF55FC74417A6F5A44D0891681CDFC47996ED40; path="/"; domain=".idmsa.apple.com"; path_spec; secure; discard; HttpOnly=None; version=0 diff --git a/manual-e2e-cookie/christorpgmailcom.session b/manual-e2e-cookie/christorpgmailcom.session new file mode 100644 index 000000000..fd6e1e242 --- /dev/null +++ b/manual-e2e-cookie/christorpgmailcom.session @@ -0,0 +1 @@ +{"client_id": "auth-10fa11fe-1721-11f1-bca8-181dea8b03bc", "session_id": "3D9348A6598A974AFF2E0B99DDAF957ED72485918522EDDFA7B4835B5155B825511C2BFA51C0A03E94012493855935388B9A26056B4E594B7492F97AC3B7FB9A396CA36ABCC82E6283AE263AEB51424E762D8D41A1A40D94E0C4F8A1375D5D5CDA6A69238EF55FC74417A6F5A44D0891681CDFC47996ED40", "scnt": "AAAA-jNEOTM0OEE2NTk4QTk3NEFGRjJFMEI5OUREQUY5NTdFRDcyNDg1OTE4NTIyRURERkE3QjQ4MzVCNTE1NUI4MjU1MTFDMkJGQTUxQzBBMDNFOTQwMTI0OTM4NTU5MzUzODhCOUEyNjA1NkI0RTU5NEI3NDkyRjk3QUMzQjdGQjlBMzk2Q0EzNkFCQ0M4MkU2MjgzQUUyNjNBRUI1MTQyNEU3NjJEOEQ0MUExQTQwRDk0RTBDNEY4QTEzNzVENUQ1Q0RBNkE2OTIzOEVGNTVGQzc0NDE3QTZGNUE0NEQwODkxNjgxQ0RGQzQ3OTk2RUQ0MHw0AAABnLSrQOXJrFwZYEAlHtl9aqdRdyqxVKX97-izomXIEuzPPncWJIV6uiVF3boHAAJ0B8lE2_qrgG4dfJHAU2MNlid8FrKEUTLCHGaMoKQNcCUH7gGnng"} \ No newline at end of file diff --git a/manual-e2e-cookie/icloudpd.sqlite b/manual-e2e-cookie/icloudpd.sqlite new file mode 100644 index 000000000..f221e763b Binary files /dev/null and b/manual-e2e-cookie/icloudpd.sqlite differ diff --git a/manual-e2e-metrics.json b/manual-e2e-metrics.json new file mode 100644 index 000000000..a2a47e00a --- /dev/null +++ b/manual-e2e-metrics.json @@ -0,0 +1,27 @@ +{ + "exit_code": 1, + "status": "runtime_error", + "users_total": 1, + "users_with_failures": 0, + "users": [ + { + "username": "christorp@gmail.com", + "run_mode": "stateful_engine", + "started_at_epoch": 1772556744.4582224, + "finished_at_epoch": 1772556748.4089096, + "elapsed_seconds": 3.9506871700286865, + "assets_considered": 0, + "downloads_attempted": 0, + "downloads_succeeded": 0, + "downloads_failed": 0, + "success_gap": 0, + "bytes_downloaded": 0, + "throughput_downloads_per_sec": 0.0, + "retries": 0, + "throttle_events": 0, + "queue_depth_max": 0, + "queue_depth_last": 0, + "low_disk_events": 0 + } + ] +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 409e2b3e5..10d0ad64d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ requires = [ build-backend = "setuptools.build_meta" [project] -version="1.32.2" +version="1.33.0" name = "icloudpd" description = "icloudpd is a command-line tool to download photos and videos from iCloud." readme = "README_PYPI.md" @@ -133,4 +133,3 @@ ignore = [ # deprecated typing namespace "UP035", ] - diff --git a/scripts/benchmark_download_chunks.py b/scripts/benchmark_download_chunks.py new file mode 100755 index 000000000..cd16c0a29 --- /dev/null +++ b/scripts/benchmark_download_chunks.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +"""Synthetic benchmark for download chunk size and verification tradeoffs.""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime +import hashlib +import json +import logging +import os +import tempfile +import threading +import time +import tracemalloc +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +from icloudpd import download + + +@dataclass(frozen=True) +class Case: + chunk_bytes: int + verify_size: bool + verify_checksum: bool + + +class FakeResponse: + def __init__(self, payload: bytes, total_bytes: int) -> None: + self._payload = payload + self._total_bytes = total_bytes + self._payload_len = len(payload) + + def iter_content(self, chunk_size: int): + sent = 0 + while sent < self._total_bytes: + remaining = self._total_bytes - sent + take = chunk_size if remaining >= chunk_size else remaining + start = sent % self._payload_len + if start + take <= self._payload_len: + chunk = self._payload[start : start + take] + else: + part1 = self._payload[start:] + left = take - len(part1) + repeats = left // self._payload_len + tail = left % self._payload_len + chunk = part1 + (self._payload * repeats) + self._payload[:tail] + sent += len(chunk) + yield chunk + + +def _expected_content(payload: bytes, total_bytes: int) -> bytes: + return (payload * ((total_bytes // len(payload)) + 1))[:total_bytes] + + +def run_single_download(case: Case, payload: bytes, total_bytes: int, root_dir: str) -> dict: + expected = _expected_content(payload, total_bytes) + expected_size = len(expected) + checksum = hashlib.md5(expected).digest() + + thread_id = threading.get_ident() + base = os.path.join( + root_dir, + f"case-{case.chunk_bytes}-{int(case.verify_size)}-{int(case.verify_checksum)}-{thread_id}-{time.time_ns()}", + ) + temp_path = f"{base}.part" + out_path = f"{base}.bin" + + if os.path.exists(temp_path): + os.remove(temp_path) + if os.path.exists(out_path): + os.remove(out_path) + + response = FakeResponse(payload=payload, total_bytes=total_bytes) + created = datetime.datetime(2026, 3, 3, tzinfo=datetime.timezone.utc) + + start_wall = time.perf_counter() + start_cpu = time.process_time() + download.set_download_chunk_bytes(case.chunk_bytes) + write_ok = download.download_response_to_path( + response, + temp_path, + append_mode=False, + download_path=out_path, + created_date=created, + ) + verify_ok = download.verify_download_integrity( + logger=logging.getLogger("benchmark"), + download_path=out_path, + expected_size=expected_size, + expected_checksum=checksum, + verify_size=case.verify_size, + verify_checksum=case.verify_checksum, + ) + elapsed = time.perf_counter() - start_wall + cpu = time.process_time() - start_cpu + + with contextlib.suppress(OSError): + os.remove(out_path) + + return { + "chunk_bytes": case.chunk_bytes, + "verify_size": case.verify_size, + "verify_checksum": case.verify_checksum, + "ok": bool(write_ok and verify_ok), + "size_bytes": expected_size, + "elapsed_seconds": elapsed, + "cpu_seconds": cpu, + "throughput_mib_per_sec": (expected_size / (1024 * 1024) / elapsed) if elapsed > 0 else 0.0, + } + + +def run_case(case: Case, *, size_mib: int, workers: int, iterations: int) -> dict: + payload = b"icloudpd-benchmark-payload-" + total_bytes = size_mib * 1024 * 1024 + + with tempfile.TemporaryDirectory(prefix="icloudpd-benchmark-chunks-") as temp_dir: + tracemalloc.start() + per_run: list[dict] = [] + for _ in range(iterations): + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = [ + pool.submit(run_single_download, case, payload, total_bytes, temp_dir) + for _ in range(workers) + ] + for future in futures: + per_run.append(future.result()) + _current, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + throughput = [run["throughput_mib_per_sec"] for run in per_run] + cpu = [run["cpu_seconds"] for run in per_run] + + return { + "chunk_bytes": case.chunk_bytes, + "verify_size": case.verify_size, + "verify_checksum": case.verify_checksum, + "workers": workers, + "iterations": iterations, + "size_mib_per_worker": size_mib, + "runs": len(per_run), + "all_ok": all(run["ok"] for run in per_run), + "throughput_mib_per_sec_avg": sum(throughput) / len(throughput), + "cpu_seconds_avg": sum(cpu) / len(cpu), + "tracemalloc_peak_bytes": peak, + "theoretical_stream_buffer_bound_bytes": workers * case.chunk_bytes, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--workers", type=int, default=4) + parser.add_argument("--size-mib", type=int, default=16) + parser.add_argument("--iterations", type=int, default=2) + parser.add_argument("--chunk-bytes", type=int, nargs="+", default=[65536, 262144, 1048576]) + parser.add_argument("--out", type=str, default=None) + args = parser.parse_args() + + cases: list[Case] = [] + for chunk in args.chunk_bytes: + cases.append(Case(chunk_bytes=chunk, verify_size=True, verify_checksum=False)) + cases.append(Case(chunk_bytes=chunk, verify_size=True, verify_checksum=True)) + + results = [ + run_case(case, size_mib=args.size_mib, workers=args.workers, iterations=args.iterations) + for case in cases + ] + + payload = { + "timestamp_epoch": time.time(), + "workers": args.workers, + "size_mib": args.size_mib, + "iterations": args.iterations, + "results": results, + } + + rendered = json.dumps(payload, indent=2) + if args.out: + with open(args.out, "w", encoding="utf-8") as handle: + handle.write(rendered) + print(rendered) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/benchmark_download_workers.py b/scripts/benchmark_download_workers.py new file mode 100644 index 000000000..9f3a14d1f --- /dev/null +++ b/scripts/benchmark_download_workers.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +"""Synthetic benchmark for download-worker limiter settings.""" + +from __future__ import annotations + +import argparse +import json +import random +import time +from concurrent.futures import ThreadPoolExecutor + +from icloudpd.limiter import AdaptiveDownloadLimiter + + +def run_case(workers: int, tasks: int, latency_seconds: float, throttle_probability: float) -> dict: + limiter = AdaptiveDownloadLimiter( + max_workers=workers, + min_workers=1, + cooldown_seconds=max(0.0, latency_seconds * 2), + increase_every=5, + ) + errors = 0 + done = 0 + + def unit_of_work() -> bool: + nonlocal errors + with limiter.slot(timeout=2.0): + time.sleep(latency_seconds) + throttled = random.random() < throttle_probability + if throttled: + limiter.on_throttle() + errors += 1 + return False + limiter.on_success() + return True + + start = time.perf_counter() + with ThreadPoolExecutor(max_workers=max(1, workers * 2)) as pool: + futures = [pool.submit(unit_of_work) for _ in range(tasks)] + for future in futures: + if future.result(): + done += 1 + elapsed = time.perf_counter() - start + + return { + "workers": workers, + "tasks": tasks, + "successes": done, + "errors": errors, + "error_rate": (errors / tasks) if tasks else 0.0, + "throughput_tasks_per_sec": (tasks / elapsed) if elapsed > 0 else 0.0, + "elapsed_seconds": elapsed, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--tasks", type=int, default=200) + parser.add_argument("--latency-seconds", type=float, default=0.01) + parser.add_argument("--throttle-probability", type=float, default=0.03) + parser.add_argument("--seed", type=int, default=1337) + parser.add_argument("--workers", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument("--out", type=str, default=None) + args = parser.parse_args() + + random.seed(args.seed) + results = [ + run_case( + workers=workers, + tasks=args.tasks, + latency_seconds=args.latency_seconds, + throttle_probability=args.throttle_probability, + ) + for workers in args.workers + ] + payload = { + "timestamp_epoch": time.time(), + "seed": args.seed, + "tasks": args.tasks, + "latency_seconds": args.latency_seconds, + "throttle_probability": args.throttle_probability, + "results": results, + } + rendered = json.dumps(payload, indent=2) + if args.out: + with open(args.out, "w", encoding="utf-8") as f: + f.write(rendered) + print(rendered) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/foundation/__init__.py b/src/foundation/__init__.py index 226634e75..80fd680d9 100644 --- a/src/foundation/__init__.py +++ b/src/foundation/__init__.py @@ -27,7 +27,7 @@ class VersionInfo(NamedTuple): # will be updated by CI version_info = VersionInfo( - version="0.0.1", + version="1.33.0", commit_sha="abcdefgh", commit_timestamp=1234567890, ) diff --git a/src/icloudpd/authentication.py b/src/icloudpd/authentication.py index d6a7d1568..e249d04db 100644 --- a/src/icloudpd/authentication.py +++ b/src/icloudpd/authentication.py @@ -248,6 +248,8 @@ def request_2fa_web( # wait for input while True: + if status_exchange.get_progress().cancel: + raise KeyboardInterrupt() status = status_exchange.get_status() if status == Status.NEED_MFA: time.sleep(1) diff --git a/src/icloudpd/base.py b/src/icloudpd/base.py index 2c9fe24dc..05298526e 100644 --- a/src/icloudpd/base.py +++ b/src/icloudpd/base.py @@ -7,15 +7,18 @@ import json import logging import os +import re +import signal import subprocess import sys import time import typing import urllib +import uuid from functools import partial, singledispatch from logging import Logger from multiprocessing import freeze_support -from threading import Thread +from threading import Event, Thread, current_thread, main_thread from typing import ( Any, Callable, @@ -45,15 +48,42 @@ from icloudpd.counter import Counter from icloudpd.email_notifications import send_2sa_notification from icloudpd.filename_policies import build_filename_with_policies, create_filename_builder +from icloudpd.limiter import AdaptiveDownloadLimiter from icloudpd.log_level import LogLevel +from icloudpd.metrics import RunMetrics, write_metrics_json from icloudpd.mfa_provider import MFAProvider from icloudpd.password_provider import PasswordProvider from icloudpd.paths import local_download_path, remove_unicode_chars +from icloudpd.retry_utils import ( + RetryConfig, + is_fatal_auth_config_error, + is_throttle_error, + is_transient_error, +) from icloudpd.server import serve_app +from icloudpd.state_db import ( + checkpoint_wal, + clear_asset_tasks_need_url_refresh, + initialize_state_db, + load_checkpoint, + mark_asset_tasks_need_url_refresh, + prune_completed_tasks, + record_asset_checksum_result, + requeue_in_progress_tasks, + requeue_stale_leases, + resolve_state_db_path, + save_checkpoint, + upsert_asset_tasks, + vacuum_state_db, +) from icloudpd.status import Status, StatusExchange from icloudpd.string_helpers import parse_timestamp_or_timedelta, truncate_middle from icloudpd.xmp_sidecar import generate_xmp_file -from pyicloud_ipd.asset_version import add_suffix_to_filename, calculate_version_filename +from pyicloud_ipd.asset_version import ( + AssetVersion, + add_suffix_to_filename, + calculate_version_filename, +) from pyicloud_ipd.base import PyiCloudService from pyicloud_ipd.exceptions import ( PyiCloudAPIResponseException, @@ -81,6 +111,78 @@ from pyicloud_ipd.version_size import AssetVersionSize, LivePhotoVersionSize freeze_support() # fmt: skip # fixing tqdm on macos +EXIT_CANCELLED = 130 +ENGINE_MODE_LEGACY_STATELESS = "legacy_stateless" +ENGINE_MODE_STATEFUL = "stateful_engine" +THROTTLE_ALERT_THRESHOLD = 5 + + +def determine_engine_mode(state_db_path: str | None) -> str: + return ENGINE_MODE_STATEFUL if state_db_path else ENGINE_MODE_LEGACY_STATELESS + + +def emit_throttle_alert_if_needed( + logger: logging.Logger, + run_metrics: RunMetrics, + threshold: int = THROTTLE_ALERT_THRESHOLD, +) -> bool: + if run_metrics.throttle_events < threshold: + return False + logger.warning( + "Repeated throttling detected for user %s (%d events). Consider lowering --download-workers, increasing --watch-with-interval, and keeping --no-remote-count enabled.", + run_metrics.username, + run_metrics.throttle_events, + ) + return True + + +class ShutdownController: + def __init__(self, status_exchange: StatusExchange): + self._status_exchange = status_exchange + self._event = Event() + self._signal_name: str | None = None + self._installed = False + self._previous_handlers: dict[int, typing.Any] = {} + + def install(self) -> None: + if self._installed or current_thread() is not main_thread(): + return + for sig in (signal.SIGINT, signal.SIGTERM): + self._previous_handlers[sig] = signal.getsignal(sig) + signal.signal(sig, self._build_handler(sig)) + self._installed = True + + def restore(self) -> None: + if not self._installed: + return + for sig, previous in self._previous_handlers.items(): + signal.signal(sig, previous) + self._previous_handlers.clear() + self._installed = False + + def _build_handler(self, sig: int) -> Callable[[int, typing.Any], None]: + def handler(_signum: int, _frame: typing.Any) -> None: + self._signal_name = signal.Signals(sig).name + self._event.set() + self._status_exchange.get_progress().cancel = True + + return handler + + def request_stop(self, reason: str | None = None) -> None: + self._signal_name = reason or self._signal_name + self._event.set() + self._status_exchange.get_progress().cancel = True + + def requested(self) -> bool: + return self._event.is_set() or self._status_exchange.get_progress().cancel + + def signal_name(self) -> str | None: + return self._signal_name + + def sleep_or_stop(self, seconds: float) -> bool: + if seconds <= 0: + return not self.requested() + return not self._event.wait(timeout=seconds) def build_filename_cleaner(keep_unicode: bool) -> Callable[[str], str]: @@ -138,6 +240,10 @@ def get_password_from_webui( # wait for input while True: + if status_exchange.get_progress().cancel: + logger.info("Password input cancelled") + status_exchange.replace_status(Status.NEED_PASSWORD, Status.NO_INPUT_NEEDED) + return None status = status_exchange.get_status() if status == Status.NEED_PASSWORD: time.sleep(1) @@ -209,12 +315,81 @@ def ensure_tzinfo(tz: datetime.tzinfo, input: datetime.datetime) -> datetime.dat def create_logger(config: GlobalConfig) -> logging.Logger: - logging.basicConfig( - format="%(asctime)s %(levelname)-8s %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - stream=sys.stdout, - ) logger = logging.getLogger("icloudpd") + logger.handlers.clear() + logger.propagate = True + + class RunContextFilter(logging.Filter): + def __init__(self, run_id: str): + super().__init__() + self._run_id = run_id + + def filter(self, record: logging.LogRecord) -> bool: + if not hasattr(record, "run_id"): + record.run_id = self._run_id # type: ignore[attr-defined] + if not hasattr(record, "asset_id"): + record.asset_id = None # type: ignore[attr-defined] + if not hasattr(record, "attempt"): + record.attempt = None # type: ignore[attr-defined] + if not hasattr(record, "http_status"): + record.http_status = None # type: ignore[attr-defined] + return True + + class SensitiveDataRedactionFilter(logging.Filter): + _KEY_VALUE_PATTERNS = [ + # JSON payload style: "password": "value" + re.compile( + r'(?i)"(password|passphrase|session_token|trust_token|token|authorization|cookie|scnt)"\s*:\s*"([^"]+)"' + ), + # key=value style in plain logs + re.compile( + r"(?i)\b(password|passphrase|session_token|trust_token|token|authorization|cookie|scnt)\b\s*[:=]\s*(Bearer\s+[^\s,;]+|[^\s,;]+)" + ), + # Authorization bearer tokens + re.compile(r"(?i)\bBearer\s+[A-Za-z0-9\-._~+/]+=*"), + ] + + def _redact(self, message: str) -> str: + redacted = message + redacted = self._KEY_VALUE_PATTERNS[0].sub(r'"\1":"REDACTED"', redacted) + redacted = self._KEY_VALUE_PATTERNS[1].sub(r"\1=REDACTED", redacted) + redacted = self._KEY_VALUE_PATTERNS[2].sub("Bearer REDACTED", redacted) + return redacted + + def filter(self, record: logging.LogRecord) -> bool: + redacted = self._redact(record.getMessage()) + record.msg = redacted + record.args = () # type: ignore[assignment] + return True + + class JsonLogFormatter(logging.Formatter): + def format(self, record: logging.LogRecord) -> str: + payload = { + "timestamp": self.formatTime(record, "%Y-%m-%dT%H:%M:%S"), + "level": record.levelname, + "logger": record.name, + "message": record.getMessage(), + "run_id": getattr(record, "run_id", None), + "asset_id": getattr(record, "asset_id", None), + "attempt": getattr(record, "attempt", None), + "http_status": getattr(record, "http_status", None), + } + return json.dumps(payload, ensure_ascii=True) + + handler = logging.StreamHandler(sys.stdout) + if config.log_format == "json": + handler.setFormatter(JsonLogFormatter()) + else: + handler.setFormatter( + logging.Formatter( + "%(asctime)s %(levelname)-8s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + ) + logger.addHandler(handler) + logger.addFilter(RunContextFilter(uuid.uuid4().hex)) + logger.addFilter(SensitiveDataRedactionFilter()) + if config.only_print_filenames: logger.disabled = True else: @@ -241,6 +416,8 @@ def run_with_configs(global_config: GlobalConfig, user_configs: Sequence[UserCon # Create shared status exchange for web server and progress tracking shared_status_exchange = StatusExchange() + shutdown = ShutdownController(shared_status_exchange) + shutdown.install() # Check if any user needs web server (webui for MFA or passwords) needs_web_server = global_config.mfa_provider == MFAProvider.WEBUI or any( @@ -258,7 +435,16 @@ def run_with_configs(global_config: GlobalConfig, user_configs: Sequence[UserCon if not watch_interval: # No watch mode - process each user once and exit - return _process_all_users_once(global_config, user_configs, logger, shared_status_exchange) + try: + return _process_all_users_once( + global_config, + user_configs, + logger, + shared_status_exchange, + shutdown, + ) + finally: + shutdown.restore() else: # Watch mode - infinite loop processing all users, then wait skip_bar = not os.environ.get("FORCE_TQDM") and ( @@ -267,46 +453,63 @@ def run_with_configs(global_config: GlobalConfig, user_configs: Sequence[UserCon or not sys.stdout.isatty() ) - while True: - # Process all user configs in this iteration - result = _process_all_users_once( - global_config, user_configs, logger, shared_status_exchange - ) - - # If any critical operation (auth-only, list commands) succeeded, exit - if result == 0: - first_user = user_configs[0] if user_configs else None - if first_user and ( - first_user.auth_only or first_user.list_albums or first_user.list_libraries - ): - return 0 - - # Wait for the watch interval before next iteration - # Clear current user during wait period to avoid misleading UI - shared_status_exchange.clear_current_user() - logger.info(f"Waiting for {watch_interval} sec...") - interval: Sequence[int] = range(1, watch_interval) - iterable: Sequence[int] = ( - interval - if skip_bar - else typing.cast( - Sequence[int], - tqdm( - iterable=interval, - desc="Waiting...", - ascii=True, - leave=False, - dynamic_ncols=True, - ), + try: + while True: + if shutdown.requested(): + signal_name = shutdown.signal_name() or "cancellation signal" + logger.info("Run cancelled by %s", signal_name) + return EXIT_CANCELLED + + # Process all user configs in this iteration + result = _process_all_users_once( + global_config, + user_configs, + logger, + shared_status_exchange, + shutdown, ) - ) - for counter in iterable: - # Update shared status exchange with wait progress - shared_status_exchange.get_progress().waiting = watch_interval - counter - if shared_status_exchange.get_progress().resume: - shared_status_exchange.get_progress().reset() - break - time.sleep(1) + if result == EXIT_CANCELLED: + return EXIT_CANCELLED + + # If any critical operation (auth-only, list commands) succeeded, exit + if result == 0: + first_user = user_configs[0] if user_configs else None + if first_user and ( + first_user.auth_only or first_user.list_albums or first_user.list_libraries + ): + return 0 + + # Wait for the watch interval before next iteration + # Clear current user during wait period to avoid misleading UI + shared_status_exchange.clear_current_user() + logger.info(f"Waiting for {watch_interval} sec...") + interval: Sequence[int] = range(1, watch_interval) + iterable: Sequence[int] = ( + interval + if skip_bar + else typing.cast( + Sequence[int], + tqdm( + iterable=interval, + desc="Waiting...", + ascii=True, + leave=False, + dynamic_ncols=True, + ), + ) + ) + for counter in iterable: + # Update shared status exchange with wait progress + shared_status_exchange.get_progress().waiting = watch_interval - counter + if shared_status_exchange.get_progress().resume: + shared_status_exchange.get_progress().reset() + break + if not shutdown.sleep_or_stop(1): + signal_name = shutdown.signal_name() or "cancellation signal" + logger.info("Run cancelled by %s", signal_name) + return EXIT_CANCELLED + finally: + shutdown.restore() def _process_all_users_once( @@ -314,14 +517,54 @@ def _process_all_users_once( user_configs: Sequence[UserConfig], logger: logging.Logger, shared_status_exchange: StatusExchange, + shutdown: ShutdownController, ) -> int: """Process all user configs once (used by both single run and watch mode)""" # Set global config and all user configs to status exchange once, before processing shared_status_exchange.set_global_config(global_config) shared_status_exchange.set_user_configs(user_configs) + user_metrics_snapshots: list[dict[str, float | int | str]] = [] + + def write_run_summary(exit_code: int) -> None: + if not global_config.metrics_json: + return + users_with_failures = 0 + for user in user_metrics_snapshots: + failed = user.get("downloads_failed", 0) + if isinstance(failed, int) and failed > 0: + users_with_failures += 1 + if exit_code == 0 and users_with_failures > 0: + status = "partial_success" + elif exit_code == 0: + status = "success" + elif exit_code == EXIT_CANCELLED: + status = "cancelled" + elif exit_code == 2: + status = "cli_error" + else: + status = "runtime_error" + write_metrics_json( + global_config.metrics_json, + { + "exit_code": exit_code, + "status": status, + "users_total": len(user_metrics_snapshots), + "users_with_failures": users_with_failures, + "users": user_metrics_snapshots, + }, + ) + + if shutdown.requested(): + write_run_summary(EXIT_CANCELLED) + return EXIT_CANCELLED for user_config in user_configs: + if shutdown.requested(): + signal_name = shutdown.signal_name() or "cancellation signal" + logger.info("Stopping before user %s due to %s", user_config.username, signal_name) + write_run_summary(EXIT_CANCELLED) + return EXIT_CANCELLED with logging_redirect_tqdm(): # Use shared status exchange instead of creating new ones per user status_exchange = shared_status_exchange @@ -386,12 +629,52 @@ def password_provider(_username: str) -> str | None: ) # Set up function builders + state_db_path = resolve_state_db_path(user_config.state_db, user_config.cookie_directory) + engine_mode = determine_engine_mode(state_db_path) + if state_db_path: + logger.info("Initializing state DB at %s", state_db_path) + initialize_state_db(state_db_path) + logger.info( + "Engine mode: %s (persistent task/checkpoint state enabled)", + engine_mode, + ) + else: + logger.info( + "Engine mode: %s (filesystem skip semantics, no state DB required)", + engine_mode, + ) + + download.set_retry_config( + RetryConfig( + max_retries=global_config.max_retries, + backoff_base_seconds=global_config.backoff_base_seconds, + backoff_max_seconds=global_config.backoff_max_seconds, + respect_retry_after=global_config.respect_retry_after, + throttle_cooldown_seconds=global_config.throttle_cooldown_seconds, + ) + ) + download.set_download_chunk_bytes(user_config.download_chunk_bytes) + download.set_download_verification( + verify_size=user_config.verify_size, + verify_checksum=user_config.verify_checksum, + ) + download.set_download_limiter( + AdaptiveDownloadLimiter( + max_workers=user_config.download_workers, + cooldown_seconds=global_config.throttle_cooldown_seconds, + ) + ) + run_metrics = RunMetrics(username=user_config.username, run_mode=engine_mode) + run_metrics.start() + download.set_metrics_collector(run_metrics) passer = partial( where_builder, logger, user_config.skip_videos, user_config.skip_created_before, user_config.skip_created_after, + user_config.skip_added_before, + user_config.skip_added_after, user_config.skip_photos, filename_builder, ) @@ -440,16 +723,34 @@ def password_provider(_username: str) -> str | None: status_exchange, global_config, user_config, + shutdown, password_providers_dict, + run_metrics, passer, downloader, notificator, lp_filename_generator, ) + if state_db_path: + if user_config.state_db_prune_completed_days is not None: + pruned = prune_completed_tasks( + state_db_path, + older_than_days=user_config.state_db_prune_completed_days, + ) + if pruned > 0: + logger.info("Pruned %d completed/failed state DB task(s)", pruned) + checkpoint_wal(state_db_path, mode="PASSIVE") + if user_config.state_db_vacuum: + logger.info("Running state DB VACUUM at %s", state_db_path) + vacuum_state_db(state_db_path) + run_metrics.finish() + emit_throttle_alert_if_needed(logger, run_metrics) + user_metrics_snapshots.append(run_metrics.snapshot()) # If any user config fails and we're not in watch mode, return the error code if result != 0: if not global_config.watch_with_interval: + write_run_summary(result) return result else: # In watch mode, log error and continue with next user @@ -457,6 +758,7 @@ def password_provider(_username: str) -> str | None: f"Error processing user {user_config.username}, continuing with next user..." ) + write_run_summary(0) return 0 @@ -517,6 +819,8 @@ def where_builder( skip_videos: bool, skip_created_before: datetime.datetime | datetime.timedelta | None, skip_created_after: datetime.datetime | datetime.timedelta | None, + skip_added_before: datetime.datetime | datetime.timedelta | None, + skip_added_after: datetime.datetime | datetime.timedelta | None, skip_photos: bool, filename_builder: Callable[[PhotoAsset], str], photo: PhotoAsset, @@ -540,6 +844,26 @@ def where_builder( logger.debug(skip_created_after_message(temp_created_after, photo, filename_builder)) return False + if skip_added_before is not None: + temp_added_before = offset_to_datetime(skip_added_before) + try: + added_date = photo.added_date.astimezone(get_localzone()) + except (KeyError, TypeError, ValueError, OSError): + added_date = None + if added_date is not None and added_date < temp_added_before: + logger.debug(skip_added_before_message(temp_added_before, photo, filename_builder)) + return False + + if skip_added_after is not None: + temp_added_after = offset_to_datetime(skip_added_after) + try: + added_date = photo.added_date.astimezone(get_localzone()) + except (KeyError, TypeError, ValueError, OSError): + added_date = None + if added_date is not None and added_date > temp_added_after: + logger.debug(skip_added_after_message(temp_added_after, photo, filename_builder)) + return False + return True @@ -561,6 +885,24 @@ def skip_created_after_message( return f"Skipping {filename}, as it was created {photo.created}, after {target_created_date}." +def skip_added_before_message( + target_added_date: datetime.datetime, + photo: PhotoAsset, + filename_builder: Callable[[PhotoAsset], str], +) -> str: + filename = filename_builder(photo) + return f"Skipping {filename}, as it was added {photo.added_date}, before {target_added_date}." + + +def skip_added_after_message( + target_added_date: datetime.datetime, + photo: PhotoAsset, + filename_builder: Callable[[PhotoAsset], str], +) -> str: + filename = filename_builder(photo) + return f"Skipping {filename}, as it was added {photo.added_date}, after {target_added_date}." + + def download_builder( logger: logging.Logger, folder_structure: str, @@ -634,6 +976,13 @@ def download_builder( download_dir = os.path.normpath(os.path.join(directory, date_path)) success = False + def refresh_asset_version(target_size: AssetVersionSize | LivePhotoVersionSize) -> AssetVersion | None: + # Drop cached versions and rebuild from latest asset metadata snapshot. + if hasattr(photo, "_versions"): + photo._versions = None # type: ignore[attr-defined] + refreshed_versions = photo.versions_with_raw_policy(raw_policy) + return refreshed_versions.get(target_size) + for download_size in primary_sizes: if download_size not in versions and download_size != AssetVersionSize.ORIGINAL: if force_size: @@ -700,6 +1049,7 @@ def download_builder( version, download_size, filename_builder, + refresh_version=partial(refresh_asset_version, download_size), ) success = download_result @@ -798,6 +1148,7 @@ def download_builder( version, lp_size, filename_builder, + refresh_version=partial(refresh_asset_version, lp_size), ) success = download_result and success if download_result: @@ -877,9 +1228,11 @@ def core_single_run( status_exchange: StatusExchange, global_config: GlobalConfig, user_config: UserConfig, + shutdown: ShutdownController, password_providers_dict: Dict[ PasswordProvider, Tuple[Callable[[str], str | None], Callable[[str, str], None]] ], + run_metrics: RunMetrics, passer: Callable[[PhotoAsset], bool], downloader: Callable[[PyiCloudService, Counter, PhotoAsset], bool], notificator: Callable[[], None], @@ -892,7 +1245,23 @@ def core_single_run( or global_config.no_progress_bar or not sys.stdout.isatty() ) + retry_config = RetryConfig( + max_retries=global_config.max_retries, + backoff_base_seconds=global_config.backoff_base_seconds, + backoff_max_seconds=global_config.backoff_max_seconds, + respect_retry_after=global_config.respect_retry_after, + throttle_cooldown_seconds=global_config.throttle_cooldown_seconds, + ) + state_db_path = resolve_state_db_path(user_config.state_db, user_config.cookie_directory) + retry_count = 0 + stale_requeue_done = False while True: # retry loop (not watch - only for immediate retries) + if shutdown.requested(): + if state_db_path: + requeue_in_progress_tasks(state_db_path) + signal_name = shutdown.signal_name() or "cancellation signal" + logger.info("Run cancelled before authentication due to %s", signal_name) + return EXIT_CANCELLED captured_responses: List[Mapping[str, Any]] = [] def append_response(captured: List[Mapping[str, Any]], response: Mapping[str, Any]) -> None: @@ -973,21 +1342,50 @@ def append_response(captured: List[Mapping[str, Any]], response: Mapping[str, An album_phrase = f" from albums {','.join(user_config.albums)}" logger.debug(f"Looking up all {photo_video_phrase}{album_phrase}...") + if state_db_path and not stale_requeue_done: + requeued = requeue_stale_leases(state_db_path) + stale_requeue_done = True + if requeued > 0: + logger.info("Requeued %d stale in-progress task(s) from prior run", requeued) albums: Iterable[PhotoAlbum] = ( list(map_(library_object.albums.__getitem__, user_config.albums)) if len(user_config.albums) > 0 else [library_object.all] ) - album_lengths: Callable[[Iterable[PhotoAlbum]], Iterable[int]] = partial_1_1( - map_, len + for album in albums: + album.page_size = user_config.album_page_size + + should_fetch_remote_count = ( + not user_config.no_remote_count + and user_config.until_found is None ) + if not should_fetch_remote_count: + photos_count = None + else: + album_lengths: Callable[[Iterable[PhotoAlbum]], Iterable[int]] = partial_1_1( + map_, len + ) - def sum_(inp: Iterable[int]) -> int: - return sum(inp) + def sum_(inp: Iterable[int]) -> int: + return sum(inp) - photos_count: int | None = compose(sum_, album_lengths)(albums) + photos_count = compose(sum_, album_lengths)(albums) for photo_album in albums: + album_name = str(photo_album) + library_name = user_config.library + if state_db_path: + checkpoint = load_checkpoint( + state_db_path, library=library_name, album=album_name + ) + if checkpoint is not None: + logger.debug( + "Resuming album %s from checkpoint offset %d", + album_name, + checkpoint, + ) + photo_album.offset = checkpoint + photos_enumerator: Iterable[PhotoAsset] = photo_album # Optional: Only download the x most recent photos. @@ -1070,7 +1468,10 @@ def should_break(counter: Counter) -> bool: download_photo = partial(downloader, icloud) for item in photos_bar: + if shutdown.requested(): + break try: + run_metrics.on_asset_considered() if should_break(consecutive_files_found): logger.info( "Found %s consecutive previously downloaded photos. Exiting", @@ -1080,10 +1481,51 @@ def should_break(counter: Counter) -> bool: # item = next(photos_iterator) should_delete = False + if state_db_path: + upsert_asset_tasks( + state_db_path, + photo=item, + library=library_name, + album=album_name, + ) + save_checkpoint( + state_db_path, + library=library_name, + album=album_name, + start_rank=int(photo_album.offset), + ) + passer_result = passer(item) download_result = passer_result and download_photo( consecutive_files_found, item ) + url_refresh_needed = download.consume_url_refresh_needed_signal() + if state_db_path and passer_result and url_refresh_needed: + mark_asset_tasks_need_url_refresh( + state_db_path, + asset_id=item.id, + library=library_name, + album=album_name, + ) + elif state_db_path and passer_result and download_result: + clear_asset_tasks_need_url_refresh( + state_db_path, + asset_id=item.id, + library=library_name, + album=album_name, + ) + if state_db_path and passer_result and download_result: + record_asset_checksum_result( + state_db_path, + asset_id=item.id, + library=library_name, + album=album_name, + checksum_result=( + "verified" + if user_config.verify_checksum + else "not_checked" + ), + ) if download_result and user_config.delete_after_download: should_delete = True @@ -1144,6 +1586,8 @@ def should_break(counter: Counter) -> bool: photos_counter += 1 status_exchange.get_progress().photos_counter = photos_counter + if photos_count is not None: + run_metrics.set_queue_depth(photos_count - photos_counter) if status_exchange.get_progress().cancel: break @@ -1157,10 +1601,15 @@ def should_break(counter: Counter) -> bool: pass if status_exchange.get_progress().cancel: - logger.info("Iteration was cancelled") + if state_db_path: + requeue_in_progress_tasks(state_db_path) + signal_name = shutdown.signal_name() or "cancellation signal" + logger.info("Iteration was cancelled by %s", signal_name) status_exchange.get_progress().photos_last_message = ( "Iteration was cancelled" ) + status_exchange.get_progress().reset() + return EXIT_CANCELLED else: if user_config.skip_photos or user_config.skip_videos: photo_video_phrase = ( @@ -1207,6 +1656,10 @@ def should_break(counter: Counter) -> bool: PyiCloudServiceUnavailableException, PyiCloudAPIResponseException, PyiCloudConnectionErrorException, + ChunkedEncodingError, + ContentDecodingError, + StreamConsumedError, + UnrewindableBodyError, ) as error: logger.info(error) dump_responses(logger.debug, captured_responses) @@ -1222,18 +1675,36 @@ def should_break(counter: Counter) -> bool: pass else: pass + if is_fatal_auth_config_error(error): + return 1 + if is_transient_error(error) and retry_count < retry_config.max_retries: + retry_count += 1 + wait_seconds = retry_config.next_delay_seconds( + retry_count, + throttle_error=is_throttle_error(error), + ) + logger.info( + "Transient error (%s). Retrying in %.1f seconds (%d/%d)...", + type(error).__name__, + wait_seconds, + retry_count, + retry_config.max_retries, + ) + if not shutdown.sleep_or_stop(wait_seconds): + if state_db_path: + requeue_in_progress_tasks(state_db_path) + signal_name = shutdown.signal_name() or "cancellation signal" + logger.info("Run cancelled during retry backoff by %s", signal_name) + return EXIT_CANCELLED + continue # In single run mode, return error after webui retry attempts return 1 - except ( - ChunkedEncodingError, - ContentDecodingError, - StreamConsumedError, - UnrewindableBodyError, - ) as error: - logger.debug(error) - logger.debug("Retrying...") - # these errors we can safely retry - continue + except KeyboardInterrupt: + shutdown.request_stop("KeyboardInterrupt") + if state_db_path: + requeue_in_progress_tasks(state_db_path) + logger.info("Run cancelled by KeyboardInterrupt") + return EXIT_CANCELLED except Exception: dump_responses(logger.debug, captured_responses) raise diff --git a/src/icloudpd/cli.py b/src/icloudpd/cli.py index 9c33ac1af..eec79eb64 100644 --- a/src/icloudpd/cli.py +++ b/src/icloudpd/cli.py @@ -12,6 +12,7 @@ import foundation from foundation.core import chain_from_iterable, compose, map_, partial_1_1, skip from foundation.string_utils import lower +from icloudpd import constants from icloudpd.base import ensure_tzinfo, run_with_configs from icloudpd.config import GlobalConfig, UserConfig from icloudpd.log_level import LogLevel @@ -23,6 +24,10 @@ from pyicloud_ipd.raw_policy import RawTreatmentPolicy from pyicloud_ipd.version_size import AssetVersionSize, LivePhotoVersionSize +EXIT_SUCCESS = 0 +EXIT_RUNTIME_ERROR = 1 +EXIT_USAGE_ERROR = 2 + def map_align_raw_to_enum(align_raw_str: str) -> RawTreatmentPolicy: """Map user-friendly CLI strings to RawTreatmentPolicy enum values.""" @@ -251,11 +256,80 @@ def add_options_for_user(parser: argparse.ArgumentParser) -> argparse.ArgumentPa default=None, type=parse_timestamp_or_timedelta_tz_error, ) + cloned.add_argument( + "--skip-added-before", + help="Do not process assets added to iCloud before the specified timestamp in ISO format (2025-01-02) or interval backwards from now (20d = 20 days ago)", + default=None, + type=parse_timestamp_or_timedelta_tz_error, + ) + cloned.add_argument( + "--skip-added-after", + help="Do not process assets added to iCloud after the specified timestamp in ISO format (2025-01-02) or interval backwards from now (20d = 20 days ago)", + default=None, + type=parse_timestamp_or_timedelta_tz_error, + ) cloned.add_argument( "--skip-photos", help="Don't download any photos (default: download all photos and videos)", action="store_true", ) + cloned.add_argument( + "--download-chunk-bytes", + help="Chunk size in bytes for streamed downloads. Default: %(default)s", + type=int, + default=262144, + ) + cloned.add_argument( + "--download-workers", + help="Number of download worker slots. Metadata enumeration remains single-threaded. Default: %(default)s", + type=int, + default=4, + ) + cloned.add_argument( + "--verify-size", + help="Validate downloaded file size against iCloud metadata (default: disabled)", + action=argparse.BooleanOptionalAction, + default=False, + ) + cloned.add_argument( + "--verify-checksum", + help="Validate downloaded file checksum when metadata checksum is available (default: disabled)", + action=argparse.BooleanOptionalAction, + default=False, + ) + cloned.add_argument( + "--album-page-size", + help="Album enumeration page size (results per request). Recommended range: 50-500. Default: %(default)s", + type=int, + default=100, + ) + cloned.add_argument( + "--no-remote-count", + help="Skip remote album count lookups and run without known total progress", + action="store_true", + default=False, + ) + cloned.add_argument( + "--state-db", + metavar="PATH", + nargs="?", + const="auto", + default=None, + help="Enable persistent state DB for resumable runs. If PATH is omitted, defaults to /icloudpd.sqlite", + ) + cloned.add_argument( + "--state-db-prune-completed-days", + metavar="DAYS", + type=int, + default=None, + help="Delete completed/failed state DB tasks older than DAYS (run-end maintenance)", + ) + cloned.add_argument( + "--state-db-vacuum", + help="Run VACUUM on state DB at run end (can be slow on large DBs)", + action="store_true", + default=False, + ) return cloned @@ -314,6 +388,13 @@ def add_global_options(parser: argparse.ArgumentParser) -> argparse.ArgumentPars default="debug", type=lower, ) + cloned.add_argument( + "--log-format", + help="Log output format. Default: %(default)s", + choices=["text", "json"], + default="text", + type=lower, + ) cloned.add_argument( "--no-progress-bar", help="Disable the one-line progress bar and print log messages on separate lines " @@ -325,7 +406,7 @@ def add_global_options(parser: argparse.ArgumentParser) -> argparse.ArgumentPars deprecated_kwargs["deprecated"] = True cloned.add_argument( "--threads-num", - help="Number of CPU threads - deprecated & always 1. To be removed in a future version", + help="Number of CPU threads - deprecated & always 1. Use --download-workers for download concurrency. To be removed in a future version", type=int, default=1, **deprecated_kwargs, @@ -358,6 +439,42 @@ def add_global_options(parser: argparse.ArgumentParser) -> argparse.ArgumentPars default="console", type=lower, ) + cloned.add_argument( + "--max-retries", + help="Maximum retry attempts for transient metadata and download failures. Default: %(default)s", + type=int, + default=constants.MAX_RETRIES, + ) + cloned.add_argument( + "--backoff-base-seconds", + help="Base retry delay in seconds for exponential backoff. Default: %(default)s", + type=float, + default=5.0, + ) + cloned.add_argument( + "--backoff-max-seconds", + help="Maximum retry delay in seconds. Default: %(default)s", + type=float, + default=300.0, + ) + cloned.add_argument( + "--respect-retry-after", + help="Respect server Retry-After headers when available (default: enabled)", + action=argparse.BooleanOptionalAction, + default=True, + ) + cloned.add_argument( + "--throttle-cooldown-seconds", + help="Minimum cool-down delay in seconds for throttling errors. Default: %(default)s", + type=float, + default=60.0, + ) + cloned.add_argument( + "--metrics-json", + help="Write run metrics as JSON to this path", + default=None, + type=str, + ) return cloned @@ -471,7 +588,18 @@ def map_to_config(user_ns: argparse.Namespace) -> UserConfig: file_match_policy=FileMatchPolicy(user_ns.file_match_policy), skip_created_before=user_ns.skip_created_before, skip_created_after=user_ns.skip_created_after, + skip_added_before=user_ns.skip_added_before, + skip_added_after=user_ns.skip_added_after, skip_photos=user_ns.skip_photos, + download_chunk_bytes=user_ns.download_chunk_bytes, + download_workers=user_ns.download_workers, + verify_size=user_ns.verify_size, + verify_checksum=user_ns.verify_checksum, + album_page_size=user_ns.album_page_size, + no_remote_count=user_ns.no_remote_count, + state_db=user_ns.state_db, + state_db_prune_completed_days=user_ns.state_db_prune_completed_days, + state_db_vacuum=user_ns.state_db_vacuum, ) @@ -515,6 +643,7 @@ def parse(args: Sequence[str]) -> Tuple[GlobalConfig, Sequence[UserConfig]]: use_os_locale=global_ns.use_os_locale, only_print_filenames=global_ns.only_print_filenames, log_level=log_level(global_ns.log_level), + log_format=global_ns.log_format, no_progress_bar=global_ns.no_progress_bar, threads_num=global_ns.threads_num, domain=global_ns.domain, @@ -528,6 +657,12 @@ def parse(args: Sequence[str]) -> Tuple[GlobalConfig, Sequence[UserConfig]]: ) ), mfa_provider=MFAProvider(global_ns.mfa_provider), + max_retries=global_ns.max_retries, + backoff_base_seconds=global_ns.backoff_base_seconds, + backoff_max_seconds=global_ns.backoff_max_seconds, + respect_retry_after=global_ns.respect_retry_after, + throttle_cooldown_seconds=global_ns.throttle_cooldown_seconds, + metrics_json=global_ns.metrics_json, ), user_nses, ) @@ -538,7 +673,7 @@ def cli() -> int: global_ns, user_nses = parse(sys.argv[1:]) except argparse.ArgumentError as error: print(error) - return 2 + return EXIT_USAGE_ERROR if global_ns.use_os_locale: from locale import LC_ALL, setlocale @@ -547,17 +682,17 @@ def cli() -> int: pass if global_ns.help: print(format_help()) - return 0 + return EXIT_SUCCESS elif global_ns.version: print(foundation.version_info_formatted()) - return 0 + return EXIT_SUCCESS else: # check param compatibility if [user_ns for user_ns in user_nses if user_ns.skip_videos and user_ns.skip_photos]: print( "Only one of --skip-videos and --skip-photos can be used at a time for each configuration" ) - return 2 + return EXIT_USAGE_ERROR # check required directory param only if not list albums elif [ @@ -571,7 +706,31 @@ def cli() -> int: print( "--auth-only, --directory, --list-libraries, or --list-albums are required for each configuration" ) - return 2 + return EXIT_USAGE_ERROR + elif [ + user_ns + for user_ns in user_nses + if user_ns.directory is not None and not pathlib.Path(user_ns.directory).exists() + ]: + print("Directory specified in --directory does not exist") + return EXIT_USAGE_ERROR + elif [user_ns for user_ns in user_nses if user_ns.download_chunk_bytes <= 0]: + print("--download-chunk-bytes must be greater than 0") + return EXIT_USAGE_ERROR + elif [user_ns for user_ns in user_nses if user_ns.download_workers <= 0]: + print("--download-workers must be greater than 0") + return EXIT_USAGE_ERROR + elif [user_ns for user_ns in user_nses if user_ns.album_page_size < 1]: + print("--album-page-size must be greater than 0") + return EXIT_USAGE_ERROR + elif [ + user_ns + for user_ns in user_nses + if user_ns.state_db_prune_completed_days is not None + and user_ns.state_db_prune_completed_days <= 0 + ]: + print("--state-db-prune-completed-days must be greater than 0") + return EXIT_USAGE_ERROR elif [ user_ns @@ -581,7 +740,7 @@ def cli() -> int: print( "--auto-delete and --delete-after-download are mutually exclusive per configuration" ) - return 2 + return EXIT_USAGE_ERROR elif [ user_ns @@ -591,7 +750,7 @@ def cli() -> int: print( "--keep-icloud-recent-days and --delete-after-download should not be used together in one configuration" ) - return 2 + return EXIT_USAGE_ERROR elif global_ns.watch_with_interval and ( [ @@ -604,7 +763,22 @@ def cli() -> int: print( "--watch-with-interval is not compatible with --list-albums, --list-libraries, --only-print-filenames, and --auth-only" ) - return 2 + return EXIT_USAGE_ERROR + elif global_ns.max_retries < 0: + print("--max-retries must be 0 or greater") + return EXIT_USAGE_ERROR + elif global_ns.backoff_base_seconds <= 0: + print("--backoff-base-seconds must be greater than 0") + return EXIT_USAGE_ERROR + elif global_ns.backoff_max_seconds <= 0: + print("--backoff-max-seconds must be greater than 0") + return EXIT_USAGE_ERROR + elif global_ns.backoff_max_seconds < global_ns.backoff_base_seconds: + print("--backoff-max-seconds must be greater than or equal to --backoff-base-seconds") + return EXIT_USAGE_ERROR + elif global_ns.throttle_cooldown_seconds < 0: + print("--throttle-cooldown-seconds must be 0 or greater") + return EXIT_USAGE_ERROR else: return run_with_configs(global_ns, user_nses) diff --git a/src/icloudpd/config.py b/src/icloudpd/config.py index 019843864..5787bde9c 100644 --- a/src/icloudpd/config.py +++ b/src/icloudpd/config.py @@ -50,6 +50,17 @@ class _DefaultConfig: skip_created_before: datetime.datetime | datetime.timedelta | None skip_created_after: datetime.datetime | datetime.timedelta | None skip_photos: bool + skip_added_before: datetime.datetime | datetime.timedelta | None = None + skip_added_after: datetime.datetime | datetime.timedelta | None = None + download_chunk_bytes: int = 262144 + verify_size: bool = False + verify_checksum: bool = False + download_workers: int = 4 + album_page_size: int = 100 + no_remote_count: bool = False + state_db: str | None = None + state_db_prune_completed_days: int | None = None + state_db_vacuum: bool = False @dataclass(kw_only=True) @@ -65,9 +76,16 @@ class GlobalConfig: use_os_locale: bool only_print_filenames: bool log_level: LogLevel + log_format: str = "text" no_progress_bar: bool threads_num: int domain: str watch_with_interval: int | None password_providers: Sequence[PasswordProvider] mfa_provider: MFAProvider + max_retries: int + backoff_base_seconds: float + backoff_max_seconds: float + respect_retry_after: bool + throttle_cooldown_seconds: float + metrics_json: str | None = None diff --git a/src/icloudpd/counter.py b/src/icloudpd/counter.py index 06b1e80ac..d01eeaf01 100644 --- a/src/icloudpd/counter.py +++ b/src/icloudpd/counter.py @@ -1,22 +1,22 @@ -"""Atomic counter""" +"""Thread-safe counter.""" -from multiprocessing import Lock, RawValue +from threading import Lock class Counter: def __init__(self, value: int = 0): self.initial_value = value - self.val = RawValue("i", value) + self.val = value self.lock = Lock() def increment(self) -> None: with self.lock: - self.val.value += 1 + self.val += 1 def reset(self) -> None: with self.lock: - self.val = RawValue("i", self.initial_value) + self.val = self.initial_value def value(self) -> int: with self.lock: - return int(self.val.value) + return int(self.val) diff --git a/src/icloudpd/download.py b/src/icloudpd/download.py index 78da134b5..320d8e836 100644 --- a/src/icloudpd/download.py +++ b/src/icloudpd/download.py @@ -2,8 +2,10 @@ import base64 import datetime +import hashlib import logging import os +import shutil import time from functools import partial from typing import Callable @@ -11,14 +13,91 @@ from requests import Response from tzlocal import get_localzone -# Import the constants object so that we can mock WAIT_SECONDS in tests from icloudpd import constants +from icloudpd.limiter import AdaptiveDownloadLimiter +from icloudpd.metrics import RunMetrics +from icloudpd.retry_utils import ( + RetryConfig, + is_session_invalid_error, + is_throttle_error, + is_transient_error, +) from pyicloud_ipd.asset_version import AssetVersion, calculate_version_filename from pyicloud_ipd.base import PyiCloudService from pyicloud_ipd.exceptions import PyiCloudAPIResponseException from pyicloud_ipd.services.photos import PhotoAsset from pyicloud_ipd.version_size import VersionSize +_RETRY_CONFIG: RetryConfig | None = None +_DOWNLOAD_CHUNK_BYTES: int = 262144 +_VERIFY_SIZE: bool = True +_VERIFY_CHECKSUM: bool = False +_DOWNLOAD_LIMITER: AdaptiveDownloadLimiter | None = None +_METRICS: RunMetrics | None = None +_URL_REFRESH_NEEDED: bool = False + + +def set_retry_config(retry_config: RetryConfig) -> None: + global _RETRY_CONFIG + _RETRY_CONFIG = retry_config + + +def set_download_chunk_bytes(chunk_bytes: int) -> None: + global _DOWNLOAD_CHUNK_BYTES + _DOWNLOAD_CHUNK_BYTES = chunk_bytes + + +def get_download_chunk_bytes() -> int: + return _DOWNLOAD_CHUNK_BYTES + + +def set_download_limiter(limiter: AdaptiveDownloadLimiter | None) -> None: + global _DOWNLOAD_LIMITER + _DOWNLOAD_LIMITER = limiter + + +def get_download_limiter() -> AdaptiveDownloadLimiter | None: + return _DOWNLOAD_LIMITER + + +def set_metrics_collector(metrics: RunMetrics | None) -> None: + global _METRICS + _METRICS = metrics + + +def get_metrics_collector() -> RunMetrics | None: + return _METRICS + + +def consume_url_refresh_needed_signal() -> bool: + global _URL_REFRESH_NEEDED + value = _URL_REFRESH_NEEDED + _URL_REFRESH_NEEDED = False + return value + + +def set_download_verification(*, verify_size: bool, verify_checksum: bool) -> None: + global _VERIFY_SIZE + global _VERIFY_CHECKSUM + _VERIFY_SIZE = verify_size + _VERIFY_CHECKSUM = verify_checksum + + +def get_download_verification() -> tuple[bool, bool]: + return (_VERIFY_SIZE, _VERIFY_CHECKSUM) + + +def get_retry_config() -> RetryConfig: + if _RETRY_CONFIG is not None: + return _RETRY_CONFIG + return RetryConfig( + max_retries=constants.MAX_RETRIES, + backoff_base_seconds=float(constants.WAIT_SECONDS), + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, + ) + def update_mtime(created: datetime.datetime, download_path: str) -> None: """Set the modification time of the downloaded file to the photo creation date""" @@ -79,7 +158,7 @@ def download_response_to_path( ) -> bool: """Saves response content into file with desired created date""" with open(temp_download_path, ("ab" if append_mode else "wb")) as file_obj: - for chunk in response.iter_content(chunk_size=1024): + for chunk in response.iter_content(chunk_size=get_download_chunk_bytes()): if chunk: file_obj.write(chunk) os.rename(temp_download_path, download_path) @@ -87,6 +166,59 @@ def download_response_to_path( return True +def _calculate_digest(path: str, hash_name: str) -> bytes: + hash_obj = hashlib.new(hash_name) + with open(path, "rb") as file_obj: + for chunk in iter(lambda: file_obj.read(1024 * 1024), b""): + hash_obj.update(chunk) + return hash_obj.digest() + + +def _matches_checksum(path: str, expected_checksum: bytes) -> bool: + digest_length_to_algorithm = { + 16: "md5", + 20: "sha1", + 32: "sha256", + 48: "sha384", + 64: "sha512", + } + preferred_algorithm = digest_length_to_algorithm.get(len(expected_checksum)) + if preferred_algorithm is not None: + return _calculate_digest(path, preferred_algorithm) == expected_checksum + + for hash_name in ["md5", "sha1", "sha256", "sha384", "sha512"]: + if _calculate_digest(path, hash_name) == expected_checksum: + return True + return False + + +def verify_download_integrity( + logger: logging.Logger, + download_path: str, + *, + expected_size: int, + expected_checksum: bytes, + verify_size: bool, + verify_checksum: bool, +) -> bool: + if verify_size: + actual_size = os.path.getsize(download_path) + if actual_size != expected_size: + logger.error( + "File size mismatch for %s (expected %d bytes, got %d)", + download_path, + expected_size, + actual_size, + ) + return False + + if verify_checksum and not _matches_checksum(download_path, expected_checksum): + logger.error("Checksum mismatch for %s", download_path) + return False + + return True + + def download_response_to_path_dry_run( logger: logging.Logger, _response: Response, @@ -112,8 +244,12 @@ def download_media( version: AssetVersion, size: VersionSize, filename_builder: Callable[[PhotoAsset], str], + refresh_version: Callable[[], AssetVersion | None] | None = None, ) -> bool: """Download the photo to path, with retries and error handling""" + retry_config = get_retry_config() + limiter = get_download_limiter() + metrics = get_metrics_collector() mkdirs_local = mkdirs_for_path_dry_run if dry_run else mkdirs_for_path if not mkdirs_local(logger, download_path): @@ -129,19 +265,129 @@ def download_media( ) retries = 0 + exhausted_retries = False + download_failed = False + refreshed_url_once = False + range_restart_attempted = False + global _URL_REFRESH_NEEDED + _URL_REFRESH_NEEDED = False + verify_size, verify_checksum = get_download_verification() while True: + retry_after_header: str | None = None try: append_mode = os.path.exists(temp_download_path) current_size = os.path.getsize(temp_download_path) if append_mode else 0 if append_mode: logger.debug(f"Resuming downloading of {download_path} from {current_size}") + if ( + not dry_run + and version.size is not None + and version.size > 0 + and not append_mode + and not has_disk_space_for_download(download_path, version.size) + ): + logger.error( + "Low disk space for %s (required: %d bytes). Skipping download.", + download_path, + version.size, + ) + if metrics is not None: + metrics.on_low_disk() + metrics.on_download_failed() + return False - photo_response = photo.download(icloud.photos.session, version.url, current_size) + if metrics is not None: + metrics.on_download_attempt() + if limiter is None: + photo_response = photo.download(icloud.photos.session, version.url, current_size) + else: + with limiter.slot(): + photo_response = photo.download(icloud.photos.session, version.url, current_size) if photo_response.ok: - return download_local( + if append_mode and photo_response.status_code != 206: + logger.warning( + "Range resume unsupported for %s (HTTP %d). Restarting partial download.", + download_path, + photo_response.status_code, + ) + append_mode = False + + saved = download_local( photo_response, temp_download_path, append_mode, download_path, photo.created ) + if not saved: + return False + + if dry_run: + return True + + if not verify_download_integrity( + logger, + download_path, + expected_size=version.size, + expected_checksum=checksum, + verify_size=verify_size, + verify_checksum=verify_checksum, + ): + try: + os.remove(download_path) + except OSError: + logger.error("Could not remove failed download %s", download_path) + return False + + if limiter is not None: + limiter.on_success() + if metrics is not None: + try: + metrics.on_download_success(os.path.getsize(download_path)) + except OSError: + metrics.on_download_success(0) + return True else: + status_code = str(photo_response.status_code) + if append_mode and status_code == "416": + logger.warning( + "Range resume rejected for %s (HTTP 416). Restarting partial download.", + download_path, + ) + if range_restart_attempted: + break + range_restart_attempted = True + try: + os.remove(temp_download_path) + except OSError: + logger.error("Could not remove stale partial download %s", temp_download_path) + break + continue + if status_code in {"401", "403", "410"}: + if refresh_version is not None and not refreshed_url_once: + logger.info( + "Download URL may be expired for %s. Refreshing asset metadata and retrying once.", + filename_builder(photo), + ) + try: + refreshed = refresh_version() + except Exception: + refreshed = None + if refreshed is not None and refreshed.url and refreshed.url != version.url: + version = refreshed + refreshed_url_once = True + continue + _URL_REFRESH_NEEDED = True + raise PyiCloudAPIResponseException( + f"Download URL expired or denied (HTTP {status_code})", + status_code, + ) + if status_code in {"429", "500", "502", "503", "504"}: + retry_after_header = photo_response.headers.get("Retry-After") + if status_code == "429" and limiter is not None: + limiter.on_throttle() + if status_code == "429" and metrics is not None: + metrics.on_throttle() + raise PyiCloudAPIResponseException( + f"Download request failed with HTTP {status_code}", + status_code, + ) # Use the standard original filename generator for error logging from icloudpd.base import lp_filename_original as simple_lp_filename_generator @@ -158,29 +404,41 @@ def download_media( break except PyiCloudAPIResponseException as ex: - if "Invalid global session" in str(ex): + download_failed = True + if is_session_invalid_error(ex): logger.error("Session error, re-authenticating...") - if retries > 0: - # If the first re-authentication attempt failed, - # start waiting a few seconds before retrying in case - # there are some issues with the Apple servers - time.sleep(constants.WAIT_SECONDS) - icloud.authenticate() else: - # short circuiting 0 retries - if retries == constants.MAX_RETRIES: + if limiter is not None and is_throttle_error(ex): + limiter.on_throttle() + if not is_transient_error(ex): break - # you end up here when p.e. throttling by Apple happens - wait_time = (retries + 1) * constants.WAIT_SECONDS - # Get the proper filename for error messages - error_filename = filename_builder(photo) - logger.error( - "Error downloading %s, retrying after %s seconds...", error_filename, wait_time - ) - time.sleep(wait_time) + # short circuiting 0 retries + if retries >= retry_config.max_retries: + exhausted_retries = True + break + + retries += 1 + if metrics is not None: + metrics.on_retry() + wait_time = retry_config.next_delay_seconds( + retries, + retry_after=retry_after_header, + throttle_error=is_throttle_error(ex), + ) + error_filename = filename_builder(photo) + logger.error( + "Error downloading %s, retrying after %.1f seconds... (%d/%d)", + error_filename, + wait_time, + retries, + retry_config.max_retries, + ) + time.sleep(wait_time) + continue except OSError: + download_failed = True logger.error( "IOError while writing file to %s. " + "You might have run out of disk space, or the file " @@ -189,10 +447,10 @@ def download_media( download_path, ) break - retries = retries + 1 - if retries >= constants.MAX_RETRIES: - break - if retries >= constants.MAX_RETRIES: + + if exhausted_retries or download_failed: + if metrics is not None: + metrics.on_download_failed() # Get the proper filename for error messages error_filename = filename_builder(photo) logger.error( @@ -201,3 +459,12 @@ def download_media( ) return False + + +def has_disk_space_for_download(download_path: str, required_bytes: int, reserve_bytes: int = 50 * 1024 * 1024) -> bool: + target_dir = os.path.dirname(download_path) or "." + try: + free_bytes = shutil.disk_usage(target_dir).free + except OSError: + return True + return free_bytes >= (required_bytes + reserve_bytes) diff --git a/src/icloudpd/limiter.py b/src/icloudpd/limiter.py new file mode 100644 index 000000000..d654c661e --- /dev/null +++ b/src/icloudpd/limiter.py @@ -0,0 +1,137 @@ +"""Adaptive limiter for download request concurrency.""" + +from __future__ import annotations + +import threading +import time +from contextlib import contextmanager +from typing import Callable, Iterator + + +class AdaptiveDownloadLimiter: + def __init__( + self, + *, + max_workers: int, + cooldown_seconds: float, + min_workers: int = 1, + increase_every: int = 10, + clock: Callable[[], float] = time.monotonic, + ) -> None: + if max_workers < 1: + raise ValueError("max_workers must be >= 1") + if min_workers < 1: + raise ValueError("min_workers must be >= 1") + if min_workers > max_workers: + raise ValueError("min_workers must be <= max_workers") + if increase_every < 1: + raise ValueError("increase_every must be >= 1") + if cooldown_seconds < 0: + raise ValueError("cooldown_seconds must be >= 0") + + self._max_workers = max_workers + self._min_workers = min_workers + self._cooldown_seconds = cooldown_seconds + self._increase_every = increase_every + self._clock = clock + + self._condition = threading.Condition() + self._in_flight = 0 + self._current_limit = max_workers + self._success_streak = 0 + self._cooldown_until = 0.0 + self._stopped = False + + @property + def max_workers(self) -> int: + return self._max_workers + + @property + def current_limit(self) -> int: + with self._condition: + return self._current_limit + + @property + def cooldown_remaining_seconds(self) -> float: + with self._condition: + return max(0.0, self._cooldown_until - self._clock()) + + def acquire(self, timeout: float | None = None) -> bool: + deadline = None if timeout is None else (self._clock() + timeout) + with self._condition: + while True: + if self._stopped: + return False + now = self._clock() + cooldown_remaining = max(0.0, self._cooldown_until - now) + can_enter = cooldown_remaining <= 0 and self._in_flight < self._current_limit + if can_enter: + self._in_flight += 1 + return True + + wait_for = cooldown_remaining if cooldown_remaining > 0 else None + if deadline is not None: + remaining = deadline - now + if remaining <= 0: + return False + wait_for = remaining if wait_for is None else min(wait_for, remaining) + self._condition.wait(wait_for) + + def release(self) -> None: + with self._condition: + if self._in_flight > 0: + self._in_flight -= 1 + self._condition.notify_all() + + def stop(self, wait: bool = True, timeout: float | None = None) -> bool: + deadline = None if timeout is None else (self._clock() + timeout) + with self._condition: + self._stopped = True + self._condition.notify_all() + if not wait: + return True + while self._in_flight > 0: + if deadline is not None: + remaining = deadline - self._clock() + if remaining <= 0: + return False + self._condition.wait(remaining) + else: + self._condition.wait() + return True + + def start(self) -> None: + with self._condition: + self._stopped = False + self._condition.notify_all() + + def on_success(self) -> None: + with self._condition: + if self._clock() < self._cooldown_until: + return + if self._current_limit >= self._max_workers: + self._success_streak = 0 + return + self._success_streak += 1 + if self._success_streak >= self._increase_every: + self._current_limit += 1 + self._success_streak = 0 + self._condition.notify_all() + + def on_throttle(self) -> None: + with self._condition: + self._success_streak = 0 + self._current_limit = max(self._min_workers, self._current_limit // 2) + self._cooldown_until = max( + self._cooldown_until, self._clock() + self._cooldown_seconds + ) + self._condition.notify_all() + + @contextmanager + def slot(self, timeout: float | None = None) -> Iterator[None]: + if not self.acquire(timeout=timeout): + raise TimeoutError("Could not acquire download slot within timeout") + try: + yield + finally: + self.release() diff --git a/src/icloudpd/metrics.py b/src/icloudpd/metrics.py new file mode 100644 index 000000000..0744e69b2 --- /dev/null +++ b/src/icloudpd/metrics.py @@ -0,0 +1,90 @@ +"""Run metrics collection and export helpers.""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import dataclass + + +@dataclass +class RunMetrics: + username: str + run_mode: str = "legacy_stateless" + started_at_epoch: float = 0.0 + finished_at_epoch: float = 0.0 + assets_considered: int = 0 + downloads_attempted: int = 0 + downloads_succeeded: int = 0 + downloads_failed: int = 0 + bytes_downloaded: int = 0 + retries: int = 0 + throttle_events: int = 0 + low_disk_events: int = 0 + queue_depth_max: int = 0 + queue_depth_last: int = 0 + + def start(self) -> None: + self.started_at_epoch = time.time() + + def finish(self) -> None: + self.finished_at_epoch = time.time() + + def on_asset_considered(self) -> None: + self.assets_considered += 1 + + def on_download_attempt(self) -> None: + self.downloads_attempted += 1 + + def on_download_success(self, bytes_written: int) -> None: + self.downloads_succeeded += 1 + self.bytes_downloaded += max(0, bytes_written) + + def on_download_failed(self) -> None: + self.downloads_failed += 1 + + def on_retry(self) -> None: + self.retries += 1 + + def on_throttle(self) -> None: + self.throttle_events += 1 + + def on_low_disk(self) -> None: + self.low_disk_events += 1 + + def set_queue_depth(self, depth: int) -> None: + normalized = max(0, depth) + self.queue_depth_last = normalized + if normalized > self.queue_depth_max: + self.queue_depth_max = normalized + + def snapshot(self) -> dict[str, float | int | str]: + elapsed = max(0.0, self.finished_at_epoch - self.started_at_epoch) + return { + "username": self.username, + "run_mode": self.run_mode, + "started_at_epoch": self.started_at_epoch, + "finished_at_epoch": self.finished_at_epoch, + "elapsed_seconds": elapsed, + "assets_considered": self.assets_considered, + "downloads_attempted": self.downloads_attempted, + "downloads_succeeded": self.downloads_succeeded, + "downloads_failed": self.downloads_failed, + "success_gap": self.downloads_attempted - self.downloads_succeeded, + "bytes_downloaded": self.bytes_downloaded, + "throughput_downloads_per_sec": ( + self.downloads_succeeded / elapsed if elapsed > 0 else 0.0 + ), + "retries": self.retries, + "throttle_events": self.throttle_events, + "queue_depth_max": self.queue_depth_max, + "queue_depth_last": self.queue_depth_last, + "low_disk_events": self.low_disk_events, + } + + +def write_metrics_json(path: str, payload: dict[str, object]) -> None: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=True, indent=2) diff --git a/src/icloudpd/retry_utils.py b/src/icloudpd/retry_utils.py new file mode 100644 index 000000000..0828df264 --- /dev/null +++ b/src/icloudpd/retry_utils.py @@ -0,0 +1,141 @@ +"""Shared retry policy and error classification utilities.""" + +from __future__ import annotations + +import datetime +import random +from dataclasses import dataclass +from email.utils import parsedate_to_datetime + +from requests.exceptions import ( + ChunkedEncodingError, + ContentDecodingError, + StreamConsumedError, + UnrewindableBodyError, +) + +from pyicloud_ipd.exceptions import ( + PyiCloud2SARequiredException, + PyiCloudAPIResponseException, + PyiCloudConnectionErrorException, + PyiCloudFailedLoginException, + PyiCloudFailedMFAException, + PyiCloudNoStoredPasswordAvailableException, + PyiCloudServiceNotActivatedException, + PyiCloudServiceUnavailableException, +) + +_TRANSIENT_CODES = {"421", "429", "450", "500", "502", "503", "504", "ACCESS_DENIED"} +_THROTTLE_CODES = {"429", "ACCESS_DENIED"} +_SESSION_INVALID_MARKER = "invalid global session" + + +@dataclass(frozen=True) +class RetryConfig: + max_retries: int + backoff_base_seconds: float + backoff_max_seconds: float + respect_retry_after: bool + throttle_cooldown_seconds: float + jitter_fraction: float = 0.1 + + def next_delay_seconds( + self, + retry_number: int, + *, + retry_after: str | None = None, + throttle_error: bool = False, + ) -> float: + retry_after_seconds = ( + parse_retry_after_seconds(retry_after) + if self.respect_retry_after and retry_after is not None + else None + ) + if retry_after_seconds is not None: + delay = retry_after_seconds + else: + delay = self.backoff_base_seconds * (2 ** max(0, retry_number - 1)) + + delay = min(delay, self.backoff_max_seconds) + if throttle_error: + delay = max(delay, self.throttle_cooldown_seconds) + + jitter_upper = max(0.0, delay * self.jitter_fraction) + if jitter_upper > 0: + delay += random.uniform(0.0, jitter_upper) + return min(delay, self.backoff_max_seconds) + + +def parse_retry_after_seconds(retry_after: str | None) -> float | None: + if retry_after is None: + return None + candidate = retry_after.strip() + if not candidate: + return None + + try: + return max(0.0, float(candidate)) + except ValueError: + pass + + try: + parsed = parsedate_to_datetime(candidate) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=datetime.timezone.utc) + now = datetime.datetime.now(tz=datetime.timezone.utc) + return max(0.0, (parsed - now).total_seconds()) + except (TypeError, ValueError): + return None + + +def is_session_invalid_error(error: Exception) -> bool: + if not isinstance(error, PyiCloudAPIResponseException): + return False + return _SESSION_INVALID_MARKER in str(error).lower() + + +def is_throttle_error(error: Exception) -> bool: + if not isinstance(error, PyiCloudAPIResponseException): + return False + code = str(error.code or "").upper() + message = str(error).lower() + return code in _THROTTLE_CODES or "throttl" in message or "rate limit" in message + + +def is_fatal_auth_config_error(error: Exception) -> bool: + return isinstance( + error, + ( + PyiCloudFailedLoginException, + PyiCloudFailedMFAException, + PyiCloud2SARequiredException, + PyiCloudNoStoredPasswordAvailableException, + PyiCloudServiceNotActivatedException, + ), + ) + + +def is_transient_error(error: Exception) -> bool: + if is_fatal_auth_config_error(error): + return False + if isinstance( + error, + ( + ChunkedEncodingError, + ContentDecodingError, + StreamConsumedError, + UnrewindableBodyError, + ), + ): + return True + if isinstance(error, (PyiCloudServiceUnavailableException, PyiCloudConnectionErrorException)): + return True + if not isinstance(error, PyiCloudAPIResponseException): + return False + + if is_session_invalid_error(error): + return True + code = str(error.code or "").upper() + if code in _TRANSIENT_CODES: + return True + return "timed out" in str(error).lower() diff --git a/src/icloudpd/state_db.py b/src/icloudpd/state_db.py new file mode 100644 index 000000000..66b1f822a --- /dev/null +++ b/src/icloudpd/state_db.py @@ -0,0 +1,431 @@ +"""State DB support for resumable processing.""" + +from __future__ import annotations + +import os +import sqlite3 +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pyicloud_ipd.services.photos import PhotoAsset + + +def resolve_state_db_path(state_db_option: str | None, cookie_directory: str) -> str | None: + if state_db_option is None: + return None + if state_db_option == "auto": + return os.path.join(os.path.expanduser(cookie_directory), "icloudpd.sqlite") + return os.path.expanduser(state_db_option) + + +def initialize_state_db(db_path: str) -> None: + os.makedirs(os.path.dirname(db_path), exist_ok=True) + + with sqlite3.connect(db_path) as conn: + conn.execute("PRAGMA journal_mode=WAL;") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS assets ( + asset_id TEXT NOT NULL, + library TEXT NOT NULL, + album TEXT NOT NULL, + added_date TEXT, + asset_date TEXT, + item_type TEXT, + metadata_json TEXT, + PRIMARY KEY (asset_id, library, album) + ) + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS tasks ( + asset_id TEXT NOT NULL, + library TEXT NOT NULL, + album TEXT NOT NULL, + version TEXT NOT NULL, + expected_size INTEGER, + checksum TEXT, + url TEXT, + local_path TEXT, + status TEXT NOT NULL DEFAULT 'pending', + attempts INTEGER NOT NULL DEFAULT 0, + lease_owner TEXT, + lease_expires_at TEXT, + last_error TEXT, + checksum_result TEXT, + needs_url_refresh INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (asset_id, library, album, version) + ) + """ + ) + _ensure_column(conn, "tasks", "checksum_result", "TEXT") + _ensure_column(conn, "tasks", "needs_url_refresh", "INTEGER NOT NULL DEFAULT 0") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS checkpoints ( + library TEXT NOT NULL, + album TEXT NOT NULL, + start_rank INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (library, album) + ) + """ + ) + + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_tasks_status_updated ON tasks(status, updated_at)" + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_tasks_lease_expires ON tasks(lease_expires_at)" + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_tasks_library_album_status ON tasks(library, album, status)" + ) + + +def utc_now_iso() -> str: + return datetime.now(tz=timezone.utc).isoformat() + + +def _ensure_column(conn: sqlite3.Connection, table: str, column: str, definition: str) -> None: + existing = {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + if column not in existing: + conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}") + + +def enqueue_task( + db_path: str, + *, + asset_id: str, + library: str, + album: str, + version: str, + expected_size: int | None, + checksum: str | None, + url: str | None, + local_path: str | None, +) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + INSERT INTO tasks ( + asset_id, library, album, version, expected_size, checksum, url, local_path, + status, attempts, needs_url_refresh, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', 0, 0, CURRENT_TIMESTAMP) + ON CONFLICT(asset_id, library, album, version) DO UPDATE SET + expected_size=excluded.expected_size, + checksum=excluded.checksum, + url=excluded.url, + local_path=excluded.local_path, + needs_url_refresh=0, + updated_at=CURRENT_TIMESTAMP + """, + (asset_id, library, album, version, expected_size, checksum, url, local_path), + ) + + +def requeue_stale_leases(db_path: str, now_iso: str | None = None) -> int: + now = now_iso or utc_now_iso() + with sqlite3.connect(db_path) as conn: + cursor = conn.execute( + """ + UPDATE tasks + SET status='pending', + lease_owner=NULL, + lease_expires_at=NULL, + updated_at=CURRENT_TIMESTAMP + WHERE status='in_progress' AND lease_expires_at IS NOT NULL AND lease_expires_at < ? + """, + (now,), + ) + return cursor.rowcount + + +def requeue_in_progress_tasks(db_path: str, lease_owner: str | None = None) -> int: + with sqlite3.connect(db_path) as conn: + if lease_owner is None: + cursor = conn.execute( + """ + UPDATE tasks + SET status='pending', + lease_owner=NULL, + lease_expires_at=NULL, + updated_at=CURRENT_TIMESTAMP + WHERE status='in_progress' + """ + ) + else: + cursor = conn.execute( + """ + UPDATE tasks + SET status='pending', + lease_owner=NULL, + lease_expires_at=NULL, + updated_at=CURRENT_TIMESTAMP + WHERE status='in_progress' AND lease_owner=? + """, + (lease_owner,), + ) + return cursor.rowcount + + +def prune_completed_tasks(db_path: str, *, older_than_days: int) -> int: + if older_than_days <= 0: + raise ValueError("older_than_days must be greater than 0") + cutoff = f"-{older_than_days} days" + with sqlite3.connect(db_path) as conn: + cursor = conn.execute( + """ + DELETE FROM tasks + WHERE status IN ('done', 'failed') + AND updated_at < datetime('now', ?) + """, + (cutoff,), + ) + return cursor.rowcount + + +def checkpoint_wal(db_path: str, *, mode: str = "PASSIVE") -> tuple[int, int, int]: + with sqlite3.connect(db_path) as conn: + row = conn.execute(f"PRAGMA wal_checkpoint({mode})").fetchone() + if row is None: + return (0, 0, 0) + return (int(row[0]), int(row[1]), int(row[2])) + + +def vacuum_state_db(db_path: str) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute("VACUUM") + + +def lease_next_task( + db_path: str, *, lease_owner: str, lease_seconds: int = 300, now_iso: str | None = None +) -> tuple[str, str, str, str] | None: + now_dt = datetime.fromisoformat(now_iso) if now_iso else datetime.now(tz=timezone.utc) + lease_expires = (now_dt + timedelta(seconds=lease_seconds)).isoformat() + now = now_dt.isoformat() + with sqlite3.connect(db_path) as conn: + conn.execute("BEGIN IMMEDIATE") + conn.execute( + """ + UPDATE tasks + SET status='pending', + lease_owner=NULL, + lease_expires_at=NULL, + updated_at=CURRENT_TIMESTAMP + WHERE status='in_progress' AND lease_expires_at IS NOT NULL AND lease_expires_at < ? + """, + (now,), + ) + row = conn.execute( + """ + SELECT asset_id, library, album, version + FROM tasks + WHERE status='pending' + ORDER BY updated_at ASC + LIMIT 1 + """ + ).fetchone() + if row is None: + conn.commit() + return None + conn.execute( + """ + UPDATE tasks + SET status='in_progress', + lease_owner=?, + lease_expires_at=?, + attempts=attempts+1, + updated_at=CURRENT_TIMESTAMP + WHERE asset_id=? AND library=? AND album=? AND version=? + """, + (lease_owner, lease_expires, row[0], row[1], row[2], row[3]), + ) + conn.commit() + return (row[0], row[1], row[2], row[3]) + + +def mark_task_done(db_path: str, *, asset_id: str, library: str, album: str, version: str) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + UPDATE tasks + SET status='done', + lease_owner=NULL, + lease_expires_at=NULL, + last_error=NULL, + updated_at=CURRENT_TIMESTAMP + WHERE asset_id=? AND library=? AND album=? AND version=? + """, + (asset_id, library, album, version), + ) + + +def record_asset_checksum_result( + db_path: str, + *, + asset_id: str, + library: str, + album: str, + checksum_result: str, +) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + UPDATE tasks + SET checksum_result=?, + updated_at=CURRENT_TIMESTAMP + WHERE asset_id=? AND library=? AND album=? + """, + (checksum_result, asset_id, library, album), + ) + + +def mark_task_failed( + db_path: str, *, asset_id: str, library: str, album: str, version: str, error: str +) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + UPDATE tasks + SET status='failed', + lease_owner=NULL, + lease_expires_at=NULL, + last_error=?, + updated_at=CURRENT_TIMESTAMP + WHERE asset_id=? AND library=? AND album=? AND version=? + """, + (error, asset_id, library, album, version), + ) + + +def mark_asset_tasks_need_url_refresh( + db_path: str, + *, + asset_id: str, + library: str, + album: str, + error: str = "expired_download_url", +) -> int: + with sqlite3.connect(db_path) as conn: + cursor = conn.execute( + """ + UPDATE tasks + SET needs_url_refresh=1, + last_error=?, + updated_at=CURRENT_TIMESTAMP + WHERE asset_id=? AND library=? AND album=? + """, + (error, asset_id, library, album), + ) + return cursor.rowcount + + +def clear_asset_tasks_need_url_refresh( + db_path: str, + *, + asset_id: str, + library: str, + album: str, +) -> int: + with sqlite3.connect(db_path) as conn: + cursor = conn.execute( + """ + UPDATE tasks + SET needs_url_refresh=0, + updated_at=CURRENT_TIMESTAMP + WHERE asset_id=? AND library=? AND album=? + """, + (asset_id, library, album), + ) + return cursor.rowcount + + +def save_checkpoint(db_path: str, *, library: str, album: str, start_rank: int) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + INSERT INTO checkpoints (library, album, start_rank, created_at, updated_at) + VALUES (?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + ON CONFLICT(library, album) DO UPDATE SET + start_rank=excluded.start_rank, + updated_at=CURRENT_TIMESTAMP + """, + (library, album, start_rank), + ) + + +def load_checkpoint(db_path: str, *, library: str, album: str) -> int | None: + with sqlite3.connect(db_path) as conn: + row = conn.execute( + "SELECT start_rank FROM checkpoints WHERE library=? AND album=?", + (library, album), + ).fetchone() + if row is None: + return None + return int(row[0]) + + +def upsert_asset( + db_path: str, + *, + asset_id: str, + library: str, + album: str, + added_date: str | None, + asset_date: str | None, + item_type: str | None, + metadata_json: str | None, +) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + INSERT INTO assets ( + asset_id, library, album, added_date, asset_date, item_type, metadata_json + ) VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(asset_id, library, album) DO UPDATE SET + added_date=excluded.added_date, + asset_date=excluded.asset_date, + item_type=excluded.item_type, + metadata_json=excluded.metadata_json + """, + (asset_id, library, album, added_date, asset_date, item_type, metadata_json), + ) + + +def upsert_asset_tasks(db_path: str, *, photo: PhotoAsset, library: str, album: str) -> None: + item_type = photo.item_type.value if photo.item_type is not None else None + try: + added_date = photo.added_date.isoformat() + except Exception: + added_date = None + try: + asset_date = photo.asset_date.isoformat() + except Exception: + asset_date = None + upsert_asset( + db_path, + asset_id=photo.id, + library=library, + album=album, + added_date=added_date, + asset_date=asset_date, + item_type=item_type, + metadata_json=None, + ) + for version, version_value in photo.versions.items(): + enqueue_task( + db_path, + asset_id=photo.id, + library=library, + album=album, + version=str(version.value), + expected_size=version_value.size, + checksum=version_value.checksum, + url=version_value.url, + local_path=None, + ) diff --git a/src/pyicloud_ipd/session.py b/src/pyicloud_ipd/session.py index 23791c0d4..a30fd60aa 100644 --- a/src/pyicloud_ipd/session.py +++ b/src/pyicloud_ipd/session.py @@ -1,6 +1,9 @@ import inspect import json import logging +import os +import tempfile +import threading import typing from typing import Any, Callable, Dict, Mapping, NoReturn, Sequence @@ -16,6 +19,8 @@ from pyicloud_ipd.utils import handle_connection_error, throw_on_503 LOGGER = logging.getLogger(__name__) +_PATH_LOCKS: dict[str, threading.Lock] = {} +_PATH_LOCKS_GUARD = threading.Lock() HEADER_DATA = { "X-Apple-ID-Account-Country": "account_country", @@ -88,14 +93,12 @@ def request(self, method: str, url, **kwargs): # type: ignore session_arg = value self.service.session_data.update({session_arg: response.headers.get(header)}) - # Save session_data to file - with open(self.service.session_path, "w", encoding="utf-8") as outfile: - json.dump(self.service.session_data, outfile) - LOGGER.debug("Saved session data to file") - - # Save cookies to file - self.cookies.save(ignore_discard=True, ignore_expires=True) # type: ignore[attr-defined] - LOGGER.debug("Cookies saved to %s", self.service.cookiejar_path) + persist_session_and_cookies( + self.service.session_path, + self.service.cookiejar_path, + self.service.session_data, + self.cookies, + ) if not response.ok and ( content_type not in json_mimetypes or response.status_code in [421, 450, 500] @@ -174,3 +177,61 @@ def _raise_error(self, code: str, reason: str) -> NoReturn: api_error = PyiCloudAPIResponseException(reason, code) LOGGER.error(api_error) raise api_error + + +def _lock_for_path(path: str) -> threading.Lock: + with _PATH_LOCKS_GUARD: + lock = _PATH_LOCKS.get(path) + if lock is None: + lock = threading.Lock() + _PATH_LOCKS[path] = lock + return lock + + +def _atomic_write_json(path: str, payload: Mapping[str, Any]) -> None: + directory = os.path.dirname(path) or "." + os.makedirs(directory, exist_ok=True) + fd, temp_path = tempfile.mkstemp(prefix=".icloudpd-session-", dir=directory, text=True) + try: + with os.fdopen(fd, "w", encoding="utf-8") as outfile: + json.dump(payload, outfile) + outfile.flush() + os.fsync(outfile.fileno()) + os.replace(temp_path, path) + finally: + if os.path.exists(temp_path): + os.remove(temp_path) + + +def _atomic_save_cookiejar(cookiejar: Any, cookiejar_path: str) -> None: + directory = os.path.dirname(cookiejar_path) or "." + os.makedirs(directory, exist_ok=True) + fd, temp_path = tempfile.mkstemp(prefix=".icloudpd-cookiejar-", dir=directory, text=True) + os.close(fd) + try: + cookiejar.save( + filename=temp_path, + ignore_discard=True, + ignore_expires=True, + ) + os.replace(temp_path, cookiejar_path) + finally: + if os.path.exists(temp_path): + os.remove(temp_path) + + +def persist_session_and_cookies( + session_path: str, cookiejar_path: str, session_data: Mapping[str, Any], cookies: Any +) -> None: + lock_paths = sorted({os.path.abspath(session_path), os.path.abspath(cookiejar_path)}) + locks = [_lock_for_path(path) for path in lock_paths] + for lock in locks: + lock.acquire() + try: + _atomic_write_json(session_path, session_data) + LOGGER.debug("Saved session data to file") + _atomic_save_cookiejar(cookies, cookiejar_path) + LOGGER.debug("Cookies saved to %s", cookiejar_path) + finally: + for lock in reversed(locks): + lock.release() diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index 5650440ff..48d32d1ac 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -3,6 +3,7 @@ import os import shutil import sys +import time import traceback from contextlib import redirect_stderr, redirect_stdout from functools import partial @@ -13,6 +14,11 @@ from foundation.core import compose, flip, partial_1_1, partial_2_1 from icloudpd.cli import cli +# Keep tests deterministic across environments that run with non-UTC local time. +os.environ["TZ"] = "UTC" +if hasattr(time, "tzset"): + time.tzset() + vcr = VCR(decode_compressed_response=True, record_mode="none") @@ -115,7 +121,11 @@ def assert_files( ) -DEFAULT_ENV: Mapping[str, str | None] = {"CLIENT_ID": "DE309E26-942E-11E8-92F5-14109FE0B321"} +DEFAULT_ENV: Mapping[str, str | None] = { + "CLIENT_ID": "DE309E26-942E-11E8-92F5-14109FE0B321", + "HOME": "/tmp/icloudpd-tests-home", + "TZ": "UTC", +} def clean_boolean_args(params: Sequence[str]) -> list[str]: @@ -170,6 +180,10 @@ def run_main_env( ) -> TestResult: """Run the new argparse-based CLI with environment variables""" + home_dir = env.get("HOME") + if home_dir: + os.makedirs(home_dir, exist_ok=True) + # Set environment variables original_env = {} for key, value in env.items(): @@ -179,6 +193,8 @@ def run_main_env( del os.environ[key] else: os.environ[key] = value + if hasattr(time, "tzset"): + time.tzset() # Capture stdout and stderr stdout_capture = io.StringIO() @@ -263,6 +279,8 @@ def readline(self, size: int = -1) -> str: # type: ignore[override] del os.environ[key] else: os.environ[key] = original_value + if hasattr(time, "tzset"): + time.tzset() # Clean the output to remove log prefixes for compatibility with old tests raw_output = stdout_capture.getvalue() @@ -278,16 +296,10 @@ def readline(self, size: int = -1) -> str: # type: ignore[override] # Create both original and cleaned output cleaned_output = "\n".join(cleaned_lines) - # For compatibility with old tests, adjust exit codes for specific error conditions - adjusted_exit_code = exit_code - if exit_code == 1 and "Invalid email/password combination" in raw_output: - # Authentication failure - old tests expect exit code 2 - adjusted_exit_code = 2 - # For compatibility, provide the cleaned output as the primary output # but keep raw_output available if needed result = TestResult( - exit_code=adjusted_exit_code, + exit_code=exit_code, output=cleaned_output, exception=exception, stderr=stderr_capture.getvalue(), diff --git a/tests/test_alerts.py b/tests/test_alerts.py new file mode 100644 index 000000000..8c818166e --- /dev/null +++ b/tests/test_alerts.py @@ -0,0 +1,26 @@ +import logging +from unittest import TestCase +from unittest.mock import patch + +from icloudpd.base import emit_throttle_alert_if_needed +from icloudpd.metrics import RunMetrics + + +class AlertsTestCase(TestCase): + def test_emit_throttle_alert_below_threshold(self) -> None: + logger = logging.getLogger("icloudpd-test-alerts") + metrics = RunMetrics(username="u1") + metrics.throttle_events = 2 + with patch.object(logger, "warning") as warning_mock: + emitted = emit_throttle_alert_if_needed(logger, metrics, threshold=3) + self.assertFalse(emitted) + warning_mock.assert_not_called() + + def test_emit_throttle_alert_at_threshold(self) -> None: + logger = logging.getLogger("icloudpd-test-alerts") + metrics = RunMetrics(username="u1") + metrics.throttle_events = 3 + with patch.object(logger, "warning") as warning_mock: + emitted = emit_throttle_alert_if_needed(logger, metrics, threshold=3) + self.assertTrue(emitted) + warning_mock.assert_called_once() diff --git a/tests/test_cli.py b/tests/test_cli.py index 7bc6996b1..a65fcf219 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,13 +1,14 @@ import datetime import inspect +import json import os import shutil -import zoneinfo from argparse import ArgumentError from typing import Sequence, Tuple from unittest import TestCase import pytest +from tzlocal import get_localzone from icloudpd.cli import format_help, parse from icloudpd.config import GlobalConfig, UserConfig @@ -90,6 +91,11 @@ def test_cli_parser(self) -> None: PasswordProvider.CONSOLE, ], mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, ), [], ), @@ -114,6 +120,11 @@ def test_cli_parser(self) -> None: PasswordProvider.CONSOLE, ], mfa_provider=MFAProvider.WEBUI, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, ), [], ), @@ -143,6 +154,11 @@ def test_cli_parser(self) -> None: watch_with_interval=None, password_providers=[PasswordProvider.WEBUI, PasswordProvider.CONSOLE], mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, ), [], ), @@ -167,11 +183,110 @@ def test_cli_parser(self) -> None: PasswordProvider.CONSOLE, ], mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, ), [], ), "--version --use-os-locale", ) + self.assertEqual( + parse( + [ + "--log-format", + "json", + "--metrics-json", + "/tmp/metrics.json", + "--max-retries", + "7", + "--backoff-base-seconds", + "2", + "--backoff-max-seconds", + "120", + "--no-respect-retry-after", + "--throttle-cooldown-seconds", + "30", + ] + ), + ( + GlobalConfig( + help=False, + version=False, + use_os_locale=False, + only_print_filenames=False, + log_level=LogLevel.DEBUG, + log_format="json", + no_progress_bar=False, + threads_num=1, + domain="com", + watch_with_interval=None, + password_providers=[ + PasswordProvider.PARAMETER, + PasswordProvider.KEYRING, + PasswordProvider.CONSOLE, + ], + mfa_provider=MFAProvider.CONSOLE, + max_retries=7, + backoff_base_seconds=2.0, + backoff_max_seconds=120.0, + respect_retry_after=False, + throttle_cooldown_seconds=30.0, + metrics_json="/tmp/metrics.json", + ), + [], + ), + "retry options parse", + ) + _, users_with_chunk_size = parse( + ["--directory", "abc", "--username", "u1", "--download-chunk-bytes", "65536"] + ) + self.assertEqual(users_with_chunk_size[0].download_chunk_bytes, 65536) + _, users_with_download_workers = parse( + ["--directory", "abc", "--username", "u1", "--download-workers", "6"] + ) + self.assertEqual(users_with_download_workers[0].download_workers, 6) + _, users_with_verification = parse( + [ + "--directory", + "abc", + "--username", + "u1", + "--no-verify-size", + "--verify-checksum", + ] + ) + self.assertFalse(users_with_verification[0].verify_size) + self.assertTrue(users_with_verification[0].verify_checksum) + _, users_with_paging_options = parse( + [ + "--directory", + "abc", + "--username", + "u1", + "--album-page-size", + "250", + "--no-remote-count", + ] + ) + self.assertEqual(users_with_paging_options[0].album_page_size, 250) + self.assertTrue(users_with_paging_options[0].no_remote_count) + _, users_with_added_date_filters = parse( + [ + "--directory", + "abc", + "--username", + "u1", + "--skip-added-before", + "2025-01-01", + "--skip-added-after", + "2d", + ] + ) + self.assertIsNotNone(users_with_added_date_filters[0].skip_added_before) + self.assertIsNotNone(users_with_added_date_filters[0].skip_added_after) self.assertEqual( parse( ["--directory", "abc", "--username", "u1", "--username", "u2", "--directory", "def"] @@ -193,6 +308,11 @@ def test_cli_parser(self) -> None: PasswordProvider.CONSOLE, ], mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, ), [ UserConfig( @@ -309,6 +429,11 @@ def test_cli_parser(self) -> None: PasswordProvider.CONSOLE, ], mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, ), [ UserConfig( @@ -348,7 +473,7 @@ def test_cli_parser(self) -> None: align_raw=RawTreatmentPolicy.AS_IS, file_match_policy=FileMatchPolicy.NAME_SIZE_DEDUP_WITH_SUFFIX, skip_created_before=datetime.datetime( - year=2025, month=1, day=2, tzinfo=zoneinfo.ZoneInfo(key="Etc/UTC") + year=2025, month=1, day=2, tzinfo=get_localzone() ), skip_created_after=datetime.timedelta(days=2), skip_photos=False, @@ -389,6 +514,29 @@ def test_cli_parser(self) -> None: "2", ] ) + _, users_with_auto_state_db = parse( + ["--directory", "abc", "--username", "u1", "--state-db"] + ) + self.assertEqual(users_with_auto_state_db[0].state_db, "auto") + + _, users_with_custom_state_db = parse( + ["--directory", "abc", "--username", "u1", "--state-db", "/tmp/icloudpd.sqlite"] + ) + self.assertEqual(users_with_custom_state_db[0].state_db, "/tmp/icloudpd.sqlite") + _, users_with_state_db_maintenance = parse( + [ + "--directory", + "abc", + "--username", + "u1", + "--state-db", + "--state-db-prune-completed-days", + "30", + "--state-db-vacuum", + ] + ) + self.assertEqual(users_with_state_db_maintenance[0].state_db_prune_completed_days, 30) + self.assertTrue(users_with_state_db_maintenance[0].state_db_vacuum) def test_cli(self) -> None: result = run_main(["--help"]) @@ -551,3 +699,84 @@ def test_conflict_options_delete_after_download_and_keep_icloud_recent_days(self self.assertEqual(result.exit_code, 2, "exit code") self.assertFalse(os.path.exists(base_dir), f"{base_dir} exists") + + def test_invalid_album_page_size(self) -> None: + result = run_main( + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "-d", + "/tmp", + "--album-page-size", + "0", + ], + ) + self.assertEqual(result.exit_code, 2, "exit code") + + def test_invalid_download_workers(self) -> None: + result = run_main( + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "-d", + "/tmp", + "--download-workers", + "0", + ], + ) + self.assertEqual(result.exit_code, 2, "exit code") + + def test_invalid_state_db_prune_completed_days(self) -> None: + result = run_main( + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "-d", + "/tmp", + "--state-db-prune-completed-days", + "0", + ], + ) + self.assertEqual(result.exit_code, 2, "exit code") + + def test_metrics_json_export(self) -> None: + base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) + metrics_path = os.path.join(base_dir, "metrics.json") + if os.path.exists(base_dir): + shutil.rmtree(base_dir) + os.makedirs(base_dir, exist_ok=True) + + _data_dir, result = run_icloudpd_test( + self.assertEqual, + self.root_path, + base_dir, + "listing_photos.yml", + [], + [], + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "--auth-only", + "--metrics-json", + metrics_path, + ], + ) + self.assertEqual(result.exit_code, 0, "exit code") + self.assertTrue(os.path.exists(metrics_path)) + with open(metrics_path, encoding="utf-8") as f: + payload = json.load(f) + self.assertIn("exit_code", payload) + self.assertIn("status", payload) + self.assertEqual(payload["exit_code"], 0) + self.assertIn(payload["status"], ["success", "partial_success"]) + self.assertIn("users", payload) + self.assertTrue(len(payload["users"]) >= 1) + self.assertIn("username", payload["users"][0]) diff --git a/tests/test_download_config.py b/tests/test_download_config.py new file mode 100644 index 000000000..66150ed2f --- /dev/null +++ b/tests/test_download_config.py @@ -0,0 +1,433 @@ +import datetime +import os +import tempfile +import tracemalloc +from base64 import b32encode, b64encode +from unittest import TestCase, mock +from unittest.mock import MagicMock + +from icloudpd import download +from icloudpd.limiter import AdaptiveDownloadLimiter +from pyicloud_ipd.asset_version import AssetVersion +from pyicloud_ipd.version_size import AssetVersionSize + + +class _FakeResponse: + def __init__(self, *, ok: bool, status_code: int, body_chunks: list[bytes] | None = None) -> None: + self.ok = ok + self.status_code = status_code + self.headers: dict[str, str] = {} + self._body_chunks = body_chunks or [] + + def iter_content(self, chunk_size: int) -> list[bytes]: # noqa: ARG002 + return self._body_chunks + + +class _FakePhoto: + def __init__(self, responses: list[_FakeResponse]) -> None: + self.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + self._responses = list(responses) + self.download_calls: list[tuple[str, int]] = [] + + def download(self, _session: object, url: str, current_size: int) -> _FakeResponse: + self.download_calls.append((url, current_size)) + return self._responses.pop(0) + + +class _LargeStreamingResponse: + def __init__(self, total_bytes: int) -> None: + self._total_bytes = total_bytes + + def iter_content(self, chunk_size: int): + sent = 0 + while sent < self._total_bytes: + take = min(chunk_size, self._total_bytes - sent) + sent += take + yield b"x" * take + + +class DownloadConfigTestCase(TestCase): + def setUp(self) -> None: + self._tmpdir = tempfile.mkdtemp(prefix="icloudpd-download-config-") + download.set_download_chunk_bytes(262144) + download.set_download_verification(verify_size=True, verify_checksum=False) + download.set_download_limiter(None) + + def tearDown(self) -> None: + for entry in os.listdir(self._tmpdir): + os.remove(os.path.join(self._tmpdir, entry)) + os.rmdir(self._tmpdir) + + def test_download_response_uses_configured_chunk_size(self) -> None: + download.set_download_chunk_bytes(65536) + + response = MagicMock() + response.iter_content.return_value = [b"abc", b"def"] + + temp_path = "/tmp/icloudpd-download-chunk-test.part" + final_path = "/tmp/icloudpd-download-chunk-test.bin" + for path in [temp_path, final_path]: + if os.path.exists(path): + os.remove(path) + + ok = download.download_response_to_path( + response, + temp_path, + append_mode=False, + download_path=final_path, + created_date=datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc), + ) + self.assertTrue(ok) + response.iter_content.assert_called_once_with(chunk_size=65536) + self.assertTrue(os.path.exists(final_path)) + + def test_download_response_streaming_memory_is_bounded(self) -> None: + download.set_download_chunk_bytes(65536) + response = _LargeStreamingResponse(total_bytes=64 * 1024 * 1024) + + temp_path = os.path.join(self._tmpdir, "bounded-stream.part") + final_path = os.path.join(self._tmpdir, "bounded-stream.bin") + + tracemalloc.start() + ok = download.download_response_to_path( + response, + temp_path, + append_mode=False, + download_path=final_path, + created_date=datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc), + ) + _current, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + self.assertTrue(ok) + self.assertEqual(os.path.getsize(final_path), 64 * 1024 * 1024) + self.assertLess(peak, 8 * 1024 * 1024) + + def test_download_media_verifies_size(self) -> None: + download.set_download_verification(verify_size=True, verify_checksum=False) + + photo = MagicMock() + photo.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + response = MagicMock() + response.ok = True + response.status_code = 200 + response.iter_content.return_value = [b"abc"] + photo.download.return_value = response + + icloud = MagicMock() + download_path = os.path.join(self._tmpdir, "test-size.jpg") + checksum = b64encode(b"0123456789abcdef").decode() + + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=download_path, + version=AssetVersion(size=3, url="https://example.test/file", type="public.jpeg", checksum=checksum), + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "test-size.jpg", + ) + self.assertTrue(ok) + self.assertTrue(os.path.exists(download_path)) + + bad_path = os.path.join(self._tmpdir, "test-size-mismatch.jpg") + bad_ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=bad_path, + version=AssetVersion(size=4, url="https://example.test/file", type="public.jpeg", checksum=checksum), + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "test-size-mismatch.jpg", + ) + self.assertFalse(bad_ok) + self.assertFalse(os.path.exists(bad_path)) + + def test_download_media_verifies_checksum(self) -> None: + download.set_download_verification(verify_size=False, verify_checksum=True) + + photo = MagicMock() + photo.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + response = MagicMock() + response.ok = True + response.status_code = 200 + response.iter_content.return_value = [b"abc"] + photo.download.return_value = response + + icloud = MagicMock() + download_path = os.path.join(self._tmpdir, "test-checksum.jpg") + matching_checksum = b64encode(b"\x90\x01P\x98<\xd2O\xb0\xd6\x96?}(\xe1\x7fr").decode() + + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=download_path, + version=AssetVersion( + size=3, + url="https://example.test/file", + type="public.jpeg", + checksum=matching_checksum, + ), + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "test-checksum.jpg", + ) + self.assertTrue(ok) + self.assertTrue(os.path.exists(download_path)) + + bad_path = os.path.join(self._tmpdir, "test-checksum-mismatch.jpg") + bad_checksum = b64encode(b"0000000000000000").decode() + bad_ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=bad_path, + version=AssetVersion( + size=3, + url="https://example.test/file", + type="public.jpeg", + checksum=bad_checksum, + ), + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "test-checksum-mismatch.jpg", + ) + self.assertFalse(bad_ok) + self.assertFalse(os.path.exists(bad_path)) + + def test_download_media_restarts_partial_when_range_ignored(self) -> None: + download.set_download_verification(verify_size=True, verify_checksum=False) + + photo = MagicMock() + photo.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + response = MagicMock() + response.ok = True + response.status_code = 200 + response.iter_content.return_value = [b"new"] + photo.download.return_value = response + + icloud = MagicMock() + checksum = b64encode(b"0123456789abcdef").decode() + version = AssetVersion(size=3, url="https://example.test/file", type="public.jpeg", checksum=checksum) + checksum_raw = b"0123456789abcdef" + temp_download_path = os.path.join( + self._tmpdir, b32encode(checksum_raw).decode() + ".part" + ) + with open(temp_download_path, "wb") as part_file: + part_file.write(b"old-bytes") + + download_path = os.path.join(self._tmpdir, "range-restart.jpg") + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=download_path, + version=version, + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "range-restart.jpg", + ) + self.assertTrue(ok) + photo.download.assert_called_once_with(icloud.photos.session, version.url, 9) + with open(download_path, "rb") as downloaded: + self.assertEqual(downloaded.read(), b"new") + + def test_download_media_resumes_when_server_returns_partial_content(self) -> None: + download.set_download_verification(verify_size=True, verify_checksum=False) + + checksum_raw = b"0123456789abcdef" + checksum = b64encode(checksum_raw).decode() + version = AssetVersion(size=6, url="https://example.test/file", type="public.jpeg", checksum=checksum) + temp_download_path = os.path.join( + self._tmpdir, b32encode(checksum_raw).decode() + ".part" + ) + with open(temp_download_path, "wb") as part_file: + part_file.write(b"old") + + photo = _FakePhoto( + responses=[_FakeResponse(ok=True, status_code=206, body_chunks=[b"new"])] + ) + icloud = MagicMock() + download_path = os.path.join(self._tmpdir, "range-resume.jpg") + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=download_path, + version=version, + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "range-resume.jpg", + ) + self.assertTrue(ok) + self.assertEqual(photo.download_calls, [(version.url, 3)]) + with open(download_path, "rb") as downloaded: + self.assertEqual(downloaded.read(), b"oldnew") + + def test_download_media_restarts_partial_when_server_returns_416(self) -> None: + download.set_download_verification(verify_size=True, verify_checksum=False) + + checksum_raw = b"0123456789abcdef" + checksum = b64encode(checksum_raw).decode() + version = AssetVersion(size=3, url="https://example.test/file", type="public.jpeg", checksum=checksum) + temp_download_path = os.path.join( + self._tmpdir, b32encode(checksum_raw).decode() + ".part" + ) + with open(temp_download_path, "wb") as part_file: + part_file.write(b"stale-partial") + + photo = _FakePhoto( + responses=[ + _FakeResponse(ok=False, status_code=416), + _FakeResponse(ok=True, status_code=200, body_chunks=[b"new"]), + ] + ) + icloud = MagicMock() + download_path = os.path.join(self._tmpdir, "range-416-restart.jpg") + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=download_path, + version=version, + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "range-416-restart.jpg", + ) + self.assertTrue(ok) + self.assertEqual(photo.download_calls, [(version.url, 13), (version.url, 0)]) + with open(download_path, "rb") as downloaded: + self.assertEqual(downloaded.read(), b"new") + + def test_download_media_throttling_updates_limiter(self) -> None: + limiter = AdaptiveDownloadLimiter(max_workers=4, cooldown_seconds=0.0) + download.set_download_limiter(limiter) + download.set_retry_config( + retry_config=download.RetryConfig( + max_retries=0, + backoff_base_seconds=1.0, + backoff_max_seconds=1.0, + respect_retry_after=True, + throttle_cooldown_seconds=0.0, + ) + ) + + photo = MagicMock() + photo.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + response = MagicMock() + response.ok = False + response.status_code = 429 + response.headers = {} + photo.download.return_value = response + + icloud = MagicMock() + checksum = b64encode(b"0123456789abcdef").decode() + + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=os.path.join(self._tmpdir, "throttle.jpg"), + version=AssetVersion(size=3, url="https://example.test/file", type="public.jpeg", checksum=checksum), + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "throttle.jpg", + ) + self.assertFalse(ok) + self.assertEqual(limiter.current_limit, 1) + + def test_download_media_low_disk_space_classification(self) -> None: + download.set_download_verification(verify_size=False, verify_checksum=False) + photo = MagicMock() + photo.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + icloud = MagicMock() + checksum = b64encode(b"0123456789abcdef").decode() + with mock.patch("icloudpd.download.has_disk_space_for_download", return_value=False): + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=os.path.join(self._tmpdir, "low-disk.jpg"), + version=AssetVersion( + size=10_000_000, + url="https://example.test/file", + type="public.jpeg", + checksum=checksum, + ), + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "low-disk.jpg", + ) + self.assertFalse(ok) + photo.download.assert_not_called() + + def test_download_media_refreshes_expired_url_once(self) -> None: + download.set_download_verification(verify_size=False, verify_checksum=False) + photo = MagicMock() + photo.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + first = MagicMock() + first.ok = False + first.status_code = 403 + first.headers = {} + second = MagicMock() + second.ok = True + second.status_code = 200 + second.iter_content.return_value = [b"abc"] + photo.download.side_effect = [first, second] + + icloud = MagicMock() + checksum = b64encode(b"0123456789abcdef").decode() + original = AssetVersion(size=3, url="https://example.test/stale", type="public.jpeg", checksum=checksum) + refreshed = AssetVersion( + size=3, url="https://example.test/fresh", type="public.jpeg", checksum=checksum + ) + refresh_cb = MagicMock(return_value=refreshed) + + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=os.path.join(self._tmpdir, "refresh-url.jpg"), + version=original, + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "refresh-url.jpg", + refresh_version=refresh_cb, + ) + self.assertTrue(ok) + self.assertEqual(photo.download.call_args_list[0].args[1], "https://example.test/stale") + self.assertEqual(photo.download.call_args_list[1].args[1], "https://example.test/fresh") + refresh_cb.assert_called_once() + self.assertFalse(download.consume_url_refresh_needed_signal()) + + def test_download_media_sets_url_refresh_signal_when_unrecoverable(self) -> None: + download.set_download_verification(verify_size=False, verify_checksum=False) + photo = MagicMock() + photo.created = datetime.datetime(2026, 3, 2, tzinfo=datetime.timezone.utc) + response = MagicMock() + response.ok = False + response.status_code = 410 + response.headers = {} + photo.download.return_value = response + + icloud = MagicMock() + checksum = b64encode(b"0123456789abcdef").decode() + version = AssetVersion(size=3, url="https://example.test/stale", type="public.jpeg", checksum=checksum) + refresh_cb = MagicMock(return_value=None) + + ok = download.download_media( + logger=MagicMock(), + dry_run=False, + icloud=icloud, + photo=photo, + download_path=os.path.join(self._tmpdir, "refresh-url-fail.jpg"), + version=version, + size=AssetVersionSize.ORIGINAL, + filename_builder=lambda _p: "refresh-url-fail.jpg", + refresh_version=refresh_cb, + ) + self.assertFalse(ok) + self.assertTrue(download.consume_url_refresh_needed_signal()) diff --git a/tests/test_download_photos.py b/tests/test_download_photos.py index d06974798..2c7e825da 100644 --- a/tests/test_download_photos.py +++ b/tests/test_download_photos.py @@ -21,6 +21,7 @@ from pyicloud_ipd.services.photos import PhotoAlbum, PhotoAsset, PhotoLibrary from pyicloud_ipd.version_size import AssetVersionSize, LivePhotoVersionSize from tests.helpers import ( + DEFAULT_ENV, calc_data_dir, create_files, path_from_project_root, @@ -28,6 +29,7 @@ recreate_path, run_cassette, run_icloudpd_test, + run_main_env, ) @@ -356,6 +358,7 @@ def test_until_found(self) -> None: if (f[2] == "photo" and f[1].endswith(".MOV")) else AssetVersionSize.ORIGINAL, ANY, # filename_builder + refresh_version=ANY, ), files_to_download_ext, ) @@ -892,6 +895,7 @@ def test_download_two_sizes_with_force_size(self) -> None: ANY, AssetVersionSize.THUMB, ANY, # filename_builder + refresh_version=ANY, ) assert result.exit_code == 0 @@ -1738,6 +1742,118 @@ def mock_raise_response_error() -> NoReturn: self.assertEqual(result.exit_code, 1, "Exit Code") + def test_retry_metadata_throttling_then_success(self) -> None: + base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) + data_dir = calc_data_dir(base_dir) + cookie_dir = os.path.join(base_dir, "cookie") + for directory in [base_dir, data_dir, cookie_dir]: + recreate_path(directory) + + class FakeAlbum: + def __init__(self) -> None: + self.calls = 0 + + def __iter__(self) -> Any: + if self.calls == 0: + self.calls += 1 + raise PyiCloudAPIResponseException("Rate limited", "429") + return iter([]) + + def __len__(self) -> int: + return 0 + + def increment_offset(self, _count: int) -> None: + return None + + class FakePhotos: + def __init__(self, album: FakeAlbum) -> None: + self.all = album + self.albums = {} + self.private_libraries = {"PrimarySync": self} + self.shared_libraries = {} + + class FakeCloud: + def __init__(self, photos: FakePhotos) -> None: + self.photos = photos + self.response_observer = None + + fake_album = FakeAlbum() + fake_cloud = FakeCloud(FakePhotos(fake_album)) + + with mock.patch("icloudpd.base.authenticator", return_value=fake_cloud): + result = print_result_exception( + run_main_env( + DEFAULT_ENV, + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "--directory", + data_dir, + "--cookie-directory", + cookie_dir, + "--skip-videos", + "--skip-live-photos", + "--no-progress-bar", + "--max-retries", + "1", + "--backoff-base-seconds", + "0.01", + "--backoff-max-seconds", + "0.01", + ], + ) + ) + + self.assertIn("Transient error (PyiCloudAPIResponseException). Retrying in", result.output) + self.assertIn("All photos have been downloaded", result.output) + self.assertEqual(result.exit_code, 0, "Exit Code") + + def test_retry_download_503_then_success(self) -> None: + base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) + original_download = PhotoAsset.download + calls = {"count": 0} + + def flaky_download(photo: PhotoAsset, session: Any, url: str, start: int = 0) -> Any: + if calls["count"] == 0: + calls["count"] += 1 + raise PyiCloudAPIResponseException("Service unavailable", "503") + return original_download(photo, session, url, start) + + with mock.patch("time.sleep") as sleep_mock: # noqa: SIM117 + with mock.patch.object(PhotoAsset, "download", new=flaky_download): + _, result = run_icloudpd_test( + self.assertEqual, + self.root_path, + base_dir, + "listing_photos.yml", + [], + [("2018/07/31", "IMG_7409.JPG")], + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "--recent", + "1", + "--skip-videos", + "--skip-live-photos", + "--no-progress-bar", + "--max-retries", + "1", + "--backoff-base-seconds", + "0.01", + "--backoff-max-seconds", + "0.01", + ], + ) + + self.assertIn("Error downloading IMG_7409.JPG, retrying after", result.output) + self.assertIn("Downloaded", result.output) + self.assertEqual(sleep_mock.call_count, 1, "sleep count") + self.assertEqual(result.exit_code, 0, "Exit Code") + def test_handle_io_error_mkdir(self) -> None: base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) original_makedirs = os.makedirs @@ -2514,4 +2630,65 @@ def test_resume_download(self) -> None: # check size photo_size = os.path.getsize(out_path) - self.assertEqual(617 + 1234, photo_size, "photo size") + self.assertEqual(617, photo_size, "photo size") + + def test_no_remote_count_skips_album_length_call(self) -> None: + base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) + files_to_download = [("2018/07/31", "IMG_7409.JPG")] + + with mock.patch.object(PhotoAlbum, "__len__", side_effect=AssertionError("len should not be called")): + _data_dir, result = run_icloudpd_test( + self.assertEqual, + self.root_path, + base_dir, + "listing_photos.yml", + [], + files_to_download, + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "--recent", + "1", + "--skip-videos", + "--skip-live-photos", + "--no-remote-count", + "--no-progress-bar", + "--threads-num", + "1", + ], + ) + self.assertEqual(result.exit_code, 0) + + def test_until_found_skips_album_length_call(self) -> None: + base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) + files_to_download = [] + + with mock.patch.object( + PhotoAlbum, "__len__", side_effect=AssertionError("len should not be called") + ): + _data_dir, result = run_icloudpd_test( + self.assertEqual, + self.root_path, + base_dir, + "listing_photos.yml", + [], + files_to_download, + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "--until-found", + "1", + "--recent", + "0", + "--skip-videos", + "--skip-live-photos", + "--no-progress-bar", + "--threads-num", + "1", + ], + ) + self.assertEqual(result.exit_code, 0) diff --git a/tests/test_download_photos_id.py b/tests/test_download_photos_id.py index 1a10fd86e..4d7cd5090 100644 --- a/tests/test_download_photos_id.py +++ b/tests/test_download_photos_id.py @@ -352,6 +352,7 @@ def test_until_found_name_id7(self) -> None: if (f[2] == "photo" and f[1].endswith(".MOV")) else AssetVersionSize.ORIGINAL, ANY, # filename_builder + refresh_version=ANY, ), files_to_download_ext, ) diff --git a/tests/test_email_notifications.py b/tests/test_email_notifications.py index 58d945153..857b82247 100644 --- a/tests/test_email_notifications.py +++ b/tests/test_email_notifications.py @@ -137,6 +137,36 @@ def test_2sa_required_notification_script(self) -> None: self.assertEqual(result.exit_code, 1, "exit code") subprocess_patched.assert_called_once_with(["test_script.sh"]) + @freeze_time("2018-01-01") + def test_2sa_required_notification_script_with_state_db(self) -> None: + base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3], "state_db") + cookie_dir = os.path.join(base_dir, "cookie") + data_dir = os.path.join(base_dir, "data") + + for dir in [base_dir, cookie_dir, data_dir]: + recreate_path(dir) + + with patch("subprocess.call") as subprocess_patched: + result = run_cassette( + os.path.join(self.vcr_path, "auth_requires_2fa.yml"), + [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "--notification-script", + "./test_script.sh", + "--state-db", + "-d", + data_dir, + "--cookie-directory", + cookie_dir, + ], + ) + self.assertEqual(result.exit_code, 1, "exit code") + subprocess_patched.assert_called_once_with(["test_script.sh"]) + self.assertTrue(os.path.exists(os.path.join(cookie_dir, "icloudpd.sqlite"))) + @freeze_time("2018-01-01") def test_2sa_required_email_notification_from(self) -> None: base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) diff --git a/tests/test_filters.py b/tests/test_filters.py new file mode 100644 index 000000000..9956ad5ce --- /dev/null +++ b/tests/test_filters.py @@ -0,0 +1,70 @@ +import datetime +from unittest import TestCase +from unittest.mock import MagicMock + +from icloudpd.base import where_builder +from pyicloud_ipd.item_type import AssetItemType + + +class FilterTestCase(TestCase): + def test_skip_added_before_filters_older_added_date(self) -> None: + photo = MagicMock() + photo.item_type = AssetItemType.IMAGE + photo.created = datetime.datetime(2026, 1, 10, tzinfo=datetime.timezone.utc) + photo.added_date = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + + result = where_builder( + logger=MagicMock(), + skip_videos=False, + skip_created_before=None, + skip_created_after=None, + skip_added_before=datetime.datetime(2026, 1, 5, tzinfo=datetime.timezone.utc), + skip_added_after=None, + skip_photos=False, + filename_builder=lambda _photo: "file.jpg", + photo=photo, + ) + self.assertFalse(result) + + def test_skip_added_after_filters_newer_added_date(self) -> None: + photo = MagicMock() + photo.item_type = AssetItemType.IMAGE + photo.created = datetime.datetime(2026, 1, 10, tzinfo=datetime.timezone.utc) + photo.added_date = datetime.datetime(2026, 1, 10, tzinfo=datetime.timezone.utc) + + result = where_builder( + logger=MagicMock(), + skip_videos=False, + skip_created_before=None, + skip_created_after=None, + skip_added_before=None, + skip_added_after=datetime.datetime(2026, 1, 5, tzinfo=datetime.timezone.utc), + skip_photos=False, + filename_builder=lambda _photo: "file.jpg", + photo=photo, + ) + self.assertFalse(result) + + def test_skip_added_filters_are_ignored_when_added_date_missing(self) -> None: + class _PhotoWithoutAddedDate: + item_type = AssetItemType.IMAGE + created = datetime.datetime(2026, 1, 10, tzinfo=datetime.timezone.utc) + + @property + def added_date(self) -> datetime.datetime: + raise KeyError("addedDate") + + photo = _PhotoWithoutAddedDate() + + result = where_builder( + logger=MagicMock(), + skip_videos=False, + skip_created_before=None, + skip_created_after=None, + skip_added_before=datetime.datetime(2026, 1, 5, tzinfo=datetime.timezone.utc), + skip_added_after=datetime.datetime(2026, 1, 15, tzinfo=datetime.timezone.utc), + skip_photos=False, + filename_builder=lambda _photo: "file.jpg", + photo=photo, + ) + self.assertTrue(result) diff --git a/tests/test_folder_structure.py b/tests/test_folder_structure.py index 874d76eb9..39de3e816 100644 --- a/tests/test_folder_structure.py +++ b/tests/test_folder_structure.py @@ -1,4 +1,5 @@ import inspect +import locale import os import sys from typing import List, Tuple @@ -123,6 +124,14 @@ def test_folder_structure_none(self) -> None: @pytest.mark.skipif(sys.platform == "win32", reason="local strings are not working on windows") def test_folder_structure_de_posix(self) -> None: + current_locale = locale.setlocale(locale.LC_ALL) + try: + locale.setlocale(locale.LC_ALL, "de_DE.UTF-8") + except locale.Error: + pytest.skip("de_DE.UTF-8 locale is not available in this environment") + finally: + locale.setlocale(locale.LC_ALL, current_locale) + base_dir = os.path.join(self.fixtures_path, inspect.stack()[0][3]) data_dir, result = run_icloudpd_test( diff --git a/tests/test_limiter.py b/tests/test_limiter.py new file mode 100644 index 000000000..c89bf0367 --- /dev/null +++ b/tests/test_limiter.py @@ -0,0 +1,96 @@ +import threading +import time +from unittest import TestCase + +from icloudpd.limiter import AdaptiveDownloadLimiter + + +class AdaptiveDownloadLimiterTestCase(TestCase): + def test_additive_increase(self) -> None: + limiter = AdaptiveDownloadLimiter( + max_workers=4, + min_workers=1, + cooldown_seconds=0.0, + increase_every=2, + ) + + limiter.on_throttle() # 4 -> 2 + self.assertEqual(limiter.current_limit, 2) + limiter.on_success() + self.assertEqual(limiter.current_limit, 2) + limiter.on_success() + self.assertEqual(limiter.current_limit, 3) + + def test_multiplicative_decrease_and_cooldown(self) -> None: + limiter = AdaptiveDownloadLimiter( + max_workers=8, + min_workers=1, + cooldown_seconds=0.2, + increase_every=1, + ) + limiter.on_throttle() # 8 -> 4 + self.assertEqual(limiter.current_limit, 4) + self.assertGreater(limiter.cooldown_remaining_seconds, 0.0) + + def test_slot_respects_limit(self) -> None: + limiter = AdaptiveDownloadLimiter( + max_workers=2, + min_workers=1, + cooldown_seconds=0.0, + increase_every=10, + ) + entered = threading.Event() + + def hold_slot() -> None: + with limiter.slot(): + entered.set() + time.sleep(0.15) + + with limiter.slot(): + t = threading.Thread(target=hold_slot) + t.start() + self.assertTrue(entered.wait(0.5)) + self.assertFalse(limiter.acquire(timeout=0.05)) + t.join() + + self.assertTrue(limiter.acquire(timeout=0.2)) + limiter.release() + + def test_stop_drain_and_restart(self) -> None: + limiter = AdaptiveDownloadLimiter( + max_workers=2, + min_workers=1, + cooldown_seconds=0.0, + increase_every=10, + ) + entered = threading.Event() + release_holder = threading.Event() + + def holder() -> None: + with limiter.slot(): + entered.set() + release_holder.wait(1.0) + + thread = threading.Thread(target=holder) + thread.start() + self.assertTrue(entered.wait(0.5)) + + stopped = [] + + def stop_limiter() -> None: + stopped.append(limiter.stop(wait=True, timeout=1.0)) + + stopper = threading.Thread(target=stop_limiter) + stopper.start() + time.sleep(0.05) + self.assertFalse(limiter.acquire(timeout=0.05)) + + release_holder.set() + stopper.join() + thread.join() + self.assertEqual(stopped, [True]) + + self.assertFalse(limiter.acquire(timeout=0.05)) + limiter.start() + self.assertTrue(limiter.acquire(timeout=0.2)) + limiter.release() diff --git a/tests/test_logging.py b/tests/test_logging.py new file mode 100644 index 000000000..ae7a13326 --- /dev/null +++ b/tests/test_logging.py @@ -0,0 +1,112 @@ +import io +import json +from unittest import TestCase +from unittest.mock import patch + +from icloudpd.base import create_logger +from icloudpd.config import GlobalConfig +from icloudpd.log_level import LogLevel +from icloudpd.mfa_provider import MFAProvider +from icloudpd.password_provider import PasswordProvider + + +class LoggingTestCase(TestCase): + def test_json_log_mode_emits_structured_fields(self) -> None: + config = GlobalConfig( + help=False, + version=False, + use_os_locale=False, + only_print_filenames=False, + log_level=LogLevel.INFO, + log_format="json", + no_progress_bar=True, + threads_num=1, + domain="com", + watch_with_interval=None, + password_providers=[PasswordProvider.PARAMETER], + mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, + ) + stream = io.StringIO() + with patch("sys.stdout", stream): + logger = create_logger(config) + logger.info("hello structured logs") + + payload = json.loads(stream.getvalue().strip()) + self.assertEqual(payload["level"], "INFO") + self.assertEqual(payload["message"], "hello structured logs") + self.assertEqual(payload["logger"], "icloudpd") + self.assertIsNotNone(payload["run_id"]) + self.assertIn("asset_id", payload) + self.assertIn("attempt", payload) + self.assertIn("http_status", payload) + + def test_json_log_mode_redacts_sensitive_values(self) -> None: + config = GlobalConfig( + help=False, + version=False, + use_os_locale=False, + only_print_filenames=False, + log_level=LogLevel.INFO, + log_format="json", + no_progress_bar=True, + threads_num=1, + domain="com", + watch_with_interval=None, + password_providers=[PasswordProvider.PARAMETER], + mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, + ) + stream = io.StringIO() + with patch("sys.stdout", stream): + logger = create_logger(config) + logger.info( + 'payload={"password":"secret"} token=abc123 Authorization=Bearer mytoken cookie=rawcookie' + ) + + payload = json.loads(stream.getvalue().strip()) + message = payload["message"] + self.assertNotIn("secret", message) + self.assertNotIn("abc123", message) + self.assertNotIn("mytoken", message) + self.assertNotIn("rawcookie", message) + self.assertIn("REDACTED", message) + + def test_text_log_mode_redacts_sensitive_values(self) -> None: + config = GlobalConfig( + help=False, + version=False, + use_os_locale=False, + only_print_filenames=False, + log_level=LogLevel.INFO, + log_format="text", + no_progress_bar=True, + threads_num=1, + domain="com", + watch_with_interval=None, + password_providers=[PasswordProvider.PARAMETER], + mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=5.0, + backoff_max_seconds=300.0, + respect_retry_after=True, + throttle_cooldown_seconds=60.0, + ) + stream = io.StringIO() + with patch("sys.stdout", stream): + logger = create_logger(config) + logger.info("password=topsecret token=qwerty scnt=abcdef") + + output = stream.getvalue() + self.assertNotIn("topsecret", output) + self.assertNotIn("qwerty", output) + self.assertNotIn("abcdef", output) + self.assertIn("password=REDACTED", output) diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 000000000..9f6729f08 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,39 @@ +import json +import os +import shutil +import tempfile +from unittest import TestCase + +from icloudpd.metrics import RunMetrics, write_metrics_json + + +class MetricsTestCase(TestCase): + def test_run_metrics_snapshot_includes_expected_fields(self) -> None: + metrics = RunMetrics(username="u1") + metrics.start() + metrics.on_asset_considered() + metrics.on_download_attempt() + metrics.on_download_success(123) + metrics.on_retry() + metrics.on_throttle() + metrics.set_queue_depth(7) + metrics.finish() + + snapshot = metrics.snapshot() + self.assertEqual(snapshot["username"], "u1") + self.assertEqual(snapshot["run_mode"], "legacy_stateless") + self.assertEqual(snapshot["downloads_attempted"], 1) + self.assertEqual(snapshot["downloads_succeeded"], 1) + self.assertEqual(snapshot["bytes_downloaded"], 123) + self.assertIn("throughput_downloads_per_sec", snapshot) + self.assertIn("success_gap", snapshot) + + def test_write_metrics_json(self) -> None: + tmpdir = tempfile.mkdtemp(prefix="icloudpd-metrics-") + self.addCleanup(lambda: shutil.rmtree(tmpdir, ignore_errors=True)) + out_path = os.path.join(tmpdir, "metrics.json") + payload = {"users": [{"username": "u1"}]} + write_metrics_json(out_path, payload) + with open(out_path, encoding="utf-8") as f: + parsed = json.load(f) + self.assertEqual(parsed, payload) diff --git a/tests/test_mode_contract.py b/tests/test_mode_contract.py new file mode 100644 index 000000000..f2fe3d9a3 --- /dev/null +++ b/tests/test_mode_contract.py @@ -0,0 +1,84 @@ +import glob +import inspect +import os +import sqlite3 +from unittest import TestCase + +import pytest + +from tests.helpers import path_from_project_root, run_icloudpd_test + + +class ModeContractTestCase(TestCase): + @pytest.fixture(autouse=True) + def inject_fixtures(self) -> None: + self.root_path = path_from_project_root(__file__) + self.fixtures_path = os.path.join(self.root_path, "fixtures") + + def _run_copy_mode(self, *, use_state_db: bool, base_dir_suffix: str) -> tuple[str, str]: + base_dir = os.path.join( + self.fixtures_path, + inspect.stack()[0][3], + base_dir_suffix, + ) + files_to_create = [ + ("2018/07/30", "IMG_7408.JPG", 1151066), + ("2018/07/30", "IMG_7407.JPG", 656257), + ] + files_to_download = [("2018/07/31", "IMG_7409.JPG")] + params = [ + "--username", + "jdoe@gmail.com", + "--password", + "password1", + "--recent", + "5", + "--skip-videos", + "--skip-live-photos", + "--set-exif-datetime", + "--no-progress-bar", + "--threads-num", + "1", + ] + if use_state_db: + params.append("--state-db") + data_dir, result = run_icloudpd_test( + self.assertEqual, + self.root_path, + base_dir, + "listing_photos.yml", + files_to_create, + files_to_download, + params, + ) + self.assertEqual(result.exit_code, 0) + cookie_dir = os.path.join(base_dir, "cookie") + return data_dir, cookie_dir + + def test_legacy_stateless_mode_contract(self) -> None: + _data_dir, cookie_dir = self._run_copy_mode(use_state_db=False, base_dir_suffix="legacy") + self.assertFalse(os.path.exists(os.path.join(cookie_dir, "icloudpd.sqlite"))) + + def test_stateful_engine_mode_contract(self) -> None: + _data_dir, cookie_dir = self._run_copy_mode(use_state_db=True, base_dir_suffix="stateful") + db_path = os.path.join(cookie_dir, "icloudpd.sqlite") + self.assertTrue(os.path.exists(db_path)) + with sqlite3.connect(db_path) as conn: + task_count = conn.execute("SELECT COUNT(*) FROM tasks").fetchone()[0] + checkpoint_count = conn.execute("SELECT COUNT(*) FROM checkpoints").fetchone()[0] + self.assertGreater(task_count, 0) + self.assertGreater(checkpoint_count, 0) + + def test_mode_parity_for_downloaded_files(self) -> None: + legacy_data_dir, _legacy_cookie_dir = self._run_copy_mode( + use_state_db=False, base_dir_suffix="parity-legacy" + ) + stateful_data_dir, _stateful_cookie_dir = self._run_copy_mode( + use_state_db=True, base_dir_suffix="parity-stateful" + ) + + def relative_tree(root: str) -> list[str]: + files = glob.glob(os.path.join(root, "**/*.*"), recursive=True) + return sorted(os.path.relpath(path, root) for path in files) + + self.assertEqual(relative_tree(legacy_data_dir), relative_tree(stateful_data_dir)) diff --git a/tests/test_retry_utils.py b/tests/test_retry_utils.py new file mode 100644 index 000000000..225bd878f --- /dev/null +++ b/tests/test_retry_utils.py @@ -0,0 +1,88 @@ +import datetime +from unittest import TestCase + +from requests.exceptions import ChunkedEncodingError + +from icloudpd.retry_utils import ( + RetryConfig, + is_fatal_auth_config_error, + is_session_invalid_error, + is_throttle_error, + is_transient_error, + parse_retry_after_seconds, +) +from pyicloud_ipd.exceptions import ( + PyiCloud2SARequiredException, + PyiCloudAPIResponseException, + PyiCloudConnectionErrorException, + PyiCloudFailedLoginException, + PyiCloudFailedMFAException, + PyiCloudNoStoredPasswordAvailableException, + PyiCloudServiceNotActivatedException, + PyiCloudServiceUnavailableException, +) + + +class RetryUtilsTestCase(TestCase): + def test_parse_retry_after_numeric(self) -> None: + self.assertEqual(parse_retry_after_seconds("120"), 120.0) + + def test_parse_retry_after_http_date(self) -> None: + future = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=42) + parsed = parse_retry_after_seconds(future.strftime("%a, %d %b %Y %H:%M:%S GMT")) + assert parsed is not None + self.assertTrue(0 < parsed <= 42) + + def test_transient_error_classification(self) -> None: + self.assertTrue(is_transient_error(PyiCloudConnectionErrorException("test"))) + self.assertTrue(is_transient_error(PyiCloudServiceUnavailableException("test"))) + self.assertTrue(is_transient_error(PyiCloudAPIResponseException("throttle", "429"))) + self.assertTrue(is_transient_error(ChunkedEncodingError("chunk error"))) + self.assertFalse(is_transient_error(PyiCloudAPIResponseException("bad auth", "401"))) + + def test_session_invalid_classification(self) -> None: + self.assertTrue( + is_session_invalid_error(PyiCloudAPIResponseException("Invalid global session", "500")) + ) + self.assertFalse(is_session_invalid_error(PyiCloudAPIResponseException("other", "500"))) + + def test_fatal_auth_config_classification(self) -> None: + self.assertTrue(is_fatal_auth_config_error(PyiCloudFailedLoginException("bad credentials"))) + self.assertTrue(is_fatal_auth_config_error(PyiCloudFailedMFAException("mfa unavailable"))) + self.assertTrue(is_fatal_auth_config_error(PyiCloud2SARequiredException("user@example.com"))) + self.assertTrue( + is_fatal_auth_config_error(PyiCloudNoStoredPasswordAvailableException("no password")) + ) + self.assertTrue( + is_fatal_auth_config_error(PyiCloudServiceNotActivatedException("web disabled", "X")) + ) + self.assertFalse(is_fatal_auth_config_error(PyiCloudConnectionErrorException("test"))) + + def test_throttle_classification(self) -> None: + self.assertTrue(is_throttle_error(PyiCloudAPIResponseException("ACCESS_DENIED", "ACCESS_DENIED"))) + self.assertTrue(is_throttle_error(PyiCloudAPIResponseException("request throttled", "500"))) + self.assertFalse(is_throttle_error(PyiCloudAPIResponseException("other", "500"))) + + def test_next_delay_respects_retry_after(self) -> None: + config = RetryConfig( + max_retries=2, + backoff_base_seconds=1, + backoff_max_seconds=300, + respect_retry_after=True, + throttle_cooldown_seconds=60, + jitter_fraction=0, + ) + delay = config.next_delay_seconds(1, retry_after="120", throttle_error=False) + self.assertEqual(delay, 120) + + def test_next_delay_applies_throttle_cooldown(self) -> None: + config = RetryConfig( + max_retries=2, + backoff_base_seconds=1, + backoff_max_seconds=300, + respect_retry_after=True, + throttle_cooldown_seconds=60, + jitter_fraction=0, + ) + delay = config.next_delay_seconds(1, throttle_error=True) + self.assertEqual(delay, 60) diff --git a/tests/test_session_persistence.py b/tests/test_session_persistence.py new file mode 100644 index 000000000..bd907b6df --- /dev/null +++ b/tests/test_session_persistence.py @@ -0,0 +1,52 @@ +import json +import os +import shutil +import tempfile +import threading +import time +from unittest import TestCase + +from pyicloud_ipd.session import persist_session_and_cookies + + +class _FakeCookieJar: + def save(self, filename=None, ignore_discard=True, ignore_expires=True): # type: ignore[no-untyped-def] + assert filename is not None + # Delay to increase overlap pressure in concurrent writers. + time.sleep(0.01) + with open(filename, "w", encoding="utf-8") as f: + f.write("cookiejar-ok\n") + + +class SessionPersistenceTestCase(TestCase): + def test_concurrent_persistence_does_not_corrupt_session_or_cookies(self) -> None: + tmpdir = tempfile.mkdtemp(prefix="icloudpd-session-persist-") + self.addCleanup(lambda: shutil.rmtree(tmpdir, ignore_errors=True)) + session_path = os.path.join(tmpdir, "session.json") + cookiejar_path = os.path.join(tmpdir, "cookies") + + def writer(index: int) -> None: + payload = {"client_id": f"client-{index}", "session_id": f"session-{index}"} + persist_session_and_cookies( + session_path=session_path, + cookiejar_path=cookiejar_path, + session_data=payload, + cookies=_FakeCookieJar(), + ) + + threads = [threading.Thread(target=writer, args=(idx,)) for idx in range(10)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + with open(session_path, encoding="utf-8") as session_file: + parsed = json.load(session_file) + self.assertIn("client_id", parsed) + self.assertIn("session_id", parsed) + self.assertTrue(parsed["client_id"].startswith("client-")) + self.assertTrue(parsed["session_id"].startswith("session-")) + + with open(cookiejar_path, encoding="utf-8") as cookie_file: + cookie_contents = cookie_file.read() + self.assertEqual(cookie_contents, "cookiejar-ok\n") diff --git a/tests/test_shutdown.py b/tests/test_shutdown.py new file mode 100644 index 000000000..5d9ed81ae --- /dev/null +++ b/tests/test_shutdown.py @@ -0,0 +1,69 @@ +import json +import logging +import os +import shutil +import tempfile +from unittest import TestCase + +from icloudpd.base import ( + EXIT_CANCELLED, + ShutdownController, + _process_all_users_once, +) +from icloudpd.config import GlobalConfig +from icloudpd.log_level import LogLevel +from icloudpd.mfa_provider import MFAProvider +from icloudpd.password_provider import PasswordProvider +from icloudpd.status import StatusExchange + + +class ShutdownTestCase(TestCase): + def test_shutdown_controller_requests_stop(self) -> None: + status_exchange = StatusExchange() + shutdown = ShutdownController(status_exchange) + self.assertFalse(shutdown.requested()) + + shutdown.request_stop("SIGINT") + self.assertTrue(shutdown.requested()) + self.assertEqual(shutdown.signal_name(), "SIGINT") + self.assertTrue(status_exchange.get_progress().cancel) + self.assertFalse(shutdown.sleep_or_stop(0.01)) + + def test_cancelled_run_writes_cancelled_summary(self) -> None: + tmpdir = tempfile.mkdtemp(prefix="icloudpd-cancelled-summary-") + self.addCleanup(lambda: shutil.rmtree(tmpdir, ignore_errors=True)) + metrics_path = os.path.join(tmpdir, "metrics.json") + + global_config = GlobalConfig( + help=False, + version=False, + use_os_locale=False, + only_print_filenames=False, + log_level=LogLevel.INFO, + log_format="text", + no_progress_bar=True, + threads_num=1, + domain="com", + watch_with_interval=None, + password_providers=[PasswordProvider.PARAMETER], + mfa_provider=MFAProvider.CONSOLE, + max_retries=0, + backoff_base_seconds=1.0, + backoff_max_seconds=1.0, + respect_retry_after=True, + throttle_cooldown_seconds=0.0, + metrics_json=metrics_path, + ) + + status_exchange = StatusExchange() + shutdown = ShutdownController(status_exchange) + shutdown.request_stop("SIGTERM") + logger = logging.getLogger("icloudpd-test-shutdown") + + result = _process_all_users_once(global_config, [], logger, status_exchange, shutdown) + self.assertEqual(result, EXIT_CANCELLED) + with open(metrics_path, encoding="utf-8") as f: + summary = json.load(f) + self.assertEqual(summary["exit_code"], EXIT_CANCELLED) + self.assertEqual(summary["status"], "cancelled") + self.assertEqual(summary["users_total"], 0) diff --git a/tests/test_state_db.py b/tests/test_state_db.py new file mode 100644 index 000000000..7a16abd02 --- /dev/null +++ b/tests/test_state_db.py @@ -0,0 +1,486 @@ +import os +import sqlite3 +from datetime import datetime, timezone +from unittest import TestCase + +from icloudpd.state_db import ( + checkpoint_wal, + clear_asset_tasks_need_url_refresh, + enqueue_task, + initialize_state_db, + lease_next_task, + load_checkpoint, + mark_asset_tasks_need_url_refresh, + mark_task_done, + mark_task_failed, + prune_completed_tasks, + record_asset_checksum_result, + requeue_in_progress_tasks, + requeue_stale_leases, + resolve_state_db_path, + save_checkpoint, + upsert_asset_tasks, + vacuum_state_db, +) + + +class StateDbTestCase(TestCase): + def test_resolve_state_db_path_disabled(self) -> None: + self.assertIsNone(resolve_state_db_path(None, "~/.pyicloud")) + + def test_resolve_state_db_path_auto(self) -> None: + resolved = resolve_state_db_path("auto", "~/.pyicloud") + assert resolved is not None + self.assertTrue(resolved.endswith("/.pyicloud/icloudpd.sqlite")) + + def test_initialize_state_db_creates_schema(self) -> None: + db_path = "/tmp/icloudpd-test-state/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + + initialize_state_db(db_path) + + with sqlite3.connect(db_path) as conn: + tables = { + row[0] + for row in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name IN ('assets','tasks','checkpoints')" + ) + } + indexes = { + row[0] + for row in conn.execute( + "SELECT name FROM sqlite_master WHERE type='index' AND name LIKE 'idx_tasks_%'" + ) + } + task_columns = {row[1] for row in conn.execute("PRAGMA table_info(tasks)")} + self.assertEqual(tables, {"assets", "tasks", "checkpoints"}) + self.assertIn("idx_tasks_status_updated", indexes) + self.assertIn("idx_tasks_lease_expires", indexes) + self.assertIn("idx_tasks_library_album_status", indexes) + self.assertIn("checksum_result", task_columns) + self.assertIn("needs_url_refresh", task_columns) + + def test_initialize_state_db_is_idempotent(self) -> None: + db_path = "/tmp/icloudpd-test-state-idempotent/icloudpd.sqlite" + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + initialize_state_db(db_path) + + with sqlite3.connect(db_path) as conn: + row = conn.execute("SELECT COUNT(*) FROM sqlite_master WHERE type='table'").fetchone() + assert row is not None + self.assertGreaterEqual(row[0], 3) + + def test_lease_and_requeue_flow(self) -> None: + db_path = "/tmp/icloudpd-test-state-lease/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + enqueue_task( + db_path, + asset_id="a1", + library="PrimarySync", + album="all", + version="original", + expected_size=123, + checksum="abc", + url="https://example.com/a1", + local_path="/tmp/a1.jpg", + ) + + leased = lease_next_task( + db_path, + lease_owner="worker-1", + lease_seconds=60, + now_iso="2026-03-02T12:00:00+00:00", + ) + self.assertEqual(leased, ("a1", "PrimarySync", "all", "original")) + + # No second pending task. + self.assertIsNone( + lease_next_task( + db_path, + lease_owner="worker-2", + lease_seconds=60, + now_iso="2026-03-02T12:00:01+00:00", + ) + ) + + # Requeue after lease expiry. + requeued = requeue_stale_leases(db_path, now_iso="2026-03-02T12:02:00+00:00") + self.assertEqual(requeued, 1) + + leased_again = lease_next_task( + db_path, + lease_owner="worker-2", + lease_seconds=60, + now_iso="2026-03-02T12:03:00+00:00", + ) + self.assertEqual(leased_again, ("a1", "PrimarySync", "all", "original")) + + def test_mark_task_status(self) -> None: + db_path = "/tmp/icloudpd-test-state-status/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + enqueue_task( + db_path, + asset_id="a2", + library="PrimarySync", + album="all", + version="original", + expected_size=456, + checksum=None, + url=None, + local_path=None, + ) + lease_next_task( + db_path, + lease_owner="worker-1", + lease_seconds=60, + now_iso="2026-03-02T12:00:00+00:00", + ) + + mark_task_failed( + db_path, + asset_id="a2", + library="PrimarySync", + album="all", + version="original", + error="network error", + ) + with sqlite3.connect(db_path) as conn: + status, last_error = conn.execute( + "SELECT status, last_error FROM tasks WHERE asset_id='a2'" + ).fetchone() + self.assertEqual(status, "failed") + self.assertEqual(last_error, "network error") + + mark_task_done( + db_path, + asset_id="a2", + library="PrimarySync", + album="all", + version="original", + ) + with sqlite3.connect(db_path) as conn: + status, last_error = conn.execute( + "SELECT status, last_error FROM tasks WHERE asset_id='a2'" + ).fetchone() + self.assertEqual(status, "done") + self.assertIsNone(last_error) + + def test_requeue_in_progress_tasks(self) -> None: + db_path = "/tmp/icloudpd-test-state-requeue-in-progress/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + enqueue_task( + db_path, + asset_id="a4", + library="PrimarySync", + album="all", + version="original", + expected_size=100, + checksum=None, + url=None, + local_path=None, + ) + leased = lease_next_task( + db_path, + lease_owner="worker-x", + lease_seconds=300, + now_iso="2026-03-03T00:00:00+00:00", + ) + self.assertEqual(leased, ("a4", "PrimarySync", "all", "original")) + + requeued = requeue_in_progress_tasks(db_path, lease_owner="worker-x") + self.assertEqual(requeued, 1) + with sqlite3.connect(db_path) as conn: + status, owner = conn.execute( + "SELECT status, lease_owner FROM tasks WHERE asset_id='a4'" + ).fetchone() + self.assertEqual(status, "pending") + self.assertIsNone(owner) + + def test_prune_completed_tasks(self) -> None: + db_path = "/tmp/icloudpd-test-state-prune-completed/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + enqueue_task( + db_path, + asset_id="a5", + library="PrimarySync", + album="all", + version="original", + expected_size=10, + checksum=None, + url=None, + local_path=None, + ) + mark_task_done( + db_path, + asset_id="a5", + library="PrimarySync", + album="all", + version="original", + ) + with sqlite3.connect(db_path) as conn: + conn.execute( + "UPDATE tasks SET updated_at=datetime('now', '-5 days') WHERE asset_id='a5'" + ) + + pruned = prune_completed_tasks(db_path, older_than_days=1) + self.assertEqual(pruned, 1) + with sqlite3.connect(db_path) as conn: + remaining = conn.execute("SELECT COUNT(*) FROM tasks WHERE asset_id='a5'").fetchone()[0] + self.assertEqual(remaining, 0) + + def test_checkpoint_wal_and_vacuum(self) -> None: + db_path = "/tmp/icloudpd-test-state-maintenance/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + + result = checkpoint_wal(db_path) + self.assertEqual(len(result), 3) + vacuum_state_db(db_path) + + def test_mark_and_clear_url_refresh_marker(self) -> None: + db_path = "/tmp/icloudpd-test-state-url-refresh-marker/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + enqueue_task( + db_path, + asset_id="a6", + library="PrimarySync", + album="all", + version="original", + expected_size=1, + checksum=None, + url="https://example.test/a6", + local_path=None, + ) + marked = mark_asset_tasks_need_url_refresh( + db_path, + asset_id="a6", + library="PrimarySync", + album="all", + ) + self.assertEqual(marked, 1) + with sqlite3.connect(db_path) as conn: + value = conn.execute( + "SELECT needs_url_refresh FROM tasks WHERE asset_id='a6'" + ).fetchone()[0] + self.assertEqual(value, 1) + + cleared = clear_asset_tasks_need_url_refresh( + db_path, + asset_id="a6", + library="PrimarySync", + album="all", + ) + self.assertEqual(cleared, 1) + with sqlite3.connect(db_path) as conn: + value = conn.execute( + "SELECT needs_url_refresh FROM tasks WHERE asset_id='a6'" + ).fetchone()[0] + self.assertEqual(value, 0) + + def test_record_asset_checksum_result(self) -> None: + db_path = "/tmp/icloudpd-test-state-checksum-result/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + enqueue_task( + db_path, + asset_id="a3", + library="PrimarySync", + album="all", + version="original", + expected_size=789, + checksum="abc", + url="https://example.com/a3", + local_path="/tmp/a3.jpg", + ) + enqueue_task( + db_path, + asset_id="a3", + library="PrimarySync", + album="all", + version="medium", + expected_size=456, + checksum="def", + url="https://example.com/a3m", + local_path="/tmp/a3m.jpg", + ) + + record_asset_checksum_result( + db_path, + asset_id="a3", + library="PrimarySync", + album="all", + checksum_result="verified", + ) + + with sqlite3.connect(db_path) as conn: + rows = conn.execute( + "SELECT version, checksum_result FROM tasks WHERE asset_id='a3' ORDER BY version" + ).fetchall() + self.assertEqual(rows, [("medium", "verified"), ("original", "verified")]) + + def test_checkpoint_roundtrip(self) -> None: + db_path = "/tmp/icloudpd-test-state-checkpoint/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + + self.assertIsNone(load_checkpoint(db_path, library="PrimarySync", album="all")) + save_checkpoint(db_path, library="PrimarySync", album="all", start_rank=42) + self.assertEqual(load_checkpoint(db_path, library="PrimarySync", album="all"), 42) + save_checkpoint(db_path, library="PrimarySync", album="all", start_rank=100) + self.assertEqual(load_checkpoint(db_path, library="PrimarySync", album="all"), 100) + + def test_upsert_asset_tasks_persists_asset_and_versions(self) -> None: + db_path = "/tmp/icloudpd-test-state-asset-tasks/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + + class _ItemType: + value = "image" + + class _Version: + def __init__(self, size: int, checksum: str, url: str): + self.size = size + self.checksum = checksum + self.url = url + + class _VersionKey: + def __init__(self, value: str): + self.value = value + + class _FakePhoto: + id = "asset-1" + item_type = _ItemType() + added_date = datetime(2026, 1, 1, tzinfo=timezone.utc) + asset_date = datetime(2025, 1, 1, tzinfo=timezone.utc) + versions = { + _VersionKey("original"): _Version(100, "abc", "https://example.com/a.jpg"), + _VersionKey("medium"): _Version(50, "def", "https://example.com/m.jpg"), + } + + upsert_asset_tasks(db_path, photo=_FakePhoto(), library="PrimarySync", album="all") + + with sqlite3.connect(db_path) as conn: + assets = conn.execute("SELECT COUNT(*) FROM assets").fetchone()[0] + tasks = conn.execute("SELECT COUNT(*) FROM tasks").fetchone()[0] + self.assertEqual(assets, 1) + self.assertEqual(tasks, 2) + + def test_crash_resume_does_not_redo_completed_tasks(self) -> None: + db_path = "/tmp/icloudpd-test-state-crash-resume/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + + enqueue_task( + db_path, + asset_id="done-task", + library="PrimarySync", + album="all", + version="original", + expected_size=1, + checksum=None, + url=None, + local_path=None, + ) + enqueue_task( + db_path, + asset_id="stale-task", + library="PrimarySync", + album="all", + version="original", + expected_size=1, + checksum=None, + url=None, + local_path=None, + ) + + first = lease_next_task( + db_path, + lease_owner="worker-1", + lease_seconds=60, + now_iso="2026-03-02T12:00:00+00:00", + ) + self.assertEqual(first, ("done-task", "PrimarySync", "all", "original")) + mark_task_done( + db_path, + asset_id="done-task", + library="PrimarySync", + album="all", + version="original", + ) + + second = lease_next_task( + db_path, + lease_owner="worker-1", + lease_seconds=60, + now_iso="2026-03-02T12:00:10+00:00", + ) + self.assertEqual(second, ("stale-task", "PrimarySync", "all", "original")) + + # Simulate crash, worker disappears. Stale lease should be requeued. + self.assertEqual(requeue_stale_leases(db_path, now_iso="2026-03-02T12:05:00+00:00"), 1) + resumed = lease_next_task( + db_path, + lease_owner="worker-2", + lease_seconds=60, + now_iso="2026-03-02T12:06:00+00:00", + ) + self.assertEqual(resumed, ("stale-task", "PrimarySync", "all", "original")) + + # Completed task is not leased again. + mark_task_done( + db_path, + asset_id="stale-task", + library="PrimarySync", + album="all", + version="original", + ) + self.assertIsNone( + lease_next_task( + db_path, + lease_owner="worker-3", + lease_seconds=60, + now_iso="2026-03-02T12:07:00+00:00", + ) + ) + + def test_checkpoint_resume_after_partial_enumeration(self) -> None: + db_path = "/tmp/icloudpd-test-state-partial-enumeration/icloudpd.sqlite" + if os.path.exists(db_path): + os.remove(db_path) + os.makedirs(os.path.dirname(db_path), exist_ok=True) + initialize_state_db(db_path) + + save_checkpoint(db_path, library="PrimarySync", album="all", start_rank=100) + # Simulate restart from checkpoint and advancing. + start_rank = load_checkpoint(db_path, library="PrimarySync", album="all") + self.assertEqual(start_rank, 100) + save_checkpoint(db_path, library="PrimarySync", album="all", start_rank=150) + self.assertEqual(load_checkpoint(db_path, library="PrimarySync", album="all"), 150)