Store users locally to improve performance when querying labs/nodes (#199)

Author: tmikuska

virl2_client/models/lab.py

@@ -60,6 +60,7 @@
 
     from .annotation import AnnotationType
     from .resource_pool import ResourcePool, ResourcePoolManagement
+    from .user import UserManagement
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -93,7 +94,6 @@ class Lab:
         "connector_mappings": "labs/{lab_id}/connector_mappings",
         "resource_pools": "labs/{lab_id}/resource_pools",
         "annotations": "labs/{lab_id}/annotations",
-        "user_list": "users",
     }
 
     def __init__(
@@ -110,6 +110,7 @@ def __init__(
         wait_time: int | float = 5,
         hostname: str | None = None,
         resource_pool_manager: ResourcePoolManagement | None = None,
+        user_management: UserManagement | None = None,
     ) -> None:
         """
         A VIRL2 lab network topology.
@@ -129,6 +130,8 @@ def __init__(
         :param hostname: Hostname/IP and port for pyATS console terminal server.
         :param resource_pool_manager: ResourcePoolManagement object shared
             with parent ClientLibrary.
+        :param user_management: UserManagement object shared with parent
+            ClientLibrary for resolving user IDs to usernames.
         :raises VirlException: If the lab object is created without
             a resource pool manager.
         """
@@ -201,6 +204,12 @@ def __init__(
                 "because it is missing a resource pool manager."
             )
         self._resource_pool_manager = resource_pool_manager
+        if user_management is None:
+            raise VirlException(
+                f"Lab object for lab {title or lab_id} cannot be created "
+                "because it is missing a user management."
+            )
+        self._user_management = user_management
         self._resource_pools = []
         self._stale = False
         self._synced_configs = True
@@ -2458,10 +2467,7 @@ def _set_owner(
         :param user_name: Username.
         """
         if user_id:
-            url = self._url_for("user_list")
-            users = self._session.get(url).json()
-            for user in users:
-                if user["id"] == user_id:
-                    user_name = user["username"]
-                    break
+            resolved = self._user_management.get_username(user_id)
+            if resolved is not None:
+                user_name = resolved
         self._owner = user_name

virl2_client/models/user.py

@@ -20,6 +20,7 @@
 
 from __future__ import annotations
 
+import time
 import warnings
 from typing import TYPE_CHECKING, Any
 
@@ -36,8 +37,17 @@ class UserManagement:
         "user_id": "users/{username}/id",
     }
 
-    def __init__(self, session: httpx.Client) -> None:
+    def __init__(
+        self,
+        session: httpx.Client,
+        auto_sync: bool = True,
+        auto_sync_interval: float = 1.0,
+    ) -> None:
         self._session = session
+        self.auto_sync = auto_sync
+        self.auto_sync_interval = auto_sync_interval
+        self._last_sync_users_time = 0.0
+        self._users_by_id: dict[str, dict] = {}
 
     def _url_for(self, endpoint: str, **kwargs: str) -> str:
         """
@@ -58,6 +68,36 @@ def users(self) -> list[dict]:
         url = self._url_for("users")
         return self._session.get(url).json()
 
+    def sync_users(self) -> None:
+        """Fetch all users from the server and store them locally."""
+        user_list = self.users()
+        self._users_by_id = {user["id"]: user for user in user_list}
+        self._last_sync_users_time = time.time()
+
+    def sync_users_if_outdated(self) -> None:
+        """Synchronize the user list if the auto-sync interval has elapsed."""
+        timestamp = time.time()
+        if (
+            self.auto_sync
+            and timestamp - self._last_sync_users_time > self.auto_sync_interval
+        ):
+            self.sync_users()
+
+    def get_username(self, user_id: str) -> str | None:
+        """Look up a username by user ID from the locally synced user list.
+
+        Triggers a sync if the local list is outdated. Returns ``None`` if
+        the *user_id* is not found.
+
+        :param user_id: User UUID4.
+        :returns: The username, or ``None`` if not found.
+        """
+        self.sync_users_if_outdated()
+        user = self._users_by_id.get(user_id)
+        if user is not None:
+            return user["username"]
+        return None
+
     def get_user(self, user_id: str) -> dict:
         """
         Get user information.

virl2_client/virl2_client.py

@@ -382,7 +382,11 @@ def __init__(
 
         self.definitions = NodeImageDefinitions(self._session)
         self.licensing = Licensing(self._session)
-        self.user_management = UserManagement(self._session)
+        self.user_management = UserManagement(
+            self._session,
+            auto_sync=self.auto_sync,
+            auto_sync_interval=self.auto_sync_interval,
+        )
         self.group_management = GroupManagement(self._session)
         self.system_management = SystemManagement(
             self._session,
@@ -648,6 +652,7 @@ def _create_imported_lab(
             auto_sync=self.auto_sync,
             auto_sync_interval=self.auto_sync_interval,
             resource_pool_manager=self.resource_pool_management,
+            user_management=self.user_management,
         )
 
     @locked
@@ -797,6 +802,7 @@ def create_lab(
             wait_max_iterations=self.convergence_wait_max_iter,
             wait_time=self.convergence_wait_time,
             resource_pool_manager=self.resource_pool_management,
+            user_management=self.user_management,
         )
         lab._import_lab(result, created=True)
         self._labs[lab_id] = lab
@@ -884,10 +890,12 @@ def join_existing_lab(self, lab_id: str, sync_lab: bool = True) -> Lab:
             wait_max_iterations=self.convergence_wait_max_iter,
             wait_time=self.convergence_wait_time,
             resource_pool_manager=self.resource_pool_management,
+            user_management=self.user_management,
         )
         if sync_lab:
             lab.import_lab(topology)
             lab._initialized = True
+            lab._last_sync_topology_time = time.time()
         else:
             lab._owner = None
         self._labs[lab_id] = lab