UGOS 1.0.0__py3-none-any.whl

ugos/_client_transfer.py

@@ -0,0 +1,501 @@
+"""Download, upload and pagination implementation."""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import tempfile
+from collections.abc import Collection, Iterator, Mapping
+from pathlib import Path
+from typing import Any
+
+import httpx
+
+from ._client_base import _ClientBase
+from ._client_helpers import (
+    _capture_file_positions,
+    _DownloadSpec,
+    _error_body_limit,
+    _filename_from_response,
+    _limited_chunks,
+    _merge_params,
+    _multipart_fields,
+    _resolve_destination,
+    _restore_file_positions,
+    _RetryState,
+    _sanitize_filename,
+)
+from ._codec import _looks_like_login_html, parse_response
+from ._encoding import encode_pairs, encode_params
+from ._pagination import paginate
+from ._redaction import redact_headers, redact_text
+from ._transport import FORM_HEADERS
+from ._transport import UNSET as _UNSET
+from .exceptions import (
+    ConfigurationError,
+    CSRFError,
+    DownloadError,
+    HTTPStatusError,
+    SessionExpiredError,
+)
+from .models.common import DownloadedFile
+
+
+class _TransferMixin(_ClientBase):
+    def _build_download_request(
+        self,
+        spec: _DownloadSpec,
+    ) -> tuple[bytes | str | None, list[tuple[str, str]] | None]:
+        request_params = dict(spec.params)
+        if spec.method == "GET":
+            return None, encode_pairs(request_params, json_fields=spec.json_fields)
+        with self._lock:
+            request_params["csrfmiddlewaretoken"] = self._auth.ensure_csrf()
+        body = encode_params(request_params, json_fields=spec.json_fields)
+        return body, None
+
+    def _raise_download_status(self, response: httpx.Response, spec: _DownloadSpec) -> None:
+        status = response.status_code
+        if 300 <= status < 400:
+            location = response.headers.get("location", "")
+            if "utm/login" in location or location.rstrip("/").endswith("/login"):
+                raise SessionExpiredError(
+                    "session expired or login required",
+                    status_code=status,
+                    controller=spec.controller,
+                    action=spec.action,
+                )
+            raise HTTPStatusError(
+                f"unexpected redirect: HTTP {status}",
+                status_code=status,
+                controller=spec.controller,
+                action=spec.action,
+            )
+        if status == 401:
+            raise SessionExpiredError(
+                "unauthorized",
+                status_code=status,
+                controller=spec.controller,
+                action=spec.action,
+            )
+
+    def _raise_download_html(
+        self,
+        response: httpx.Response,
+        first: bytes,
+        raw_chunks: Iterator[bytes],
+        spec: _DownloadSpec,
+    ) -> None:
+        payload_bytes = b"".join(
+            _limited_chunks(
+                first,
+                raw_chunks,
+                limit=_error_body_limit(spec.limit),
+                controller=spec.controller,
+                action=spec.action,
+            )
+        )
+        text = payload_bytes.decode(response.encoding or "utf-8", errors="replace")
+        if _looks_like_login_html(text):
+            raise SessionExpiredError(
+                "session expired or a login page was returned",
+                status_code=response.status_code,
+                controller=spec.controller,
+                action=spec.action,
+            )
+        raise DownloadError(
+            "expected a file, got HTML",
+            status_code=response.status_code,
+            controller=spec.controller,
+            action=spec.action,
+        )
+
+    def _is_json_like(self, content_type: str, head: bytes) -> bool:
+        if "json" in content_type.lower():
+            return True
+        return head.lstrip().startswith((b"{", b"[")) and any(
+            marker in head for marker in (b'"success"', b'"errorCode"', b'"error_code"', b'"error"')
+        )
+
+    def _raise_download_payload(
+        self,
+        response: httpx.Response,
+        content_type: str,
+        first: bytes,
+        raw_chunks: Iterator[bytes],
+        spec: _DownloadSpec,
+    ) -> None:
+        payload_bytes = b"".join(
+            _limited_chunks(
+                first,
+                raw_chunks,
+                limit=_error_body_limit(spec.limit),
+                controller=spec.controller,
+                action=spec.action,
+            )
+        )
+        with self._lock:
+            self._auth.note_response(response)
+        synthetic = httpx.Response(
+            response.status_code,
+            headers={"content-type": content_type},
+            content=payload_bytes,
+        )
+        data = parse_response(
+            synthetic,
+            expect="auto",
+            controller=spec.controller,
+            action=spec.action,
+        )
+        raise DownloadError(
+            f"expected a file, got JSON: {redact_text(str(data))[:500]}",
+            controller=spec.controller,
+            action=spec.action,
+        )
+
+    def _download_to_memory(
+        self,
+        first: bytes,
+        raw_chunks: Iterator[bytes],
+        spec: _DownloadSpec,
+        content_type: str,
+        filename: str,
+        headers: Mapping[str, str],
+    ) -> DownloadedFile:
+        content = b"".join(
+            _limited_chunks(
+                first,
+                raw_chunks,
+                limit=spec.limit,
+                controller=spec.controller,
+                action=spec.action,
+            )
+        )
+        return DownloadedFile(
+            content=content,
+            filename=filename,
+            content_type=content_type or None,
+            headers=headers,
+        )
+
+    def _write_download(
+        self,
+        target: Path,
+        first: bytes,
+        raw_chunks: Iterator[bytes],
+        spec: _DownloadSpec,
+    ) -> None:
+        target.parent.mkdir(parents=True, exist_ok=True)
+        fd, temp_name = tempfile.mkstemp(
+            dir=target.parent,
+            prefix=f".{target.name}.",
+            suffix=".part",
+        )
+        temp = Path(temp_name)
+        try:
+            with os.fdopen(fd, "wb") as handle:
+                for chunk in _limited_chunks(
+                    first,
+                    raw_chunks,
+                    limit=spec.limit,
+                    controller=spec.controller,
+                    action=spec.action,
+                ):
+                    handle.write(chunk)
+            os.chmod(temp, 0o600)
+            os.replace(temp, target)
+        except Exception:
+            with contextlib.suppress(OSError):
+                temp.unlink()
+            raise
+
+    def _download_to_path(
+        self,
+        target: Path,
+        first: bytes,
+        raw_chunks: Iterator[bytes],
+        spec: _DownloadSpec,
+        content_type: str,
+        filename: str,
+        headers: Mapping[str, str],
+    ) -> DownloadedFile:
+        if target.exists() and not spec.overwrite:
+            raise FileExistsError(target)
+        self._write_download(target, first, raw_chunks, spec)
+        return DownloadedFile(
+            content=None,
+            filename=target.name,
+            content_type=content_type or None,
+            headers=headers,
+            path=target,
+        )
+
+    def _is_download_error(
+        self,
+        response: httpx.Response,
+        content_type: str,
+        first: bytes,
+    ) -> bool:
+        if response.status_code >= 400:
+            return True
+        disposition = response.headers.get("content-disposition", "").lower()
+        if "attachment" in disposition:
+            return False
+        return self._is_json_like(content_type, first.lstrip()[:512])
+
+    def _finish_download(
+        self,
+        response: httpx.Response,
+        spec: _DownloadSpec,
+        first: bytes,
+        raw_chunks: Iterator[bytes],
+        content_type: str,
+    ) -> DownloadedFile:
+        with self._lock:
+            self._auth.note_response(response)
+        name = _sanitize_filename(spec.filename or _filename_from_response(response, spec.path))
+        headers = redact_headers(dict(response.headers))
+        if spec.dest is None:
+            return self._download_to_memory(first, raw_chunks, spec, content_type, name, headers)
+        target = _resolve_destination(spec.dest, name)
+        return self._download_to_path(target, first, raw_chunks, spec, content_type, name, headers)
+
+    def _handle_download_response(
+        self,
+        response: httpx.Response,
+        spec: _DownloadSpec,
+    ) -> DownloadedFile:
+        self._raise_download_status(response, spec)
+        content_type = response.headers.get("content-type", "")
+        raw_chunks = response.iter_bytes(chunk_size=spec.chunk_size)
+        first = next(raw_chunks, b"")
+        if "html" in content_type.lower():
+            self._raise_download_html(response, first, raw_chunks, spec)
+        if self._is_download_error(response, content_type, first):
+            self._raise_download_payload(response, content_type, first, raw_chunks, spec)
+        return self._finish_download(response, spec, first, raw_chunks, content_type)
+
+    def _retry_download(
+        self,
+        exc: Exception,
+        state: _RetryState,
+        generation: int,
+        spec: _DownloadSpec,
+    ) -> bool:
+        if isinstance(exc, CSRFError):
+            return self._retry_csrf(state)
+        return self._retry_session(
+            state,
+            generation,
+            spec.action,
+            spec.method in {"GET", "HEAD", "OPTIONS"},
+            exc,
+        )
+
+    def download(
+        self,
+        controller: str,
+        action: str,
+        /,
+        *,
+        dest: str | Path | None = None,
+        filename: str | None = None,
+        overwrite: bool = False,
+        method: str = "GET",
+        payload: Mapping[str, Any] | None = None,
+        json_fields: Collection[str] = (),
+        timeout: float | httpx.Timeout | object | None = _UNSET,
+        strict: bool | None = None,
+        chunk_size: int = 64 * 1024,
+        max_bytes: int | None = None,
+        **params: Any,
+    ) -> DownloadedFile:
+        """Download a file and optionally save it to ``dest``."""
+
+        self._validate(controller, action, strict=strict)
+        limit = max_bytes if max_bytes is not None else self._config.max_download_bytes
+        if limit is not None and limit < 0:
+            raise ValueError("max_bytes must be non-negative")
+        spec = _DownloadSpec(
+            controller=controller,
+            action=action,
+            path=f"{controller}/{action}",
+            method=method.upper(),
+            params=_merge_params(payload, params),
+            json_fields=json_fields,
+            timeout=timeout,
+            chunk_size=chunk_size,
+            limit=limit,
+            dest=dest,
+            filename=filename,
+            overwrite=overwrite,
+        )
+        with self._lock:
+            self._prepare()
+        return self._download_loop(spec)
+
+    def _download_loop(self, spec: _DownloadSpec) -> DownloadedFile:
+        state = _RetryState()
+        while True:
+            with self._lock:
+                generation = self._auth.generation
+            body, query = self._build_download_request(spec)
+            try:
+                with self._transport.stream(
+                    spec.method,
+                    spec.path,
+                    params=query,
+                    data=body,
+                    headers={**(FORM_HEADERS if body else {}), "Accept": "*/*"},
+                    timeout=spec.timeout,
+                    idempotent=spec.method in {"GET", "HEAD", "OPTIONS"},
+                ) as response:
+                    return self._handle_download_response(response, spec)
+            except (CSRFError, SessionExpiredError) as exc:
+                if not self._retry_download(exc, state, generation, spec):
+                    raise
+
+    @staticmethod
+    def _restore_upload_positions(positions: list[tuple[Any, int]] | None) -> None:
+        if positions is None:
+            raise ConfigurationError("cannot retry upload with non-seekable files")
+        _restore_file_positions(positions)
+
+    def _upload_once(
+        self,
+        controller: str,
+        action: str,
+        files: Mapping[str, Any],
+        merged: Mapping[str, Any],
+        json_fields: Collection[str],
+        query: Mapping[str, Any] | None,
+        timeout: Any,
+    ) -> Any:
+        with self._lock:
+            token = self._auth.ensure_csrf()
+        data = _multipart_fields(
+            encode_pairs(
+                {**merged, "csrfmiddlewaretoken": token},
+                json_fields=json_fields,
+            )
+        )
+        query_params = encode_pairs(query, json_fields=json_fields) if query else None
+        with self._lock:
+            generation = self._auth.generation
+        response = self._transport.request(
+            "POST",
+            f"{controller}/{action}",
+            params=query_params,
+            data=data,
+            files=files,
+            headers={"Accept": "application/json, text/plain, */*"},
+            idempotent=False,
+            timeout=timeout,
+        )
+        with self._lock:
+            self._auth.note_response(response)
+        try:
+            return parse_response(
+                response,
+                expect="auto",
+                controller=controller,
+                action=action,
+            )
+        except SessionExpiredError:
+            self._refresh_session(generation)
+            raise
+
+    def _upload_loop(
+        self,
+        controller: str,
+        action: str,
+        files: Mapping[str, Any],
+        merged: Mapping[str, Any],
+        json_fields: Collection[str],
+        query: Mapping[str, Any] | None,
+        timeout: Any,
+        positions: list[tuple[Any, int]] | None,
+    ) -> Any:
+        csrf_retries = 0
+        attempt = 0
+        while True:
+            if attempt > 0:
+                self._restore_upload_positions(positions)
+            attempt += 1
+            try:
+                return self._upload_once(
+                    controller,
+                    action,
+                    files,
+                    merged,
+                    json_fields,
+                    query,
+                    timeout,
+                )
+            except CSRFError:
+                if csrf_retries >= 1:
+                    raise
+                csrf_retries += 1
+                with self._lock:
+                    self._auth.refresh_csrf()
+
+    def upload(
+        self,
+        controller: str,
+        action: str,
+        /,
+        *,
+        files: Mapping[str, Any],
+        payload: Mapping[str, Any] | None = None,
+        json_fields: Collection[str] = (),
+        query: Mapping[str, Any] | None = None,
+        timeout: float | httpx.Timeout | object | None = _UNSET,
+        strict: bool | None = None,
+        **params: Any,
+    ) -> Any:
+        """Upload files using ``multipart/form-data``."""
+
+        self._validate(controller, action, strict=strict)
+        merged = _merge_params(payload, params)
+        with self._lock:
+            self._prepare()
+        return self._upload_loop(
+            controller,
+            action,
+            files,
+            merged,
+            json_fields,
+            query,
+            timeout,
+            _capture_file_positions(files),
+        )
+
+    def iter_all(
+        self,
+        controller: str,
+        action: str,
+        /,
+        *,
+        page_size: int = 100,
+        start: int = 0,
+        max_items: int | None = None,
+        max_pages: int = 1000,
+        payload: Mapping[str, Any] | None = None,
+        strict: bool | None = None,
+        **params: Any,
+    ) -> Iterator[Any]:
+        """Iterate over all items of a paginated list action."""
+
+        self._validate(controller, action, strict=strict)
+        merged = _merge_params(payload, params)
+        return paginate(
+            self,
+            controller,
+            action,
+            page_size=page_size,
+            start=start,
+            max_items=max_items,
+            max_pages=max_pages,
+            strict=strict,
+            **merged,
+        )