aiopromql.client — aiopromql 0.1.1 documentation

# Source code for aiopromql.client

    import warnings
    from datetime import datetime
    from typing import Optional, Union

    import httpx

    from .models.prometheus import PrometheusResponseModel

    class PrometheusClientBase:
        """Base Prometheus client with common utilities."""

        def __init__(self, url: str):
            self.base_url = url

        def _parse_response(self, response: dict) -> PrometheusResponseModel:
            """Parse Prometheus JSON response into model."""
            return PrometheusResponseModel(**response)

    class PrometheusSync(PrometheusClientBase):
        """Synchronous Prometheus client using httpx."""

        def __init__(self, url: str, timeout: Optional[float] = 2.0):
            super().__init__(url)
            self.session = httpx.Client(timeout=httpx.Timeout(timeout))

        def query(self, promql: str, raw: bool = False) -> Union[PrometheusResponseModel, dict]:
            """
            Run an instant PromQL query.

            :param promql: The PromQL query string to execute.
            :type promql: str
            :param raw: If True, return raw JSON response as dict; otherwise parse into model.
            :type raw: bool
            :return: Parsed Prometheus response or raw JSON dict.
            """
            response = self.session.get(f"{self.base_url}/api/v1/query", params={"query": promql})
            response.raise_for_status()
            return response.json() if raw else self._parse_response(response.json())

        def query_range(
            self,
            promql: str,
            start: datetime,
            end: datetime,
            step: str = "30s",
            raw: bool = False,
        ) -> Union[PrometheusResponseModel, dict]:
            """
            Run a ranged PromQL query over a time window.

            :param promql: The PromQL query string to execute.
            :type promql: str
            :param start: Start datetime of the query range.
            :type start: datetime.datetime
            :param end: End datetime of the query range.
            :type end: datetime.datetime
            :param step: Query resolution step width (e.g., '30s', '1m').
            :type step: str
            :param raw: If True, return raw JSON response as dict; otherwise parse into model.
            :type raw: bool
            :return: Parsed Prometheus response or raw JSON dict.
            """
            start_ts = start.timestamp()
            end_ts = end.timestamp()
            response = self.session.get(
                f"{self.base_url}/api/v1/query_range",
                params={"query": promql, "start": start_ts, "end": end_ts, "step": step},
            )
            response.raise_for_status()
            return response.json() if raw else self._parse_response(response.json())

        def close(self):
            """Close the sync client session."""
            self.session.close()

        def __del__(self):
            self.close()

    class PrometheusAsync(PrometheusClientBase):
        """Asynchronous Prometheus client using httpx."""

        def __init__(self, url: str, timeout: Optional[float] = 2.0):
            super().__init__(url)
            self.client = httpx.AsyncClient(base_url=url, timeout=httpx.Timeout(timeout))

        async def query(self, promql: str, raw: bool = False) -> Union[PrometheusResponseModel, dict]:
            """
            Run an instant PromQL query.

            :param promql: The PromQL query string to execute.
            :type promql: str
            :param raw: If True, return raw JSON response as dict; otherwise parse into model.
            :type raw: bool
            :return: Parsed Prometheus response or raw JSON dict.
            """
            response = await self.client.get("/api/v1/query", params={"query": promql})
            response.raise_for_status()
            return response.json() if raw else self._parse_response(response.json())

        async def query_range(
            self,
            promql: str,
            start: datetime,
            end: datetime,
            step: str = "30s",
            raw: bool = False,
        ) -> Union[PrometheusResponseModel, dict]:
            """
            Run a ranged PromQL query over a time window.

            :param promql: The PromQL query string to execute.
            :type promql: str
            :param start: Start datetime of the query range.
            :type start: datetime.datetime
            :param end: End datetime of the query range.
            :type end: datetime.datetime
            :param step: Query resolution step width (e.g., '30s', '1m').
            :type step: str
            :param raw: If True, return raw JSON response as dict; otherwise parse into model.
            :type raw: bool
            :return: Parsed Prometheus response or raw JSON dict.
            """
            start_ts = start.timestamp()
            end_ts = end.timestamp()
            response = await self.client.get(
                "/api/v1/query_range",
                params={"query": promql, "start": start_ts, "end": end_ts, "step": step},
            )
            response.raise_for_status()
            return response.json() if raw else self._parse_response(response.json())

        async def aclose(self):
            """Close the async client session."""
            await self.client.aclose()

        async def __aenter__(self):
            return self

        async def __aexit__(self, exc_type, exc, tb):
            await self.aclose()

        def __del__(self):
            if not self.client.is_closed:
                warnings.warn("PrometheusAsync was not closed. Use 'async with' or call 'await .aclose()'")