refactor wip

ecmwf · Jul 1, 2024 · 88a8fd3 · 88a8fd3
1 parent 33abecd
commit 88a8fd3
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 87 deletions.
diff --git a/src/anemoi/registry/commands/base.py b/src/anemoi/registry/commands/base.py
@@ -14,6 +14,7 @@
 
 import logging
 import os
+import sys
 
 from ..entry import CatalogueEntryNotFound
 from . import Command
@@ -59,6 +60,8 @@ def run(self, args):
             LOG.info(f"Processing {self.kind} with identifier '{name_or_path}'")
             self.run_from_identifier(name_or_path, **args)
             return
+        LOG.error(f"Cannot find any {self.kind} from '{name_or_path}'")
+        sys.exit(1)
 
     def parse_location(self, location):
         for x in location:

diff --git a/src/anemoi/registry/entry/__init__.py b/src/anemoi/registry/entry/__init__.py
@@ -8,7 +8,6 @@
 import json
 import logging
 
-import requests
 from anemoi.utils.humanize import json_pretty_dump
 
 from anemoi.registry import config
@@ -63,11 +62,9 @@ def load_from_key(self, key):
     @classmethod
     def _get_record_from_catalogue(cls, key):
         try:
-            return cls.rest.get(f"{cls.collection_api}/{key}")
-        except requests.exceptions.HTTPError as e:
-            if e.response.status_code == 404:
-                raise CatalogueEntryNotFound(f"Could not find any {cls.collection_api} with key={key}")
-            raise
+            return cls.rest.get(f"{cls.collection_api}/{key}", errors={404: CatalogueEntryNotFound})
+        except CatalogueEntryNotFound:
+            raise CatalogueEntryNotFound(f"Could not find any {cls.collection_api} with key={key}")
 
     @property
     def main_key(self):
@@ -79,6 +76,7 @@ def register(self, ignore_existing=True, overwrite=False):
             return self.rest.post(self.collection_api, self.record)
         except AlreadyExists:
             if ignore_existing:
+                LOG.info(f"{self.key} already exists. Ok.")
                 return
             if overwrite is True:
                 LOG.warning(f"{self.key} already exists. Deleting existing one to overwrite it.")
@@ -95,7 +93,7 @@ def patch(self, payload):
     def unregister(self):
         if not config().get("allow_delete"):
             raise ValueError("Unregister not allowed")
-        return self.rest.delete(self.collection_api,self.key,dict(force=True))
+        return self.rest.delete(self.collection_api, self.key, dict(force=True))
 
     def __repr__(self):
         return json.dumps(self.record, indent=2)
diff --git a/src/anemoi/registry/rest.py b/src/anemoi/registry/rest.py
@@ -7,16 +7,11 @@
 
 import datetime
 import logging
-import os
-import socket
-import sys
-from getpass import getuser
 
 import requests
 from requests.exceptions import HTTPError
 
 from anemoi.registry import config
-from anemoi.registry._version import version
 
 LOG = logging.getLogger(__name__)
 # LOG.setLevel(logging.DEBUG)
@@ -53,18 +48,16 @@ def tidy(d):
 
 
 class BaseRest:
-    def get(self, collection, payload=None):
+    def get(self, collection, params=None, errors={}):
         LOG.debug(f"GET {collection}")
-        try:
-            kwargs = dict( headers={"Authorization": f"Bearer {self.token}"})
-            if payload is not None:
-                kwargs["params"] = payload
-            r = requests.get( f"{config().api_url}/{collection}", **kwargs)
-            self.raise_for_status(r)
-            return r.json()
-        except Exception as e:
-            LOG.error(e)
-            raise (e)
+
+        kwargs = dict()
+        if params is not None:
+            kwargs["params"] = params
+
+        r = requests.get(f"{config().api_url}/{collection}", **self.token_dict, **kwargs)
+        self.raise_for_status(r, errors=errors)
+        return r.json()
 
     def post(self, collection, data):
         LOG.debug(f"POST {collection} { {k:'...' for k,v in data.items()} }")
