contextbase-shared-plugins 0.0.0a1__py3-none-any.whl

shared_plugins/google_client/auth.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Collection
+from datetime import datetime, timezone
+from typing import Any
+
+from google.oauth2.credentials import Credentials
+from googleapiclient.discovery import build
+from base_client import BindingAuth_AuthenticatedAccount
+from shared_plugins.provider_token import (
+    ProviderTokenClient,
+    ProviderTokenLease,
+)
+
+
+def _normalize_google_expiry(expires_at: datetime | None) -> datetime:
+    # Google leases always carry an expiry; None means web returned a
+    # no-declared-expiry lease (a notion-shaped provider) for a google account —
+    # a contract violation, not a case to tolerate.
+    if expires_at is None:
+        raise ValueError("google lease must carry an expiry.")
+    if expires_at.tzinfo is None or expires_at.utcoffset() is None:
+        raise ValueError("expires_at must be timezone-aware.")
+
+    # google.oauth2.credentials.Credentials expects a naive UTC datetime for expiry.
+    return expires_at.astimezone(timezone.utc).replace(tzinfo=None)
+
+
+def _fetch_provider_access_token(
+    auth: BindingAuth_AuthenticatedAccount,
+    provider_token_client: ProviderTokenClient,
+    *,
+    refresh: bool = False,
+) -> ProviderTokenLease:
+    # Redeem the account's provider (Google) access token server-side: web's
+    # /provider-token mints/refreshes it from the stored OAuth grant. This is the
+    # credentials keystone — it replaces the in-process control plane removed in the
+    # ctxb thin-client collapse (Chunk 5). `refresh=True` triggers real rotation
+    # (google-auth only invokes the refresh handler when the credential is already
+    # near expiry, so rotating there is always warranted). The client is the Dagster
+    # `provider_token` resource, injected at the asset boundary. Returns a
+    # ProviderTokenLease — `.access_token` (str) + `.access_token_expires_at`
+    # (tz-aware datetime | None; always set for google).
+    return provider_token_client.redeem(
+        provider_id=auth.provider_id,
+        account_id=auth.account_id,
+        refresh=refresh,
+    )
+
+
+def _build_refresh_handler(
+    *,
+    auth: BindingAuth_AuthenticatedAccount,
+    provider_token_client: ProviderTokenClient,
+) -> Callable[..., tuple[str, datetime]]:
+    def _refresh_handler(
+        request: Any,
+        *,
+        scopes: Collection[str] | None = None,
+    ) -> tuple[str, datetime]:
+        del request
+        del scopes
+
+        # google-auth only invokes this handler when it already considers the
+        # credential stale (within REFRESH_THRESHOLD = 3m45s of expiry).
+        # `refresh=True` maps to web's real rotation path (auth.api.refreshToken),
+        # so this first stale-threshold call rotates for real and the handler goes
+        # quiet until the fresh token nears ITS expiry.
+        lease = _fetch_provider_access_token(auth, provider_token_client, refresh=True)
+        return (
+            lease.access_token,
+            _normalize_google_expiry(lease.access_token_expires_at),
+        )
+
+    return _refresh_handler
+
+
+def _build_google_credentials(
+    *,
+    auth: BindingAuth_AuthenticatedAccount,
+    provider_token_client: ProviderTokenClient,
+) -> Credentials:
+    lease = _fetch_provider_access_token(auth, provider_token_client)
+    return Credentials(
+        token=lease.access_token,
+        expiry=_normalize_google_expiry(lease.access_token_expires_at),
+        refresh_handler=_build_refresh_handler(
+            auth=auth, provider_token_client=provider_token_client
+        ),
+    )
+
+
+def build_google_service(
+    *,
+    api_name: str,
+    api_version: str,
+    auth: BindingAuth_AuthenticatedAccount,
+    provider_token_client: ProviderTokenClient,
+) -> Any:
+    credentials = _build_google_credentials(
+        auth=auth, provider_token_client=provider_token_client
+    )
+    return build(
+        api_name,
+        api_version,
+        credentials=credentials,
+        cache_discovery=False,
+    )

shared_plugins/google_client/batch_retry.py

