Merge pull request #131 from equinor/use-new-aggregation-endpoint

add aggregation client

Author: equinor-ruaj

src/sumo/wrapper/__init__.py

@@ -1,4 +1,5 @@
 from ._call_sumo_api import CallSumoApi
+from ._sumo_aggregation_client import SumoAggregationClient
 from .sumo_client import SumoClient
 
-__all__ = ["CallSumoApi", "SumoClient"]
+__all__ = ["CallSumoApi", "SumoClient", "SumoAggregationClient"]

src/sumo/wrapper/_sumo_aggregation_client.py

@@ -0,0 +1,81 @@
+import logging
+
+import requests
+
+from ._new_auth import NewAuth
+from ._request_error import raise_request_error_exception
+from .config import AGG_APP_REGISTRATION, TENANT_ID
+
+logger = logging.getLogger("sumo.wrapper")
+
+
+class SumoAggregationClient:
+    def __init__(
+        self,
+        env: str,
+        interactive: bool = False,
+        verbosity: str = "CRITICAL",
+    ):
+        """Initialize a new Sumo Aggregation object
+        Args:
+            env: Sumo environment
+            token: Access token or refresh token.
+            interactive: Enable interactive authentication (in browser).
+                If not enabled, code grant flow will be used.
+            verbosity: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+        """
+
+        logger.setLevel(verbosity)
+
+        if env not in AGG_APP_REGISTRATION:
+            raise ValueError(f"Invalid environment: {env}")
+
+        self.auth = NewAuth(
+            client_id=AGG_APP_REGISTRATION[env]["CLIENT_ID"],
+            resource_id=AGG_APP_REGISTRATION[env]["RESOURCE_ID"],
+            tenant_id=TENANT_ID,
+            interactive=interactive,
+            verbosity=verbosity,
+        )
+
+        if env == "localhost":
+            self.base_url = (
+                "https://main-sumo-surface-aggregation-service-preview"
+                + ".radix.equinor.com"
+            )
+        else:
+            self.base_url = (
+                f"https://main-sumo-surface-aggregation-service-{env}"
+                + ".radix.equinor.com"
+            )
+
+    def get_aggregate(self, json: dict):
+        """
+        Performs a POST-request to Sumo Aggregation API /fastaggregation.
+
+        Takes json with objects to aggregate and aggregation operation
+            as payload
+        Args:
+            json: Json payload
+        Returns:
+            Sumo aggregate response object
+        """
+        token = self.auth.get_token()
+
+        headers = {
+            "Content-Type": "application/json",
+            "authorization": f"Bearer {token}",
+            "Content-Length": str(len(json)),
+        }
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/fastaggregation", json=json, headers=headers
+            )
+        except requests.exceptions.ProxyError as err:
+            raise_request_error_exception(503, err)
+
+        if not response.ok:
+            raise_request_error_exception(response.status_code, response.text)
+
+        return response

src/sumo/wrapper/config.py

@@ -24,3 +24,26 @@
 TENANT_ID = "3aa4a235-b6e2-48d5-9195-7fcf05b459b0"
 
 AUTHORITY_HOST_URI = "https://login.microsoftonline.com"
+
+AGG_APP_REGISTRATION = {
+    "prod": {
+        "CLIENT_ID": "6311fb13-81e5-4393-ab8b-13b369fab92d",
+        "RESOURCE_ID": "784f0baa-013f-4fbf-b28a-76d72b5e8d5c",
+    },
+    "test": {
+        "CLIENT_ID": "6311fb13-81e5-4393-ab8b-13b369fab92d",
+        "RESOURCE_ID": "784f0baa-013f-4fbf-b28a-76d72b5e8d5c",
+    },
+    "dev": {
+        "CLIENT_ID": "6311fb13-81e5-4393-ab8b-13b369fab92d",
+        "RESOURCE_ID": "4d631410-6983-4fdb-b20c-2e2ba0c0d506",
+    },
+    "preview": {
+        "CLIENT_ID": "6311fb13-81e5-4393-ab8b-13b369fab92d",
+        "RESOURCE_ID": "4d631410-6983-4fdb-b20c-2e2ba0c0d506",
+    },
+    "localhost": {
+        "CLIENT_ID": "6311fb13-81e5-4393-ab8b-13b369fab92d",
+        "RESOURCE_ID": "4d631410-6983-4fdb-b20c-2e2ba0c0d506",
+    },
+}

src/sumo/wrapper/sumo_client.py

@@ -7,6 +7,7 @@
 from ._new_auth import NewAuth
 from ._request_error import raise_request_error_exception
 from ._blob_client import BlobClient
+from ._sumo_aggregation_client import SumoAggregationClient
 
 logger = logging.getLogger("sumo.wrapper")
 
@@ -70,6 +71,10 @@ def __init__(
         else:
             self.base_url = f"https://main-sumo-{env}.radix.equinor.com/api/v1"
 
+        self.agg_client = SumoAggregationClient(
+            env=env, interactive=interactive, verbosity=verbosity
+        )
+
     def authenticate(self) -> str:
         """Authenticate to Sumo.
 
@@ -362,3 +367,14 @@ def delete(self, path: str) -> dict:
             raise_request_error_exception(response.status_code, response.text)
 
         return response.json()
+
+    def get_aggregate(self, json: dict):
+        """Gets a surface aggregation using SumoAggregationClient:get_aggregate
+
+        Args:
+            json: Json payload
+
+        Returns:
+            Sumo aggregate response object
+        """
+        return self.agg_client.get_aggregate(json)