diff --git a/axiom/__init__.py b/axiom/__init__.py index fafa474..b358f60 100644 --- a/axiom/__init__.py +++ b/axiom/__init__.py @@ -6,3 +6,4 @@ from .client import * from .datasets import * +from .annotations import * diff --git a/axiom/annotations.py b/axiom/annotations.py new file mode 100644 index 0000000..365d3a7 --- /dev/null +++ b/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) diff --git a/axiom/client.py b/axiom/client.py index e8d8cca..6f94861 100644 --- a/axiom/client.py +++ b/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, diff --git a/axiom/datasets.py b/axiom/datasets.py index cb52577..61ba2e5 100644 --- a/axiom/datasets.py +++ b/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,14 +83,14 @@ 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}") @@ -97,7 +98,7 @@ def create(self, req: DatasetCreateRequest) -> Dataset: def get_list(self) -> List[Dataset]: """List all available datasets.""" - path = "datasets" + path = "/v1/datasets" res = self.session.get(path) datasets = [] @@ -109,7 +110,7 @@ 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}") @@ -117,7 +118,7 @@ def update(self, id: str, req: DatasetUpdateRequest) -> Dataset: 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) diff --git a/axiom/logging.py b/axiom/logging.py index ab41a4c..6916336 100644 --- a/axiom/logging.py +++ b/axiom/logging.py @@ -1,4 +1,5 @@ """Logging contains the AxiomHandler and related methods to do with logging.""" + import time import atexit diff --git a/axiom/users.py b/axiom/users.py index a5277f9..1c5d968 100644 --- a/axiom/users.py +++ b/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 diff --git a/tests/helpers.py b/tests/helpers.py index a88a541..c9bd681 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,4 +1,5 @@ """This module contains helper functions for tests.""" + import random from datetime import datetime diff --git a/tests/test_annotations.py b/tests/test_annotations.py new file mode 100644 index 0000000..ed23490 --- /dev/null +++ b/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) diff --git a/tests/test_client.py b/tests/test_client.py index a702024..385612c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,4 +1,5 @@ """This module contains the tests for the axiom client.""" + import os import unittest import gzip diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 24ae725..8375f61 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -1,4 +1,5 @@ """This module contains the tests for the DatasetsClient.""" + import os import unittest diff --git a/tests/test_logger.py b/tests/test_logger.py index 68eed92..7559ad4 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -1,4 +1,5 @@ """This module contains test for the logging Handler.""" + import os import logging import unittest