@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+import json
+import logging
+import random
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
+from typing import Any
+
+from .http_errors import extract_http_status_code
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_RETRYABLE_FORBIDDEN_REASONS = frozenset(
+    {"rateLimitExceeded", "userRateLimitExceeded"}
+)
+
+
+@dataclass(frozen=True)
+class BatchRetryPolicy:
+    max_attempts: int = 5
+    base_backoff_seconds: float = 1.0
+    max_backoff_seconds: float = 30.0
+    retryable_forbidden_reasons: frozenset[str] = _DEFAULT_RETRYABLE_FORBIDDEN_REASONS
+
+    def __post_init__(self) -> None:
+        if self.max_attempts < 1:
+            raise ValueError("max_attempts must be >= 1")
+        if self.base_backoff_seconds < 0:
+            raise ValueError("base_backoff_seconds must be >= 0")
+        if self.max_backoff_seconds < 0:
+            raise ValueError("max_backoff_seconds must be >= 0")
+
+
+@dataclass(frozen=True)
+class BatchSubrequestFailure:
+    request_id: str
+    exception: Exception
+    status_code: int | None
+    reasons: tuple[str, ...]
+
+
+class BatchRetryExhaustedError(RuntimeError):
+    def __init__(
+        self,
+        *,
+        attempts: int,
+        failures: Mapping[str, BatchSubrequestFailure],
+    ) -> None:
+        self.attempts = attempts
+        self.failures = dict(failures)
+        failed_ids = sorted(self.failures.keys())
+        sample = ", ".join(failed_ids[:5])
+        super().__init__(
+            "Batch subrequest retries exhausted after "
+            f"{attempts} attempts for {len(failed_ids)} request(s). "
+            f"Sample request IDs: {sample}"
+        )
+
+    @property
+    def failed_request_ids(self) -> tuple[str, ...]:
+        return tuple(sorted(self.failures.keys()))
+
+
+class BatchTerminalSubrequestError(RuntimeError):
+    def __init__(
+        self,
+        *,
+        request_id: str,
+        exception: Exception,
+    ) -> None:
+        self.request_id = request_id
+        self.exception = exception
+        super().__init__(
+            f"Batch subrequest {request_id} failed with a non-retryable error: {exception}"
+        )
+
+
+RequestFactory = Callable[[], Any]
+NewBatchFactory = Callable[[Callable[[str, Any, Exception | None], None]], Any]
+SleepFn = Callable[[float], None]
+RandomFn = Callable[[], float]
+
+
+def execute_batch_with_failed_subset_retries(
+    *,
+    request_factories: Mapping[str, RequestFactory],
+    new_batch: NewBatchFactory,
+    policy: BatchRetryPolicy | None = None,
+    sleep_fn: SleepFn = time.sleep,
+    random_fn: RandomFn = random.random,
+) -> dict[str, Any]:
+    """Execute a callback-based Google batch and retry only failed subrequests."""
+    if not request_factories:
+        return {}
+
+    retry_policy = policy or BatchRetryPolicy()
+    pending_ids = list(request_factories.keys())
+    responses: dict[str, Any] = {}
+    last_retryable_failures: dict[str, BatchSubrequestFailure] = {}
+
+    for attempt in range(1, retry_policy.max_attempts + 1):
+        attempt_failures: dict[str, Exception] = {}
+        retry_after_values: list[float] = []
+
+        def _callback(
+            request_id: str,
+            response: Any,
+            exception: Exception | None,
+        ) -> None:
+            if exception is not None:
+                attempt_failures[request_id] = exception
+                retry_after = _extract_retry_after_seconds(exception)
+                if retry_after is not None:
+                    retry_after_values.append(retry_after)
+                return
+            responses[request_id] = response
+
+        batch = new_batch(_callback)
+        for request_id in pending_ids:
+            batch.add(request_factories[request_id](), request_id=request_id)
+
+        try:
+            batch.execute()
+        except Exception as exc:
+            # Batch-level transport failures have no per-subrequest callback path.
+            # Treat as a transient failure for all pending requests when retryable.
+            if _is_retryable_exception(exc, retry_policy):
+                for request_id in pending_ids:
+                    attempt_failures.setdefault(request_id, exc)
+                retry_after = _extract_retry_after_seconds(exc)
+                if retry_after is not None:
+                    retry_after_values.append(retry_after)
+            else:
+                raise
+
+        if not attempt_failures:
+            return responses
+
+        status_histogram: dict[int | None, int] = {}
+        reason_histogram: dict[str, int] = {}
+        retry_after_present = 0
+        retryable_ids: list[str] = []
+        for request_id, exception in attempt_failures.items():
+            status_code = extract_http_status_code(exception)
+            reasons = tuple(sorted(_extract_error_reasons(exception)))
+            status_histogram[status_code] = status_histogram.get(status_code, 0) + 1
+            for reason in reasons:
+                reason_histogram[reason] = reason_histogram.get(reason, 0) + 1
+            if _extract_retry_after_seconds(exception) is not None:
+                retry_after_present += 1
+            if _is_retryable_exception(exception, retry_policy):
+                retryable_ids.append(request_id)
+                last_retryable_failures[request_id] = BatchSubrequestFailure(
+                    request_id=request_id,
+                    exception=exception,
+                    status_code=status_code,
+                    reasons=reasons,
+                )
+                continue
+            raise BatchTerminalSubrequestError(
+                request_id=request_id,
+                exception=exception,
+            ) from exception
+
+        if attempt >= retry_policy.max_attempts:
+            unresolved = {
+                request_id: last_retryable_failures[request_id]
+                for request_id in retryable_ids
+                if request_id in last_retryable_failures
+            }
+            raise BatchRetryExhaustedError(
+                attempts=retry_policy.max_attempts,
+                failures=unresolved,
+            )
+
+        pending_ids = retryable_ids
+        sleep_seconds = _compute_backoff_seconds(
+            attempt=attempt,
+            policy=retry_policy,
+            retry_after_values=retry_after_values,
+            random_fn=random_fn,
+        )
+        if sleep_seconds > 0:
+            logger.info(
+                "Retrying %d failed batch subrequest(s) after %.2fs (attempt %d/%d). "
+                "status_codes=%s reasons=%s retry_after_present=%d/%d",
+                len(pending_ids),
+                sleep_seconds,
+                attempt + 1,
+                retry_policy.max_attempts,
+                status_histogram,
+                reason_histogram,
+                retry_after_present,
+                len(attempt_failures),
+            )
+            sleep_fn(sleep_seconds)
+
+    return responses
+
+
+def _compute_backoff_seconds(
+    *,
+    attempt: int,
+    policy: BatchRetryPolicy,
+    retry_after_values: list[float],
+    random_fn: RandomFn,
+) -> float:
+    # Exponential backoff with jitter. attempt=1 is the first retry wait.
+    base = min(
+        policy.max_backoff_seconds,
+        policy.base_backoff_seconds * (2 ** (attempt - 1)),
+    )
+    jittered = min(policy.max_backoff_seconds, base * (1 + random_fn()))
+    retry_after = max(retry_after_values) if retry_after_values else 0.0
+    return max(jittered, retry_after)
+
+
+def _extract_error_reasons(exc: Exception) -> set[str]:
+    reasons: set[str] = set()
+
+    error_details = getattr(exc, "error_details", None)
+    _collect_reasons_from_value(error_details, reasons)
+
+    content = getattr(exc, "content", None)
+    if isinstance(content, bytes):
+        try:
+            payload = json.loads(content.decode("utf-8"))
+        except (UnicodeDecodeError, ValueError):
+            payload = None
+        _collect_reasons_from_value(payload, reasons)
+
+    return reasons
+
+
+def _collect_reasons_from_value(value: Any, out: set[str]) -> None:
+    if isinstance(value, dict):
+        reason = value.get("reason")
+        if isinstance(reason, str):
+            out.add(reason)
+        for nested in value.values():
+            _collect_reasons_from_value(nested, out)
+        return
+
+    if isinstance(value, list):
+        for nested in value:
+            _collect_reasons_from_value(nested, out)
+
+
+def _extract_retry_after_seconds(exc: Exception) -> float | None:
+    header_value = _extract_retry_after_header(exc)
+    if header_value is None:
+        return None
+
+    text = str(header_value).strip()
+    if not text:
+        return None
+
+    try:
+        seconds = float(text)
+    except ValueError:
+        seconds = _retry_after_http_date_to_seconds(text)
+    if seconds is None:
+        return None
+    return max(0.0, seconds)
+
+
+def _extract_retry_after_header(exc: Exception) -> str | None:
+    for attr in ("resp", "response"):
+        candidate = getattr(exc, attr, None)
+        if candidate is None:
+            continue
+        if hasattr(candidate, "get"):
+            for key in ("retry-after", "Retry-After"):
+                value = candidate.get(key)
+                if value is not None:
+                    return str(value)
+    return None
+
+
+def _retry_after_http_date_to_seconds(value: str) -> float | None:
+    try:
+        retry_dt = parsedate_to_datetime(value)
+    except (TypeError, ValueError, IndexError):
+        return None
+
+    if retry_dt.tzinfo is None:
+        retry_dt = retry_dt.replace(tzinfo=timezone.utc)
+    now = datetime.now(timezone.utc)
+    return (retry_dt - now).total_seconds()
+
+
+def _is_retryable_exception(exc: Exception, policy: BatchRetryPolicy) -> bool:
+    status_code = extract_http_status_code(exc)
+    if status_code is None:
+        return isinstance(exc, (TimeoutError, ConnectionError, OSError))
+
+    if status_code >= 500:
+        return True
+    if status_code in {408, 429}:
+        return True
+    if status_code == 403:
+        reasons = _extract_error_reasons(exc)
+        return any(reason in policy.retryable_forbidden_reasons for reason in reasons)
+    return False

shared_plugins/google_client/http_errors.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+
+def extract_http_status_code(exc: Exception) -> int | None:
+    status_code = _coerce_status_code(getattr(exc, "status_code", None))
+    if status_code is not None:
+        return status_code
+
+    response = getattr(exc, "response", None)
+    status_code = _coerce_status_code(getattr(response, "status_code", None))
+    if status_code is not None:
+        return status_code
+
+    resp = getattr(exc, "resp", None)
+    return _coerce_status_code(getattr(resp, "status", None))
+
+
+def _coerce_status_code(value: object) -> int | None:
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, int):
+        return value
+    if isinstance(value, str):
+        text = value.strip()
+        if text.isdigit():
+            return int(text)
+    return None