Merge pull request #112 from axiomhq/arne/axm-3951-add-annotations-to…

…-axiom-py

feat(annotations): Add AnnotationsClient

axiom/__init__.py

@@ -6,3 +6,4 @@
 
 from .client import *
 from .datasets import *
+from .annotations import *

axiom/annotations.py

@@ -0,0 +1,113 @@
+"""This package provides annotation models and methods as well as an AnnotationsClient"""
+
+import ujson
+from logging import Logger
+from requests import Session
+from typing import List, Dict, Optional
+from dataclasses import dataclass, asdict, field
+from datetime import datetime, timedelta
+from urllib.parse import urlencode
+from .util import Util
+
+
+@dataclass
+class Annotation:
+    """Represents an Axiom annotation"""
+
+    id: str = field(init=False)
+    datasets: List[str]
+    time: datetime
+    endTime: Optional[datetime]
+    title: Optional[str]
+    description: Optional[str]
+    url: Optional[str]
+    type: str
+
+
+@dataclass
+class AnnotationCreateRequest:
+    """Request used to create an annotation"""
+
+    datasets: List[str]
+    time: Optional[datetime]
+    endTime: Optional[datetime]
+    title: Optional[str]
+    description: Optional[str]
+    url: Optional[str]
+    type: str
+
+
+@dataclass
+class AnnotationUpdateRequest:
+    """Request used to update an annotation"""
+
+    datasets: Optional[List[str]]
+    time: Optional[datetime]
+    endTime: Optional[datetime]
+    title: Optional[str]
+    description: Optional[str]
+    url: Optional[str]
+    type: Optional[str]
+
+
+class AnnotationsClient:  # pylint: disable=R0903
+    """AnnotationsClient has methods to manipulate annotations."""
+
+    session: Session
+
+    def __init__(self, session: Session, logger: Logger):
+        self.session = session
+        self.logger = logger
+
+    def get(self, id: str) -> Annotation:
+        """Get a annotation by id."""
+        path = "/v2/annotations/%s" % id
+        res = self.session.get(path)
+        decoded_response = res.json()
+        return Util.from_dict(Annotation, decoded_response)
+
+    def create(self, req: AnnotationCreateRequest) -> Annotation:
+        """Create an annotation with the given properties."""
+        path = "/v2/annotations"
+        res = self.session.post(path, data=ujson.dumps(asdict(req)))
+        annotation = Util.from_dict(Annotation, res.json())
+        self.logger.info(f"created new annotation: {annotation.id}")
+        return annotation
+
+    def list(
+        self,
+        datasets: List[str] = [],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+    ) -> List[Annotation]:
+        """List all annotations."""
+        query_params = {}
+        if len(datasets) > 0:
+            query_params["datasets"] = ",".join(datasets)
+        if start != None:
+            query_params["start"] = start.isoformat()
+        if end != None:
+            query_params["end"] = end.isoformat()
+        path = f"/v2/annotations?{urlencode(query_params, doseq=True)}"
+
+        res = self.session.get(path)
+
+        annotations = []
+        for record in res.json():
+            ds = Util.from_dict(Annotation, record)
+            annotations.append(ds)
+
+        return annotations
+
+    def update(self, id: str, req: AnnotationUpdateRequest) -> Annotation:
+        """Update an annotation with the given properties."""
+        path = "/v2/annotations/%s" % id
+        res = self.session.put(path, data=ujson.dumps(asdict(req)))
+        annotation = Util.from_dict(Annotation, res.json())
+        self.logger.info(f"updated annotation({annotation.id})")
+        return annotation
+
+    def delete(self, id: str):
+        """Deletes an annotation with the given id."""
+        path = "/v2/annotations/%s" % id
+        self.session.delete(path)

axiom/client.py

@@ -1,4 +1,5 @@
 """Client provides an easy-to use client library to connect to Axiom."""
+
 import ndjson
 import dacite
 import gzip
@@ -16,6 +17,7 @@
 from requests.adapters import HTTPAdapter, Retry
 from .datasets import DatasetsClient
 from .query import QueryLegacy, QueryResult, QueryOptions, QueryLegacyResult, QueryKind
+from .annotations import AnnotationsClient
 from .users import UsersClient
 from .__init__ import __version__
 
@@ -129,6 +131,7 @@ class Client:  # pylint: disable=R0903
 
     datasets: DatasetsClient
     users: UsersClient
