pangea-sdk 6.5.0b1__py3-none-any.whl → 6.6.0__py3-none-any.whl

pangea/request.py

@@ -11,18 +11,20 @@ import json
 import logging
 import time
 from collections.abc import Iterable, Mapping
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast, overload
+from random import random
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 import requests
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel
 from pydantic_core import to_jsonable_python
-from requests.adapters import HTTPAdapter, Retry
+from requests.adapters import HTTPAdapter
 from requests_toolbelt import MultipartDecoder  # type: ignore[import-untyped]
-from typing_extensions import Literal, TypeAlias, TypeVar, override
+from typing_extensions import TypeAlias, TypeVar, override
 from yarl import URL
 
 import pangea
 import pangea.exceptions as pe
+from pangea._constants import MAX_RETRY_DELAY, RETRYABLE_HTTP_CODES
 from pangea.config import PangeaConfig
 from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
 from pangea.utils import default_encoder
@@ -46,6 +48,70 @@ _Files: TypeAlias = Union[Mapping[str, _FileSpec], Iterable[tuple[str, _FileSpec
 _HeadersUpdateMapping: TypeAlias = Mapping[str, str]
 
 
+class PangeaHTTPAdapter(HTTPAdapter):
+    """Custom HTTP adapter that keeps track of retried request IDs."""
+
+    @override
+    def __init__(self, config: PangeaConfig, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.config = config
+
+    @override
+    def send(
+        self,
+        request: requests.PreparedRequest,
+        stream: bool = False,
+        timeout: None | float | tuple[float, float] | tuple[float, None] = None,
+        verify: bool | str = True,
+        cert: None | bytes | str | tuple[bytes | str, bytes | str] = None,
+        proxies: Mapping[str, str] | None = None,
+    ) -> requests.Response:
+        max_retries = self.config.request_retries
+        request_ids = set[str]()
+        retries_taken = 0
+        for retries_taken in range(max_retries + 1):
+            remaining_retries = max_retries - retries_taken
+
+            if len(request_ids) > 0:
+                request.headers["X-Pangea-Retried-Request-Ids"] = ",".join(request_ids)
+
+            response = super().send(request, stream, timeout, verify, cert, proxies)
+
+            request_id = response.headers.get("x-request-id")
+            if request_id:
+                request_ids.add(request_id)
+
+            try:
+                response.raise_for_status()
+            except requests.HTTPError as error:
+                if remaining_retries > 0 and self._should_retry(error.response):
+                    error.response.close()
+                    self._sleep_for_retry(retries_taken=retries_taken, max_retries=max_retries, config=self.config)
+                    continue
+
+                break
+
+            break
+
+        return response
+
+    def _calculate_retry_timeout(self, remaining_retries: int, config: PangeaConfig) -> float:
+        max_retries = config.request_retries
+        nb_retries = min(max_retries - remaining_retries, 1000)
+        sleep_seconds = min(config.request_backoff * pow(2.0, nb_retries), MAX_RETRY_DELAY)
+        jitter = 1 - 0.25 * random()
+        timeout = sleep_seconds * jitter
+        return max(timeout, 0)
+
+    def _sleep_for_retry(self, *, retries_taken: int, max_retries: int, config: PangeaConfig) -> None:
+        remaining_retries = max_retries - retries_taken
+        timeout = self._calculate_retry_timeout(remaining_retries, config)
+        time.sleep(timeout)
+
+    def _should_retry(self, response: requests.Response) -> bool:
+        return response.status_code in RETRYABLE_HTTP_CODES
+
+
 class MultipartResponse:
     pangea_json: Dict[str, str]
     attached_files: List = []
@@ -111,8 +177,8 @@ class PangeaRequestBase:
 
         return self._queued_retry_enabled
 
-    def _get_delay(self, retry_count, start):
-        delay = retry_count * retry_count
+    def _get_delay(self, retry_count: int, start: float) -> float:
+        delay: float = retry_count * retry_count
         now = time.time()
         # if with this delay exceed timeout, reduce delay
         if now - start + delay >= self.config.poll_result_timeout:
@@ -120,10 +186,10 @@ class PangeaRequestBase:
 
         return delay
 
-    def _reach_timeout(self, start):
+    def _reach_timeout(self, start: float) -> bool:
         return time.time() - start >= self.config.poll_result_timeout
 
-    def _get_poll_path(self, request_id: str):
+    def _get_poll_path(self, request_id: str) -> str:
         return f"request/{request_id}"
 
     def _url(self, path: str) -> str:
@@ -220,24 +286,6 @@ class PangeaRequest(PangeaRequestBase):
     def __del__(self) -> None:
         self.session.close()
 
-    def delete(self, endpoint: str) -> None:
-        """
-        Makes a DELETE call to a Pangea endpoint.
-
-        Args:
-            endpoint: The Pangea API endpoint.
-        """
-
-        url = self._url(endpoint)
-
-        self.logger.debug(
-            json.dumps({"service": self.service, "action": "delete", "url": url}, default=default_encoder)
-        )
-
-        requests_response = self._http_delete(url, headers=self._headers())
-        self._check_http_errors(requests_response)
-
-    @overload
     def post(
         self,
         endpoint: str,
@@ -246,62 +294,18 @@ class PangeaRequest(PangeaRequestBase):
         files: Optional[list[Tuple]] = None,
         poll_result: bool = True,
         url: Optional[str] = None,
-        *,
-        pangea_response: Literal[True] = True,
     ) -> PangeaResponse[TResult]:
-        """
-        Makes the POST call to a Pangea Service endpoint.
+        """Makes the POST call to a Pangea Service endpoint.
 
         Args:
-            endpoint: The Pangea Service API endpoint.
-            data: The POST body payload object
+            endpoint(str): The Pangea Service API endpoint.
+            data(dict): The POST body payload object
 
         Returns:
             PangeaResponse which contains the response in its entirety and
                various properties to retrieve individual fields
         """
 
-    @overload
-    def post(
-        self,
-        endpoint: str,
-        result_class: Type[TResult],
-        data: str | BaseModel | Mapping[str, Any] | None = None,
-        files: Optional[list[Tuple]] = None,
-        poll_result: bool = True,
-        url: Optional[str] = None,
-        *,
-        pangea_response: Literal[False],
-    ) -> TResult:
-        """
-        Makes the POST call to a Pangea Service endpoint.
-
-        Args:
-            endpoint: The Pangea Service API endpoint.
-            data: The POST body payload object
-        """
-
-    def post(
-        self,
-        endpoint: str,
-        result_class: Type[TResult],
-        data: str | BaseModel | Mapping[str, Any] | None = None,
-        files: Optional[list[Tuple]] = None,
-        poll_result: bool = True,
-        url: Optional[str] = None,
-        *,
-        pangea_response: bool = True,
-    ) -> PangeaResponse[TResult] | TResult:
-        """
-        Makes a POST call to a Pangea Service endpoint.
-
-        Args:
-            endpoint: The Pangea Service API endpoint.
-            data: The POST body payload object
-            pangea_response: Whether or not the response body follows Pangea's
-              standard response schema
-        """
-
         if isinstance(data, BaseModel):
             data = data.model_dump(exclude_none=True)
 
@@ -339,13 +343,9 @@ class PangeaRequest(PangeaRequestBase):
 
         self._check_http_errors(requests_response)
 
-        if not pangea_response:
-            type_adapter = TypeAdapter(result_class)
-            return type_adapter.validate_python(requests_response.json())
-
         if "multipart/form-data" in requests_response.headers.get("content-type", ""):
             multipart_response = self._process_multipart_response(requests_response)
-            pangea_response_obj: PangeaResponse = PangeaResponse(
+            pangea_response: PangeaResponse = PangeaResponse(
                 requests_response,
                 result_class=result_class,
                 json=multipart_response.pangea_json,
@@ -358,14 +358,14 @@ class PangeaRequest(PangeaRequestBase):
                     json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
                 )
 
-                pangea_response_obj = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
+                pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
             except requests.exceptions.JSONDecodeError as e:
                 raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {requests_response.text}") from e
 
         if poll_result:
-            pangea_response_obj = self._handle_queued_result(pangea_response_obj)
+            pangea_response = self._handle_queued_result(pangea_response)
 
-        return self._check_response(pangea_response_obj)
+        return self._check_response(pangea_response)
 
     def _get_pangea_json(self, decoder: MultipartDecoder) -> Optional[Dict]:
         # Iterate through parts
@@ -408,14 +408,6 @@ class PangeaRequest(PangeaRequestBase):
         if resp.status_code == 503:
             raise pe.ServiceTemporarilyUnavailable(resp.json())
 
-    def _http_delete(
-        self,
-        url: str,
-        *,
-        headers: Mapping[str, str | bytes | None] = {},
-    ) -> requests.Response:
-        return self.session.delete(url, headers=headers)
-
     def _http_post(
         self,
         url: str,
@@ -460,98 +452,37 @@ class PangeaRequest(PangeaRequestBase):
 
         return response
 
-    @overload
-    def get(
-        self,
-        path: str,
-        result_class: Type[TResult],
-        check_response: bool = True,
-        *,
-        params: (
-            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
-            | None
-        ) = None,
-        pangea_response: Literal[True] = True,
-    ) -> PangeaResponse[TResult]:
-        """
-        Makes the GET call to a Pangea Service endpoint.
+    def get(self, path: str, result_class: Type[TResult], check_response: bool = True) -> PangeaResponse[TResult]:
+        """Makes the GET call to a Pangea Service endpoint.
 
         Args:
-            path: Additional URL path
-            params: Dictionary of querystring data to attach to the request
+            endpoint(str): The Pangea Service API endpoint.
+            path(str): Additional URL path
 
         Returns:
             PangeaResponse which contains the response in its entirety and
                various properties to retrieve individual fields
         """
 
-    @overload
-    def get(
-        self,
-        path: str,
-        result_class: Type[TResult],
-        check_response: bool = True,
-        *,
-        params: (
-            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
-            | None
-        ) = None,
-        pangea_response: Literal[False] = False,
-    ) -> TResult:
-        """
-        Makes the GET call to a Pangea Service endpoint.
-
-        Args:
-            path: Additional URL path
-            params: Dictionary of querystring data to attach to the request
-        """
-
-    def get(
-        self,
-        path: str,
-        result_class: Type[TResult],
-        check_response: bool = True,
-        *,
-        params: (
-            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
-            | None
-        ) = None,
-        pangea_response: bool = True,
-    ) -> PangeaResponse[TResult] | TResult:
-        """
-        Makes the GET call to a Pangea Service endpoint.
-
-        Args:
-            path: Additional URL path
-            params: Dictionary of querystring data to attach to the request
-            pangea_response: Whether or not the response body follows Pangea's
-              standard response schema
-        """
-
         url = self._url(path)
         self.logger.debug(json.dumps({"service": self.service, "action": "get", "url": url}))
-        requests_response = self.session.get(url, params=params, headers=self._headers())
+        requests_response = self.session.get(url, headers=self._headers())
         self._check_http_errors(requests_response)
-
-        if not pangea_response:
-            type_adapter = TypeAdapter(result_class)
-            return type_adapter.validate_python(requests_response.json())
-
-        pangea_response_obj: PangeaResponse = PangeaResponse(
+        pangea_response: PangeaResponse = PangeaResponse(
             requests_response, result_class=result_class, json=requests_response.json()
         )
 
         self.logger.debug(
             json.dumps(
-                {"service": self.service, "action": "get", "url": url, "response": pangea_response_obj.json},
+                {"service": self.service, "action": "get", "url": url, "response": pangea_response.json},
                 default=default_encoder,
             )
         )
 
         if check_response is False:
-            return pangea_response_obj
+            return pangea_response
 
-        return self._check_response(pangea_response_obj)
+        return self._check_response(pangea_response)
 
     def download_file(self, url: str, filename: str | None = None) -> AttachedFile:
         """
@@ -763,13 +694,7 @@ class PangeaRequest(PangeaRequestBase):
 
     @override
     def _init_session(self) -> requests.Session:
-        retry_config = Retry(
-            total=self.config.request_retries,
-            backoff_factor=self.config.request_backoff,
-            status_forcelist=[500, 502, 503, 504],
-        )
-
-        adapter = HTTPAdapter(max_retries=retry_config)
+        adapter = PangeaHTTPAdapter(config=self.config)
         session = requests.Session()
 
         session.mount("http://", adapter)

pangea/response.py

@@ -18,6 +18,8 @@ from typing_extensions import TypeVar
 
 from pangea.utils import format_datetime
 
+__all__ = ("PangeaResponse", "PangeaResponseResult", "TransferMethod")
+
 
 class AttachedFile:
     filename: str
@@ -243,7 +245,9 @@ class PangeaResponse(ResponseHeader, Generic[T]):
 
     @property
     def http_status(self) -> int:  # type: ignore[return]
-        if self.raw_response:
+        # Must be an explicit None check because Response's boolean
+        # representation is equal to whether or not the response is OK.
+        if self.raw_response is not None:
             if isinstance(self.raw_response, aiohttp.ClientResponse):
                 return self.raw_response.status
             else:

pangea/services/__init__.py

@@ -1,5 +1,3 @@
-# ruff: noqa: F401
-
 from .ai_guard import AIGuard
 from .audit.audit import Audit
 from .authn.authn import AuthN
@@ -7,9 +5,27 @@ from .authz import AuthZ
 from .embargo import Embargo
 from .file_scan import FileScan
 from .intel import DomainIntel, FileIntel, IpIntel, UrlIntel, UserIntel
-from .management import Management
 from .prompt_guard import PromptGuard
 from .redact import Redact
 from .sanitize import Sanitize
 from .share.share import Share
 from .vault.vault import Vault
+
+__all__ = (
+    "AIGuard",
+    "Audit",
+    "AuthN",
+    "AuthZ",
+    "DomainIntel",
+    "Embargo",
+    "FileIntel",
+    "FileScan",
+    "IpIntel",
+    "PromptGuard",
+    "Redact",
+    "Sanitize",
+    "Share",
+    "UrlIntel",
+    "UserIntel",
+    "Vault",
+)