Add caching for Management API

Author: kesmit13

singlestoredb/management/utils.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python
 """SingleStoreDB Cluster Management."""
 import datetime
+import functools
 import os
 import re
 import sys
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -31,37 +33,44 @@
     PathLikeABC = os.PathLike[str]
 
 
-def get_token() -> Optional[str]:
-    """Return the token for the Management API."""
-    # See if an API key is configured
-    tok = get_option('management.token')
-    if tok:
-        return tok
+class TTLProperty(object):
+    """Property with time limit."""
 
-    # See if the connection URL contains a JWT
-    url = get_option('host')
-    if not url:
-        return None
+    def __init__(self, fget: Callable[[Any], Any], ttl: datetime.timedelta):
+        self.fget = fget
+        self.ttl = ttl
+        self._last_executed = datetime.datetime(2000, 1, 1)
+        self._last_result = None
+        self.__doc__ = fget.__doc__
+        self._name = ''
 
-    urlp = urlparse(url, scheme='singlestoredb', allow_fragments=True)
-    if urlp.password:
-        try:
-            jwt.decode(urlp.password, options={'verify_signature': False})
-            return urlp.password
-        except jwt.DecodeError:
-            pass
+    def reset(self) -> None:
+        self._last_executed = datetime.datetime(2000, 1, 1)
+        self._last_result = None
 
-    # Didn't find a key anywhere
-    return None
+    def __set_name__(self, owner: Any, name: str) -> None:
+        self._name = name
 
+    def __get__(self, obj: Any, objtype: Any = None) -> Any:
+        if obj is None:
+            return self
 
-def get_organization() -> Optional[str]:
-    """Return the organization for the current token or environment."""
-    org = os.environ.get('SINGLESTOREDB_ORGANIZATION')
-    if org:
-        return org
+        if self._last_result is not None \
+                and (datetime.datetime.now() - self._last_executed) < self.ttl:
+            return self._last_result
 
-    return None
+        self._last_result = self.fget(obj)
+        self._last_executed = datetime.datetime.now()
+
+        return self._last_result
+
+
+def ttl_property(ttl: datetime.timedelta) -> Callable[[Any], Any]:
+    """Property with a time-to-live."""
+    def wrapper(func: Callable[[Any], Any]) -> Any:
+        out = TTLProperty(func, ttl=ttl)
+        return functools.wraps(func)(out)
+    return wrapper
 
 
 class NamedList(List[T]):
@@ -107,6 +116,39 @@ def get(self, name_or_id: str, *default: Any) -> Any:
             raise
 
 
+def get_token() -> Optional[str]:
+    """Return the token for the Management API."""
+    # See if an API key is configured
+    tok = get_option('management.token')
+    if tok:
+        return tok
+
+    # See if the connection URL contains a JWT
+    url = get_option('host')
+    if not url:
+        return None
+
+    urlp = urlparse(url, scheme='singlestoredb', allow_fragments=True)
+    if urlp.password:
+        try:
+            jwt.decode(urlp.password, options={'verify_signature': False})
+            return urlp.password
+        except jwt.DecodeError:
+            pass
+
+    # Didn't find a key anywhere
+    return None
+
+
+def get_organization() -> Optional[str]:
+    """Return the organization for the current token or environment."""
+    org = os.environ.get('SINGLESTOREDB_ORGANIZATION')
+    if org:
+        return org
+
+    return None
+
+
 def enable_http_tracing() -> None:
     """Enable tracing of HTTP requests."""
     import logging

singlestoredb/management/workspace.py

@@ -28,6 +28,7 @@
 from .utils import PathLike
 from .utils import snake_to_camel
 from .utils import to_datetime
+from .utils import ttl_property
 from .utils import vars_to_str
 
 
@@ -1119,7 +1120,7 @@ def from_dict(
         try:
             region = [x for x in manager.regions if x.id == obj['regionID']][0]
         except IndexError:
-            region = None
+            region = Region(obj.get('regionID', '<unknown>'), '<unknown>', '<unknown>')
         out = cls(
             name=obj['name'],
             id=obj['workspaceGroupID'],
@@ -1377,23 +1378,23 @@ class WorkspaceManager(Manager):
     #: Object type
     obj_type = 'workspace'
 
-    @ property
+    @property
     def workspace_groups(self) -> NamedList[WorkspaceGroup]:
         """Return a list of available workspace groups."""
         res = self._get('workspaceGroups')
         return NamedList([WorkspaceGroup.from_dict(item, self) for item in res.json()])
 
-    @ property
+    @property
     def organizations(self) -> Organizations:
         """Return the organizations."""
         return Organizations(self)
 
-    @ property
+    @property
     def billing(self) -> Billing:
         """Return the current billing information."""
         return Billing(self)
 
-    @ property
+    @ttl_property(datetime.timedelta(hours=1))
     def regions(self) -> NamedList[Region]:
         """Return a list of available regions."""
         res = self._get('regions')