Core Module¶

class AuthConfig(BaseModel):
    """Authentication configuration for a REST API client."""

    auth_type: AuthType = Field(AuthType.NONE, description="Authentication method")
    api_key: Optional[SecretStr] = Field(None, description="API key or Bearer token")
    api_key_header: str = Field("X-Api-Key", description="Header name for API key auth")
    api_key_param: str = Field("apikey", description="Query param name for API key auth")
    username: Optional[str] = Field(None, description="Username for Basic auth")
    password: Optional[SecretStr] = Field(None, description="Password for Basic auth")

    class Config:
        frozen = True

    def build(self) -> Tuple[dict[str, str], dict[str, str], Any]:
        """
        Resolve auth config into (headers, query_params, httpx_auth).

        Returns
        -------
        Tuple[dict, dict, Any]
            headers, query_params, and httpx auth object (or None).
        """
        headers: dict[str, str] = {}
        params: dict[str, str] = {}
        auth: Any = None

        if self.auth_type == AuthType.BEARER:
            headers["Authorization"] = f"Bearer {self.api_key.get_secret_value()}"
        elif self.auth_type == AuthType.API_KEY_HEADER:
            headers[self.api_key_header] = self.api_key.get_secret_value()
        elif self.auth_type == AuthType.API_KEY_QUERY:
            params[self.api_key_param] = self.api_key.get_secret_value()
        elif self.auth_type == AuthType.BASIC:
            auth = (self.username, self.password.get_secret_value())

        return headers, params, auth

class BasePaginationStrategy(ABC):
    """
    Abstract base for pagination strategies.

    Subclass this to implement cursor-based, offset/limit,
    page-number, or Link-header pagination.
    """

    @abstractmethod
    def next_request(
        self,
        response: httpx.Response,
        current_params: dict[str, Any],
    ) -> Optional[dict[str, Any]]:
        """
        Return updated params for the next page, or None if exhausted.

        Parameters
        ----------
        response : httpx.Response
            The most recent response from the API.
        current_params : dict
            The params used for the current request.

        Returns
        -------
        Optional[dict]
            Params for the next request, or None to stop pagination.
        """
        ...

    @abstractmethod
    def extract_records(self, response: httpx.Response) -> list[dict]:
        """Extract the list of records from a response."""
        ...

class BaseRestApiClient(ABC):
    """
    Abstract base class for REST API clients in GigaSpatial.

    Composes authentication (AuthConfig), retry logic, and pagination
    (BasePaginationStrategy) into a reusable context-managed client.
    Subclasses define the pagination strategy and any endpoint-specific methods.

    Parameters
    ----------
    config : RestApiClientConfig
        Pydantic configuration for the client.

    Examples
    --------
    >>> config = RestApiClientConfig(
    ...     base_url="https://api.example.com",
    ...     auth=AuthConfig(auth_type=AuthType.BEARER, api_key="secret"),
    ... )
    >>> with MyApiClient(config) as client:
    ...     for page in client.paginate("/schools", params={"country": "KE"}):
    ...         process(page)
    """

    def __init__(self, config: RestApiClientConfig) -> None:
        self.config = config
        self._client: Optional[httpx.Client] = None

    # ------------------------------------------------------------------
    # Pagination strategy — subclasses define this
    # ------------------------------------------------------------------

    @property
    @abstractmethod
    def pagination_strategy(self) -> BasePaginationStrategy:
        """Return the pagination strategy for this API."""
        ...

    # ------------------------------------------------------------------
    # Session lifecycle
    # ------------------------------------------------------------------

    def _build_client(self) -> httpx.Client:
        auth_headers, auth_params, httpx_auth = self.config.auth.build()
        headers = {**self.config.default_headers, **auth_headers}
        return httpx.Client(
            base_url=self.config.base_url,
            headers=headers,
            params=auth_params,
            auth=httpx_auth,
            timeout=self.config.timeout,
        )

    def __enter__(self) -> "BaseRestApiClient":
        self._client = self._build_client()
        return self

    def __exit__(self, *args: Any) -> None:
        if self._client:
            self._client.close()
            self._client = None

    # ------------------------------------------------------------------
    # Core request with retry + rate-limit handling
    # ------------------------------------------------------------------

    def request(self, method: str, endpoint: str, **kwargs: Any) -> httpx.Response:
        """
        Send an HTTP request with automatic retries and backoff.

        Parameters
        ----------
        method : str
            HTTP method (GET, POST, etc.).
        endpoint : str
            API endpoint path relative to base_url.
        **kwargs
            Passed directly to httpx.Client.request().

        Returns
        -------
        httpx.Response

        Raises
        ------
        httpx.HTTPStatusError
            If the request fails after all retries.
        """
        assert self._client is not None, "Client not started — use as a context manager."
        delay = 1.0

        for attempt in range(1, self.config.max_retries + 1):
            response = self._client.request(method, endpoint, **kwargs)

            if response.status_code == 429:
                wait = float(response.headers.get("Retry-After", delay))
                logger.warning("Rate limited. Waiting %.1fs (attempt %d/%d)", wait, attempt, self.config.max_retries)
                time.sleep(wait)
                delay *= self.config.retry_backoff
                continue

            if response.status_code >= 500 and attempt < self.config.max_retries:
                logger.warning("Server error %d. Retrying in %.1fs (attempt %d/%d)", response.status_code, delay, attempt, self.config.max_retries)
                time.sleep(delay)
                delay *= self.config.retry_backoff
                continue

            response.raise_for_status()
            return response

        response.raise_for_status()
        return response

    def get(self, endpoint: str, **kwargs: Any) -> httpx.Response:
        """Convenience wrapper for GET requests."""
        return self.request("GET", endpoint, **kwargs)

    def post(self, endpoint: str, **kwargs: Any) -> httpx.Response:
        """Convenience wrapper for POST requests."""
        return self.request("POST", endpoint, **kwargs)

    # ------------------------------------------------------------------
    # Pagination
    # ------------------------------------------------------------------

    def paginate(
        self,
        endpoint: str,
        params: Optional[dict[str, Any]] = None,
    ) -> Generator[list[dict], None, None]:
        """
        Yield pages of records using the configured pagination strategy.

        Parameters
        ----------
        endpoint : str
            API endpoint to paginate.
        params : dict, optional
            Initial query parameters.

        Yields
        ------
        list[dict]
            One page of records per iteration.
        """
        current_params = params or {}
        strategy = self.pagination_strategy

        while True:
            response = self.get(endpoint, params=current_params)
            records = strategy.extract_records(response)
            if not records:
                break
            yield records
            next_params = strategy.next_request(response, current_params)
            if next_params is None:
                break
            current_params = next_params

