Merge pull request #186 from axiomhq/aisha/support-mpl

feat: Add MPL Query Support (Sync + Async)

Author: aisha331

README.md

@@ -44,6 +44,34 @@ edge_client = axiom_py.Client(
 
 **Note:** Edge endpoints require API tokens (`xaat-`), not personal tokens.
 Edge configuration must be passed explicitly when creating the client.
+
+## Metrics Queries (MPL)
+
+To query OTel metrics using MPL (Metrics Processing Language), configure the
+client with your edge endpoint and use `mpl_query`:
+
+```python
+import axiom_py
+from axiom_py import MplOptions
+from datetime import datetime, timedelta, timezone
+
+client = axiom_py.Client(
+    token="xaat-your-api-token",
+    edge="us-east-1.aws.edge.axiom.co"
+)
+
+end = datetime.now(timezone.utc)
+start = end - timedelta(hours=1)
+
+result = client.mpl_query(
+    "`my-metrics`:`http.server.duration` | align to 5m using avg",
+    opts=MplOptions(start_time=start, end_time=end),
+)
+
+for series in result.series:
+    print(series.metric, series.tags, series.data)
+```
+
 ## Asynchronous Client
 
 The library also provides an async client for use with asyncio:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "axiom-py"
-version = "0.10.0"
+version = "0.12.0"
 description = "Official bindings for the Axiom API"
 readme = "README.md"
 requires-python = ">=3.8"

src/axiom_py/__init__.py

@@ -13,10 +13,13 @@
     ContentEncoding,
     WrongQueryKindException,
     PersonalTokenNotSupportedForEdgeError,
+    EdgeResolutionError,
     AplOptions,
+    MplOptions,
     Client,
     AXIOM_URL,
 )
+from .query import MplResult, MplSeriesItem, MplMetadata
 from .datasets import (
     Dataset,
     TrimRequest,
@@ -40,35 +43,40 @@
 from .logging_async import AsyncAxiomHandler
 from .structlog_async import AsyncAxiomProcessor
 
-_all_ = [
-    AxiomError,
-    IngestFailure,
-    IngestStatus,
-    IngestOptions,
-    AplResultFormat,
-    ContentType,
-    ContentEncoding,
-    WrongQueryKindException,
-    PersonalTokenNotSupportedForEdgeError,
-    AplOptions,
-    AXIOM_URL,
-    Dataset,
-    TrimRequest,
-    Annotation,
-    AnnotationCreateRequest,
-    AnnotationUpdateRequest,
+__all__ = [
+    "AxiomError",
+    "IngestFailure",
+    "IngestStatus",
+    "IngestOptions",
+    "AplResultFormat",
+    "ContentType",
+    "ContentEncoding",
+    "WrongQueryKindException",
+    "PersonalTokenNotSupportedForEdgeError",
+    "EdgeResolutionError",
+    "AplOptions",
+    "MplOptions",
+    "MplResult",
+    "MplSeriesItem",
+    "MplMetadata",
+    "AXIOM_URL",
+    "Dataset",
+    "TrimRequest",
+    "Annotation",
+    "AnnotationCreateRequest",
+    "AnnotationUpdateRequest",
     # Sync API
-    Client,
-    DatasetsClient,
-    AnnotationsClient,
-    AxiomHandler,
-    AxiomProcessor,
+    "Client",
+    "DatasetsClient",
+    "AnnotationsClient",
+    "AxiomHandler",
+    "AxiomProcessor",
     # Async API
-    AsyncClient,
-    AsyncDatasetsClient,
-    AsyncAnnotationsClient,
-    AsyncTokensClient,
-    AsyncUsersClient,
-    AsyncAxiomHandler,
-    AsyncAxiomProcessor,
+    "AsyncClient",
+    "AsyncDatasetsClient",
+    "AsyncAnnotationsClient",
+    "AsyncTokensClient",
+    "AsyncUsersClient",
+    "AsyncAxiomHandler",
+    "AsyncAxiomProcessor",
 ]

src/axiom_py/client.py

@@ -8,7 +8,7 @@
 
 from enum import Enum
 from humps import decamelize
-from typing import Optional, List, Dict, Callable
+from typing import Dict, List, Optional, Callable
 from dataclasses import dataclass, field, asdict
 from datetime import datetime
 from requests_toolbelt.sessions import BaseUrlSession
@@ -20,14 +20,19 @@
     QueryOptions,
     QueryLegacyResult,
     QueryKind,
+    MplResult,
 )
 from .annotations import AnnotationsClient
 from .users import UsersClient
 from .version import __version__
-from .util import from_dict, handle_json_serialization, is_personal_token
+from .util import (
+    from_dict,
+    format_datetime_rfc3339_utc,
+    handle_json_serialization,
+    is_personal_token,
+)
 from .tokens import TokensClient
 
-
 AXIOM_URL = "https://api.axiom.co"
 
 
@@ -41,6 +46,12 @@ def __init__(self):
         )
 
 
