Add "Create API Token" and "Delete API Token" calls

src/axiom_py/client.py

@@ -5,6 +5,7 @@
 import gzip
 import ujson
 import os
+
 from enum import Enum
 from humps import decamelize
 from typing import Optional, List, Dict, Callable
@@ -24,6 +25,7 @@
 from .users import UsersClient
 from .version import __version__
 from .util import from_dict, handle_json_serialization, is_personal_token
+from .tokens import TokenAttributes, Token
 
 
 AXIOM_URL = "https://api.axiom.co"
@@ -108,6 +110,8 @@ class AplOptions:
     # IncludeCursor will return the Cursor as part of the query result, if set
     # to true.
     includeCursor: bool = field(default=False)
+    # The query limit.
+    limit: Optional[int] = field(default=None)
 
 
 class AxiomError(Exception):
@@ -308,8 +312,24 @@ def query(
         result = from_dict(QueryResult, res.json())
         query_id = res.headers.get("X-Axiom-History-Query-Id")
         result.savedQueryID = query_id
+
         return result
 
+    def create_api_token(self, opts: TokenAttributes) -> Token:
+        """Creates a new API token with permissions specified in a TokenAttributes object."""
+        res = self.session.post(
+            "/v2/tokens",
+            data=ujson.dumps(asdict(opts), default=handle_json_serialization),
+        )
+
+        # Return the new token and ID.
+        response = res.json()
+        return Token(id=response["id"], token=response["token"])
+
+    def delete_api_token(self, token_id: str) -> None:
+        """Delete an API token using its ID string."""
+        self.session.delete(f"/v2/tokens/{token_id}")
+
     def _prepare_query_options(self, opts: QueryOptions) -> Dict[str, object]:
         """returns the query options as a Dict, handles any renaming for key fields."""
         if opts is None:
@@ -350,11 +370,13 @@ def _prepare_apl_options(
         self, opts: Optional[AplOptions]
     ) -> Dict[str, object]:
         """Prepare the apl query options for the request."""
-        params = {"format": AplResultFormat.Legacy.value}
+        params: Dict[str, object] = {"format": AplResultFormat.Legacy.value}
 
         if opts is not None:
             if opts.format:
                 params["format"] = opts.format.value
+            if opts.limit is not None:
+                params["request"] = {"limit": opts.limit}
 
         return params
 

src/axiom_py/tokens.py

@@ -0,0 +1,121 @@
+from dataclasses import dataclass, field
+from typing import Literal, Optional
+
+
+@dataclass
+class TokenDatasetCapabilities:
+    # pylint: disable=unsubscriptable-object
+    """
+    TokenDatasetCapabilities describes the dataset-level permissions
+    which a token can be assigned.
+    Each token can have multiple dataset-level capability objects;
+    one per dataset.
+    """
+
+    # Ability to ingest data. Optional.
+    ingest: Optional[list[Literal["create"]]] = field(default=None)
+    # Ability to query data. Optional.
+    query: Optional[list[Literal["read"]]] = field(default=None)
+    # Ability to use starred queries. Optional.
+    starredQueries: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+    # Ability to use virtual fields. Optional.
+    virtualFields: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+
+
+@dataclass
+class TokenOrganizationCapabilities:
+    # pylint: disable=unsubscriptable-object
+    """
+    TokenOrganizationCapabilities describes the org-level permissions
+    which a token can be assigned.
+    """
+
+    # Ability to use annotations. Optional.
+    annotations: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+    # Ability to use api tokens. Optional.
+    apiTokens: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+    # Ability to access billing. Optional.
+    billing: Optional[list[Literal["read", "update"]]] = field(default=None)
+    # Ability to use dashboards. Optional.
+    dashboards: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+    # Ability to use datasets. Optional.
+    datasets: Optional[list[Literal["create", "read", "update", "delete"]]] = (
+        field(default=None)
+    )
+    # Ability to use endpoints. Optional.
+    endpoints: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+    # Ability to use flows. Optional.
+    flows: Optional[list[Literal["create", "read", "update", "delete"]]] = (
+        field(default=None)
+    )
+    # Ability to use integrations. Optional.
+    integrations: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+    # Ability to use monitors. Optional.
+    monitors: Optional[list[Literal["create", "read", "update", "delete"]]] = (
+        field(default=None)
+    )
+    # Ability to use notifiers. Optional.
+    notifiers: Optional[
+        list[Literal["create", "read", "update", "delete"]]
+    ] = field(default=None)
+    # Ability to use role-based access controls. Optional.
+    rbac: Optional[list[Literal["create", "read", "update", "delete"]]] = (
+        field(default=None)
+    )
+    # Ability to use shared access keys. Optional.
+    sharedAccessKeys: Optional[list[Literal["read", "update"]]] = field(
+        default=None
+    )
+    # Ability to use users. Optional.
+    users: Optional[list[Literal["create", "read", "update", "delete"]]] = (
+        field(default=None)
+    )
+
+
+@dataclass
+class TokenAttributes:
+    # pylint: disable=unsubscriptable-object
+    """
+    TokenAttributes describes the set of input parameters that the
+    POST /tokens API accepts.
+    """
+
+    # Name for the token. Required.
+    name: str
+    # The token's dataset-level capabilities. Keyed on dataset name. Optional.
+    datasetCapabilities: Optional[dict[str, TokenDatasetCapabilities]] = field(
+        default=None
+    )
+    # Description for the API token. Optional.
+    description: Optional[str] = field(default=None)
+    # Expiration date for the API token. Optional.
+    expiresAt: Optional[str] = field(default=None)
+    # The token's organization-level capabilities. Optional.
+    orgCapabilities: Optional[TokenOrganizationCapabilities] = field(
+        default=None
+    )
+
+
+@dataclass
+class Token:
+    """
+    Token contains the response from a call to POST /tokens.
+    It includes the API token itself, and an ID which can be used to reference it later.
+    """
+
+    id: str
+    token: str

tests/test_client.py

@@ -5,11 +5,14 @@
 import unittest
 from unittest.mock import patch
 import gzip
+import uuid
+
 import ujson
 import rfc3339
 import responses
 from logging import getLogger
 from datetime import datetime, timedelta
+
 from .helpers import get_random_name
 from axiom_py import (
     AxiomError,
@@ -33,6 +36,10 @@
     Aggregation,
     AggregationOperation,
 )
+from axiom_py.tokens import (
+    TokenAttributes,
+    TokenOrganizationCapabilities,
+)
 
 
 class TestClient(unittest.TestCase):
@@ -263,6 +270,20 @@ def test_step005_complex_query(self):
             agg = res.buckets.totals[0].aggregations[0]
             self.assertEqual("event_count", agg.op)
 
+    def test_api_tokens(self):
+        """Test creating and deleting an API token"""
+        token_attrs = TokenAttributes(
+            name=f"PytestToken-{uuid.uuid4()}",
+            orgCapabilities=TokenOrganizationCapabilities(apiTokens=["read"]),
+        )
+        token_values = self.client.create_api_token(token_attrs)
+
+        assert token_values.id
+        assert token_values.token
+
+        # (An exception will be raised if the delete call is not successful.)
+        self.client.delete_api_token(token_values.id)
+
     @patch("sys.exit")
     def test_client_shutdown_atexit(self, mock_exit):
         """Test client shutdown atexit"""