+    annotations: AnnotationsClient
 
     def __init__(
         self,
@@ -143,15 +146,14 @@ def __init__(
             org_id = os.getenv("AXIOM_ORG_ID")
         if url_base is None:
             url_base = AXIOM_URL
-        # Append /v1 to the url_base
-        url_base = url_base.rstrip("/") + "/v1/"
 
         self.logger = getLogger()
-        self.session = BaseUrlSession(url_base)
         # set exponential retries
         retries = Retry(
             total=3, backoff_factor=2, status_forcelist=[500, 502, 503, 504]
         )
+
+        self.session = BaseUrlSession(url_base.rstrip("/"))
         self.session.mount("http://", HTTPAdapter(max_retries=retries))
         self.session.mount("https://", HTTPAdapter(max_retries=retries))
         # hook on responses, raise error when response is not successfull
@@ -175,6 +177,7 @@ def __init__(
 
         self.datasets = DatasetsClient(self.session, self.logger)
         self.users = UsersClient(self.session)
+        self.annotations = AnnotationsClient(self.session, self.logger)
 
     def ingest(
         self,
@@ -185,7 +188,7 @@ def ingest(
         opts: Optional[IngestOptions] = None,
     ) -> IngestStatus:
         """Ingest the events into the named dataset and returns the status."""
-        path = "datasets/%s/ingest" % dataset
+        path = "/v1/datasets/%s/ingest" % dataset
 
         # check if passed content type and encoding are correct
         if not contentType:
@@ -231,7 +234,7 @@ def query_legacy(
                 % (opts.saveAsKind, QueryKind.ANALYTICS, QueryKind.STREAM)
             )
 
-        path = "datasets/%s/query" % id
+        path = "/v1/datasets/%s/query" % id
         payload = ujson.dumps(asdict(query), default=Util.handle_json_serialization)
         self.logger.debug("sending query %s" % payload)
         params = self._prepare_query_options(opts)
@@ -249,7 +252,7 @@ def apl_query(self, apl: str, opts: Optional[AplOptions] = None) -> QueryResult:
 
     def query(self, apl: str, opts: Optional[AplOptions] = None) -> QueryResult:
         """Executes the given apl query on the dataset identified by its id."""
-        path = "datasets/_apl"
+        path = "/v1/datasets/_apl"
         payload = ujson.dumps(
             self._prepare_apl_payload(apl, opts),
             default=Util.handle_json_serialization,

axiom/datasets.py

@@ -1,4 +1,5 @@
 """This package provides dataset models and methods as well as a DatasetClient"""
+
 import ujson
 from logging import Logger
 from requests import Session
@@ -82,22 +83,22 @@ def __init__(self, session: Session, logger: Logger):
 
     def get(self, id: str) -> Dataset:
         """Get a dataset by id."""
-        path = "datasets/%s" % id
+        path = "/v1/datasets/%s" % id
         res = self.session.get(path)
         decoded_response = res.json()
         return Util.from_dict(Dataset, decoded_response)
 
     def create(self, req: DatasetCreateRequest) -> Dataset:
         """Create a dataset with the given properties."""
-        path = "datasets"
+        path = "/v1/datasets"
         res = self.session.post(path, data=ujson.dumps(asdict(req)))
         ds = Util.from_dict(Dataset, res.json())
         self.logger.info(f"created new dataset: {ds.name}")
         return ds
 
     def get_list(self) -> List[Dataset]:
         """List all available datasets."""
-        path = "datasets"
+        path = "/v1/datasets"
         res = self.session.get(path)
 
         datasets = []
@@ -109,15 +110,15 @@ def get_list(self) -> List[Dataset]:
 
     def update(self, id: str, req: DatasetUpdateRequest) -> Dataset:
         """Update a dataset with the given properties."""
-        path = "datasets/%s" % id
+        path = "/v1/datasets/%s" % id
         res = self.session.put(path, data=ujson.dumps(asdict(req)))
         ds = Util.from_dict(Dataset, res.json())
         self.logger.info(f"updated dataset({ds.name}) with new desc: {ds.description}")
         return ds
 
     def delete(self, id: str):
         """Deletes a dataset with the given id."""
-        path = "datasets/%s" % id
+        path = "/v1/datasets/%s" % id
         self.session.delete(path)
 
     def trim(self, id: str, maxDuration: timedelta):
@@ -126,14 +127,14 @@ def trim(self, id: str, maxDuration: timedelta):
         given will mark the oldest timestamp an event can have. Older ones will be
         deleted from the dataset.
         """
-        path = "datasets/%s/trim" % id
+        path = "/v1/datasets/%s/trim" % id
         # prepare request payload and format masDuration to append time unit at the end, e.g `1s`
         req = TrimRequest(f"{maxDuration.seconds}s")
         self.session.post(path, data=ujson.dumps(asdict(req)))
 
     def info(self, id: str) -> DatasetInfo:
         """Returns the info about a dataset."""
-        path = "datasets/%s/info" % id
+        path = "/v1/datasets/%s/info" % id
         res = self.session.get(path)
         decoded_response = res.json()
         return Util.from_dict(DatasetInfo, decoded_response)

axiom/logging.py

@@ -1,4 +1,5 @@
 """Logging contains the AxiomHandler and related methods to do with logging."""
+
 import time
 import atexit
 

axiom/users.py

@@ -21,6 +21,6 @@ def __init__(self, session: Session):
 
     def current(self) -> User:
         """Get the current authenticated user."""
-        res = self.session.get("user")
+        res = self.session.get("/v1/user")
         user = Util.from_dict(User, res.json())
         return user

tests/helpers.py

@@ -1,4 +1,5 @@
 """This module contains helper functions for tests."""
+
 import random
 from datetime import datetime
 

tests/test_annotations.py

@@ -0,0 +1,87 @@
+"""This module contains the tests for the AnnotationsClient."""
+
+import os
+
+import unittest
+from typing import List, Dict, Any, Optional
+from logging import getLogger
+from requests.exceptions import HTTPError
+from datetime import timedelta
+from .helpers import get_random_name
+from axiom import (
+    Client,
+    DatasetCreateRequest,
+    AnnotationCreateRequest,
+    AnnotationUpdateRequest,
+)
+
+
+class TestAnnotations(unittest.TestCase):
+    client: Client
+    dataset_name: str
+
+    @classmethod
+    def setUpClass(cls):
+        cls.logger = getLogger()
+
+        cls.client = Client(
+            os.getenv("AXIOM_TOKEN"),
+            os.getenv("AXIOM_ORG_ID"),
+            os.getenv("AXIOM_URL"),
+        )
+
+        # create dataset
+        cls.dataset_name = get_random_name()
+        req = DatasetCreateRequest(
+            name=cls.dataset_name,
+            description="test_annotations.py (dataset_name)",
+        )
+        res = cls.client.datasets.create(req)
+
+    def test_happy_path_crud(self):
+        """Test the happy path of creating, reading, updating, and deleting an annotation."""
+        # Create annotation
+        req = AnnotationCreateRequest(
+            datasets=[self.dataset_name],
+            type="test",
+            time=None,
+            endTime=None,
+            title=None,
+            description=None,
+            url=None,
+        )
+        created_annotation = self.client.annotations.create(req)
+        self.logger.debug(created_annotation)
+
+        # Get annotation
+        annotation = self.client.annotations.get(created_annotation.id)
+        self.logger.debug(annotation)
+        assert annotation.id == created_annotation.id
+
+        # List annotations
+        annotations = self.client.annotations.list(datasets=[self.dataset_name])
+        self.logger.debug(annotations)
+        assert len(annotations) == 1
+
+        # Update
+        newTitle = "Update title"
+        updateReq = AnnotationUpdateRequest(
+            datasets=None,
+            type=None,
+            time=None,
+            endTime=None,
+            title=newTitle,
+            description=None,
+            url=None,
+        )
+        updated_annotation = self.client.annotations.update(annotation.id, updateReq)
+        self.logger.debug(updated_annotation)
+        assert updated_annotation.title == newTitle
+
+        # Delete
+        self.client.annotations.delete(annotation.id)
+
+    @classmethod
+    def tearDownClass(cls):
+        """Delete datasets"""
+        cls.client.datasets.delete(cls.dataset_name)

tests/test_client.py

@@ -1,4 +1,5 @@
 """This module contains the tests for the axiom client."""
+
 import os
 import unittest
 import gzip

tests/test_datasets.py

@@ -1,4 +1,5 @@
 """This module contains the tests for the DatasetsClient."""
+
 import os
 
 import unittest

tests/test_logger.py

@@ -1,4 +1,5 @@
 """This module contains test for the logging Handler."""
+
 import os
 import logging
 import unittest