@@ -85,93 +78,60 @@ def trace_info_dict(self):
     def token(self):
         return config().api_token
 
-    def raise_for_status(self, r):
+    @property
+    def token_dict(self):
+        return dict(headers={"Authorization": f"Bearer {self.token}"})
+
+    def raise_for_status(self, r, errors={}):
         try:
             r.raise_for_status()
         except HTTPError as e:
             # add the response text to the exception message
             text = r.text
             text = text[:1000] + "..." if len(text) > 1000 else text
             e.args = (f"{e.args[0]} : {text}",)
-            raise e
+
+            exception_handler = errors.get(e.response.status_code)
+            errcode = e.response.status_code
+            LOG.debug("HTTP error: ", errcode, exception_handler)
+            if exception_handler:
+                raise exception_handler(e)
+            else:
+                raise e
 
 
 class ReadOnlyRest(BaseRest):
     pass
 
 
 class Rest(BaseRest):
-    def raise_for_status(self, r):
-        try:
-            r.raise_for_status()
-        except HTTPError as e:
-            # add the response text to the exception message
-            text = r.text
-            text = text[:1000] + "..." if len(text) > 1000 else text
-            e.args = (f"{e.args[0]} : {text}",)
-            raise e
 
     def post(self, collection, data):
         super().post(collection, data)
-        try:
-            r = requests.post(
-                f"{config().api_url}/{collection}",
-                json=tidy(data),
-                headers={"Authorization": f"Bearer {self.token}"},
-            )
-            self.raise_for_status(r)
-            return r.json()
-        except requests.exceptions.HTTPError as e:
-            if e.response.status_code == 409:
-                raise AlreadyExists(f"{e}Already exists in {collection}")
-            else:
-                LOG.error(f"Error in post to {collection} with data:{data}")
-                LOG.error(e)
-                raise
-        except Exception as e:
-            LOG.error(f"Error in post to {collection} with data:{data}")
-            LOG.error(e)
-            raise
+        r = requests.post(f"{config().api_url}/{collection}", json=tidy(data), **self.token_dict)
+        self.raise_for_status(r, {409: AlreadyExists})
+        return r.json()
 
     def patch(self, collection, data):
         super().patch(collection, data)
-        try:
-            r = requests.patch(
-                f"{config().api_url}/{collection}",
-                json=tidy(data),
-                headers={"Authorization": f"Bearer {self.token}"},
-            )
-            self.raise_for_status(r)
-            return r.json()
-        except Exception as e:
-            LOG.error(e)
-            raise (e)
+        r = requests.patch(f"{config().api_url}/{collection}", json=tidy(data), **self.token_dict)
+        self.raise_for_status(r)
+        return r.json()
 
     def put(self, collection, data):
         super().put(collection, data)
-        try:
-            r = requests.put(
-                f"{config().api_url}/{collection}",
-                json=tidy(data),
-                headers={"Authorization": f"Bearer {self.token}"},
-            )
-            self.raise_for_status(r)
-            return r.json()
-        except Exception as e:
-            LOG.error(e)
-            raise (e)
+        r = requests.put(f"{config().api_url}/{collection}", json=tidy(data), **self.token_dict)
+        self.raise_for_status(r)
+        return r.json()
 
     def delete(self, collection, key, data=None):
         super().delete(collection, key, data)
-        try:
-            kwargs = dict( headers={"Authorization": f"Bearer {self.token}"})
-            if data is not None:
-                kwargs["params"] = data
-
-            assert key
-            r = requests.delete( f"{config().api_url}/{collection}/{key}", **kwargs)
-            self.raise_for_status(r)
-            return r.json()
-        except Exception as e:
-            LOG.error(e)
-            raise (e)
+
+        kwargs = dict()
+        if data is not None:
+            kwargs["params"] = data
+
+        assert key
+        r = requests.delete(f"{config().api_url}/{collection}/{key}", **self.token_dict, **kwargs)
+        self.raise_for_status(r)
+        return r.json()