+class EdgeResolutionError(Exception):
+    """Raised when the edge domain for a dataset cannot be resolved."""
+
+    pass
+
+
 @dataclass
 class IngestFailure:
     """The ingestion failure of a single event"""
@@ -124,6 +135,21 @@ class AplOptions:
     limit: Optional[int] = field(default=None)
 
 
+@dataclass
+class MplOptions:
+    """Options for an MPL (Metrics Processing Language) query."""
+
+    # Start time for the query range (required by the API).
+    start_time: datetime
+    # End time for the query range (required by the API).
+    end_time: datetime
+    # Additional parameters forwarded to MetricsDB (e.g. filterbar bindings).
+    params: Optional[Dict[str, str]] = field(default=None)
+    # Optional metrics query options passed through to the API payload.
+    query_options: Optional[Dict[str, str]] = field(default=None)
+    nocache: bool = field(default=False)
+
+
 class AxiomError(Exception):
     """This exception is raised on request errors."""
 
@@ -187,8 +213,8 @@ def __init__(
                 "https://eu-central-1.aws.edge.axiom.co"). When set, ingest
                 requests use `/v1/ingest/{dataset}` and query requests use
                 `/v1/query/_apl`.
-                Must be passed explicitly (not read from environment).
                 Takes precedence over `edge`.
+                Must be passed explicitly (not read from environment).
             edge: Edge domain for ingest/query operations (e.g.,
                 "eu-central-1.aws.edge.axiom.co"). When set, ingest
                 requests use `https://{edge}/v1/ingest/{dataset}` and query