class CursorPagination(BasePaginationStrategy):
    """
    Cursor-based pagination (e.g. GitHub, Mapbox APIs).

    Parameters
    ----------
    cursor_response_key : str
        JSON key in the response body containing the next cursor.
    cursor_param : str
        Query parameter name to pass the cursor on the next request.
    records_key : str
        JSON key containing the list of records.
    """

    def __init__(
        self,
        cursor_response_key: str = "next_cursor",
        cursor_param: str = "cursor",
        records_key: str = "results",
    ) -> None:
        self.cursor_response_key = cursor_response_key
        self.cursor_param = cursor_param
        self.records_key = records_key

    def next_request(
        self,
        response: httpx.Response,
        current_params: dict[str, Any],
    ) -> Optional[dict[str, Any]]:
        cursor = response.json().get(self.cursor_response_key)
        if not cursor:
            return None
        return {**current_params, self.cursor_param: cursor}

    def extract_records(self, response: httpx.Response) -> list[dict]:
        return response.json().get(self.records_key, [])

class OffsetPagination(BasePaginationStrategy):
    """
    Standard offset/limit pagination.

    Parameters
    ----------
    page_size : int
        Number of records per page.
    offset_param : str
        Query parameter name for the offset.
    limit_param : str
        Query parameter name for the limit/page size.
    records_key : str
        JSON key containing the list of records.
    """

    def __init__(
        self,
        page_size: int = 100,
        offset_param: str = "offset",
        limit_param: str = "limit",
        records_key: str = "results",
    ) -> None:
        self.page_size = page_size
        self.offset_param = offset_param
        self.limit_param = limit_param
        self.records_key = records_key

    def next_request(
        self,
        response: httpx.Response,
        current_params: dict[str, Any],
    ) -> Optional[dict[str, Any]]:
        records = self.extract_records(response)
        if len(records) < self.page_size:
            return None  # last page
        next_offset = current_params.get(self.offset_param, 0) + self.page_size
        return {**current_params, self.offset_param: next_offset, self.limit_param: self.page_size}

    def extract_records(self, response: httpx.Response) -> list[dict]:
        return response.json().get(self.records_key, [])

class PageNumberPagination(BasePaginationStrategy):
    """
    Page-number pagination (page=1, page=2, ...).

    Parameters
    ----------
    page_size : int
        Number of records per page.
    page_param : str
        Query parameter name for the page number.
    size_param : str
        Query parameter name for the page size.
    records_key : str
        JSON key containing the list of records.
    """

    def __init__(
        self,
        page_size: int = 1000,
        page_param: str = "page",
        size_param: str = "size",
        records_key: str = "data",
    ) -> None:
        self.page_size = page_size
        self.page_param = page_param
        self.size_param = size_param
        self.records_key = records_key

    def next_request(
        self,
        response: httpx.Response,
        current_params: dict[str, Any],
    ) -> Optional[dict[str, Any]]:
        records = self.extract_records(response)
        if len(records) < self.page_size:
            return None  # partial page — last page reached
        next_page = current_params.get(self.page_param, 1) + 1
        return {**current_params, self.page_param: next_page}

    def extract_records(self, response: httpx.Response) -> list[dict]:
        return response.json().get(self.records_key, [])

class RestApiClientConfig(BaseModel):
    """Top-level configuration for a REST API client."""

    base_url: str = Field(..., description="Base URL of the REST API")
    auth: AuthConfig = Field(default_factory=AuthConfig, description="Authentication config")
    timeout: float = Field(30.0, description="Request timeout in seconds")
    max_retries: int = Field(3, description="Max retries on transient errors")
    retry_backoff: float = Field(1.5, description="Exponential backoff multiplier")
    default_headers: dict[str, str] = Field(default_factory=dict)

    class Config:
        frozen = True