@@ -204,8 +230,8 @@ def __init__(
             url = os.getenv("AXIOM_URL")
         # Note: edge_url and edge are NOT auto-read from environment.
         # Edge configuration must be explicit to avoid accidentally routing
-        # all requests through edge when AXIOM_EDGE_URL or AXIOM_EDGE is set for
-        # edge-specific tests. Create a separate Client with edge_url
+        # all requests through edge when AXIOM_EDGE_URL/AXIOM_EDGE are set
+        # for edge-specific tests. Create a separate Client with edge params
         # for edge operations.
 
         # Normalize empty strings to None for edge config
@@ -269,49 +295,49 @@ def is_edge_configured(self) -> bool:
         """Check if edge is configured."""
         return self._edge_url is not None or self._edge is not None
 
-    def _get_edge_ingest_url(self, dataset: str) -> Optional[str]:
+    def _resolve_edge_url(self, path_suffix: str) -> Optional[str]:
         """
-        Get the full edge ingest URL for a dataset.
-        Returns None if edge is not configured.
+        Resolve an edge URL for the given path suffix.
+
+        If edge_url includes a custom path, it is used as-is.
+        If edge_url is just a host/base URL, the suffix is appended.
+        If only edge domain is configured, https://{edge}/{suffix} is used.
         """
         if self._edge_url is not None:
             url = self._edge_url.rstrip("/")
             parsed = urlparse(url)
             path = parsed.path
 
-            # If path is empty or just "/", append edge ingest format
             if path == "" or path == "/":
-                return f"{parsed.scheme}://{parsed.netloc}/v1/ingest/{dataset}"
+                return f"{parsed.scheme}://{parsed.netloc}/{path_suffix}"
 
-            # edge_url has a custom path, use as-is
             return url
 
         if self._edge is not None:
-            return f"https://{self._edge.rstrip('/')}/v1/ingest/{dataset}"
+            return f"https://{self._edge.rstrip('/')}/{path_suffix}"
 
         return None
 
+    def _get_edge_ingest_url(self, dataset: str) -> Optional[str]:
+        """
+        Get the full edge ingest URL for a dataset.
+        Returns None if edge is not configured.
+        """
+        return self._resolve_edge_url(f"v1/ingest/{dataset}")
+
     def _get_edge_query_url(self) -> Optional[str]:
         """
         Get the full edge query URL.
         Returns None if edge is not configured.
         """
-        if self._edge_url is not None:
-            url = self._edge_url.rstrip("/")
-            parsed = urlparse(url)
-            path = parsed.path
-
-            # If path is empty or just "/", append edge query format
-            if path == "" or path == "/":
-                return f"{parsed.scheme}://{parsed.netloc}/v1/query/_apl"
-
-            # edge_url has a custom path, use as-is
-            return url
+        return self._resolve_edge_url("v1/query/_apl")
 
-        if self._edge is not None:
-            return f"https://{self._edge.rstrip('/')}/v1/query/_apl"
-
-        return None
+    def _get_edge_mpl_url(self) -> Optional[str]:
+        """
+        Get the full edge MPL query URL.
+        Returns None if edge is not configured.
+        """
+        return self._resolve_edge_url("v1/query/_mpl")
 
     def before_shutdown(self, func: Callable):
         self.before_shutdown_funcs.append(func)
@@ -452,6 +478,61 @@ def query(
 
         return result
 
+    def mpl_query(self, mpl: str, opts: MplOptions) -> MplResult:
+        """
+        Execute an MPL (Metrics Processing Language) query.
+
+        MPL queries are edge-only. Configure the client with edge_url or
+        edge to specify the endpoint. Edge endpoints require API tokens
+        (xaat-), not personal tokens.
+
+        Args:
+            mpl: MPL query string (e.g.
+                "`my-dataset`:`http.requests` | align to 5m using avg")
+            opts: Query options including required start_time and end_time.
+
+        Returns:
+            MplResult containing the time-series data.
+        """
+        url = self._get_edge_mpl_url()
+        if url is None:
+            raise EdgeResolutionError(
+                "MPL queries require an edge endpoint; "
+                "set edge_url or edge when creating the Client"
+            )
+        if is_personal_token(self._token):
+            raise PersonalTokenNotSupportedForEdgeError()
+
+        payload = ujson.dumps(
+            self._prepare_mpl_payload(mpl, opts),
+            default=handle_json_serialization,
+        )
+        params = self._prepare_mpl_params(opts)
+        res = self.session.post(url, data=payload, params=params)
+        result = from_dict(MplResult, res.json())
+        result.savedQueryID = res.headers.get("X-Axiom-History-Query-Id")
+        return result
+
+    def _prepare_mpl_payload(
+        self, mpl: str, opts: MplOptions
+    ) -> Dict[str, object]:
+        payload: Dict[str, object] = {
+            "mpl": mpl,
+            "startTime": format_datetime_rfc3339_utc(opts.start_time),
+            "endTime": format_datetime_rfc3339_utc(opts.end_time),
+        }
+        if opts.params:
+            payload["params"] = opts.params
+        if opts.query_options:
+            payload["queryOptions"] = opts.query_options
+        return payload
+
+    def _prepare_mpl_params(self, opts: MplOptions) -> Dict[str, object]:
+        params: Dict[str, object] = {"format": "metrics-v2"}
+        if opts.nocache:
+            params["nocache"] = "true"
+        return params
+
     def _prepare_query_options(self, opts: QueryOptions) -> Dict[str, object]:
         """returns the query options as a Dict, handles any renaming for key
         fields."""