python/authn: add cluster operations

- Implemented methods for listing, retrieving, registering, updating, and deleting clusters.

Signed-off-by: Abhishek Gaikwad <[email protected]>
Signed-off-by: Aaron Wilson <[email protected]>

python/aistore/sdk/authn/__init__.py

@@ -1 +1,2 @@
 from aistore.sdk.authn.authn_client import AuthNClient
+from aistore.sdk.authn.cluster_manager import ClusterManager

python/aistore/sdk/authn/authn_client.py

@@ -2,13 +2,19 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 #
 
+import logging
 from typing import Optional, Union
 from aistore.sdk.request_client import RequestClient
 from aistore.sdk.const import (
     HTTP_METHOD_POST,
     URL_PATH_AUTHN_USERS,
 )
 from aistore.sdk.authn.authn_types import TokenMsg, LoginMsg
+from aistore.sdk.authn.cluster_manager import ClusterManager
+
+# logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 # pylint: disable=too-many-arguments, too-few-public-methods
@@ -37,9 +43,21 @@ def __init__(
         timeout: Optional[Union[float, tuple[float, float]]] = None,
         token: Optional[str] = None,
     ):
+        logger.info("Initializing AuthNClient")
         self._request_client = RequestClient(
             endpoint, skip_verify, ca_cert, timeout, token
         )
+        logger.info("AuthNClient initialized with endpoint: %s", endpoint)
+
+    @property
+    def client(self) -> RequestClient:
+        """
+        Get the request client.
+
+        Returns:
+            RequestClient: The client this AuthN client uses to make requests.
+        """
+        return self._request_client
 
     def login(
         self,
@@ -65,11 +83,29 @@ def login(
         if password.strip() == "":
             raise ValueError("Password cannot be empty or spaces only")
 
+        logger.info("Attempting to log in with username: %s", username)
         login_msg = LoginMsg(password=password, expires_in=expires_in).as_dict()
 
-        return self._request_client.request_deserialize(
-            HTTP_METHOD_POST,
-            path=f"{URL_PATH_AUTHN_USERS}/{username}",
-            json=login_msg,
-            res_model=TokenMsg,
-        ).token
+        try:
+            token = self.client.request_deserialize(
+                HTTP_METHOD_POST,
+                path=f"{URL_PATH_AUTHN_USERS}/{username}",
+                json=login_msg,
+                res_model=TokenMsg,
+            ).token
+            logger.info("Login successful for username: %s", username)
+            # Update the client token
+            self.client.token = token
+            return token
+        except Exception as err:
+            logger.error("Login failed for username: %s, error: %s", username, err)
+            raise
+
+    def cluster_manager(self) -> ClusterManager:
+        """
+        Factory method to create a ClusterManager instance.
+
+        Returns:
+            ClusterManager: An instance to manage cluster operations.
+        """
+        return ClusterManager(client=self._request_client)

python/aistore/sdk/authn/authn_types.py

@@ -3,7 +3,7 @@
 #
 
 from __future__ import annotations
-from typing import Optional, Union
+from typing import Optional, Union, Dict, List
 from pydantic import BaseModel, Field
 from aistore.sdk.const import NANOSECONDS_IN_SECOND
 
@@ -47,3 +47,51 @@ class TokenMsg(BaseModel):
     """
 
     token: str
+
+
+class ClusterInfo(BaseModel):
+    """
+    Represents information about a cluster.
+
+    Attributes:
+        id (str): The unique identifier of the cluster.
+        alias (Optional[str]): The alias name of the cluster. Defaults to None.
+        urls (List[str]): A list of URLs associated with the cluster.
+    """
+
+    id: str
+    alias: Optional[str] = None
+    urls: List[str] = []
+
+    def as_dict(self) -> Dict[str, Optional[str] | List[str]]:
+        """
+        Converts the ClusterInfo object to a dictionary.
+
+        Returns:
+            Dict[str, Optional[str] | List[str]]: A dictionary representation of the ClusterInfo object.
+        """
+        return {
+            "id": self.id,
+            "alias": self.alias,
+            "urls": self.urls,
+        }
+
+
+class ClusterList(BaseModel):
+    """
+    Represents a list of clusters.
+
+    Attributes:
+        clusters (Dict[str, ClusterInfo]): A dictionary of cluster IDs to ClusterInfo objects.
+    """
+
+    clusters: Dict[str, ClusterInfo] = {}
+
+    def as_dict(self) -> Dict[str, ClusterInfo]:
+        """
+        Converts the ClusterList object to a dictionary.
+
+        Returns:
+            Dict[str, ClusterInfo]: A dictionary representation of the ClusterList object.
+        """
+        return self.clusters

python/aistore/sdk/authn/cluster_manager.py

@@ -0,0 +1,210 @@
+#
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+#
+
+import logging
+from typing import List, Optional
+from aistore.sdk.request_client import RequestClient
+from aistore.sdk import Client as AISClient
+from aistore.sdk.const import (
+    HTTP_METHOD_GET,
+    HTTP_METHOD_POST,
+    HTTP_METHOD_PUT,
+    HTTP_METHOD_DELETE,
+    URL_PATH_AUTHN_CLUSTERS,
+)
+from aistore.sdk.authn.authn_types import ClusterInfo, ClusterList
+
+# logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class ClusterManager:
+    """
+    ClusterManager class for handling operations on clusters within the context of authentication.
+
+    This class provides methods to list, get, register, update, and delete clusters on AuthN server.
+
+    Args:
+        client (RequestClient): The request client to make HTTP requests.
+    """
+
+    def __init__(self, client: RequestClient):
+        self._client = client
+
+    @property
+    def client(self) -> RequestClient:
+        """RequestClient: The client this cluster manager uses to make requests."""
+        return self._client
+
+    def list(self) -> ClusterList:
+        """
+        Retrieve all clusters.
+
+        Returns:
+            ClusterList: A list of all clusters.
+
+        Raises:
+            AISError: If an error occurs while listing clusters.
+        """
+        logger.info("Listing all clusters")
+        return self.client.request_deserialize(
+            HTTP_METHOD_GET,
+            path=URL_PATH_AUTHN_CLUSTERS,
+            res_model=ClusterList,
+        )
+
+    def get(
+        self, cluster_id: Optional[str] = None, cluster_alias: Optional[str] = None
+    ) -> ClusterInfo:
+        """
+        Retrieve a specific cluster by ID or alias.
+
+        Args:
+            cluster_id (Optional[str]): The ID of the cluster. Defaults to None.
+            cluster_alias (Optional[str]): The alias of the cluster. Defaults to None.
+
+        Returns:
+            ClusterInfo: Information about the specified cluster.
+
+        Raises:
+            ValueError: If neither cluster_id nor cluster_alias is provided.
+            RuntimeError: If no cluster matches the provided ID or alias.
+            AISError: If an error occurs while getting the cluster.
+        """
+        if not cluster_id and not cluster_alias:
+            raise ValueError(
+                "At least one of cluster_id or cluster_alias must be provided"
+            )
+
+        cluster_id_value = cluster_id or cluster_alias
+
+        logger.info("Getting cluster with ID or alias: %s", cluster_id_value)
+
+        cluster_list = self.client.request_deserialize(
+            HTTP_METHOD_GET,
+            path=f"{URL_PATH_AUTHN_CLUSTERS}/{cluster_id_value}",
+            res_model=ClusterList,
+        )
+
+        if not cluster_list.clusters:
+            raise RuntimeError(
+                f"No cluster found with ID or alias '{cluster_id_value}'"
+            )
+
+        first_key = next(iter(cluster_list.clusters))
+        return cluster_list.clusters[first_key]
+
+    def register(self, cluster_alias: str, urls: List[str]) -> ClusterInfo:
+        """
+        Register a new cluster.
+
+        Args:
+            cluster_alias (str): The alias for the new cluster.
+            urls (List[str]): A list of URLs for the new cluster.
+
+        Returns:
+            ClusterInfo: Information about the registered cluster.
+
+        Raises:
+            ValueError: If no URLs are provided or an invalid URL is provided.
+            AISError: If an error occurs while registering the cluster.
+        """
+        if not urls:
+            raise ValueError("At least one URL must be provided")
+
+        logger.info(
+            "Registering new cluster with alias: %s and URLs: %s", cluster_alias, urls
+        )
+
+        try:
+            ais_client = AISClient(
+                endpoint=urls[0], token=self.client.token, skip_verify=True
+            )
+            uuid = ais_client.cluster().get_uuid()
+        except Exception as err:
+            raise ValueError(
+                f"Failed to retrieve UUID for the provided URL: {urls[0]}. "
+                f"Ensure the URL is correct and the endpoint is accessible."
+            ) from err
+
+        if not uuid:
+            raise ValueError("Failed to retrieve UUID for the provided URL")
+
+        self.client.request(
+            HTTP_METHOD_POST,
+            path=URL_PATH_AUTHN_CLUSTERS,
+            json={"id": uuid, "alias": cluster_alias, "urls": urls},
+        )
+
+        return self.get(uuid)
+
+    def update(
+        self,
+        cluster_id: str,
+        cluster_alias: Optional[str] = None,
+        urls: Optional[List[str]] = None,
+    ) -> ClusterInfo:
+        """
+        Update an existing cluster.
+
+        Args:
+            cluster_id (str): The ID of the cluster to update.
+            cluster_alias (Optional[str]): The new alias for the cluster. Defaults to None.
+            urls (Optional[List[str]]): The new list of URLs for the cluster. Defaults to None.
+
+        Returns:
+            ClusterInfo: Information about the updated cluster.
+
+        Raises:
+            ValueError: If neither cluster_alias nor urls are provided.
+            AISError: If an error occurs while updating the cluster
+        """
+        if not cluster_id:
+            raise ValueError("Cluster ID must be provided")
+        if cluster_alias is None and urls is None:
+            raise ValueError("At least one of cluster_alias or urls must be provided")
+
+        cluster_info = ClusterInfo(id=cluster_id)
+        if cluster_alias:
+            cluster_info.alias = cluster_alias
+        if urls:
+            cluster_info.urls = urls
+
+        logger.info("Updating cluster with ID: %s", cluster_id)
+
+        self.client.request(
+            HTTP_METHOD_PUT,
+            path=f"{URL_PATH_AUTHN_CLUSTERS}/{cluster_id}",
+            json=cluster_info.dict(),
+        )
+        return self.get(cluster_id=cluster_id)
+
+    def delete(
+        self, cluster_id: Optional[str] = None, cluster_alias: Optional[str] = None
+    ):
+        """
+        Delete a specific cluster by ID or alias.
+
+        Args:
+            cluster_id (Optional[str]): The ID of the cluster to delete. Defaults to None.
+            cluster_alias (Optional[str]): The alias of the cluster to delete. Defaults to None.
+
+        Raises:
+            ValueError: If neither cluster_id nor cluster_alias is provided.
+            AISError: If an error occurs while deleting the cluster
+        """
+        if not cluster_id and not cluster_alias:
+            raise ValueError(
+                "At least one of cluster_id or cluster_alias must be provided"
+            )
+
+        cluster_id_value = cluster_id or cluster_alias
+
+        logger.info("Deleting cluster with ID or alias: %s", cluster_id_value)
+
+        self.client.request(
+            HTTP_METHOD_DELETE,
+            path=f"{URL_PATH_AUTHN_CLUSTERS}/{cluster_id_value}",
+        )

python/aistore/sdk/cluster.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
 #
 
 from __future__ import annotations  # pylint: disable=unused-variable
@@ -269,3 +269,9 @@ def _get_smap(self):
     def _get_targets(self):
         tmap = self._get_smap().tmap
         return list(tmap.keys())
+
+    def get_uuid(self) -> str:
+        """
+        Returns: UUID of AIStore Cluster
+        """
+        return self._get_smap().uuid

python/aistore/sdk/const.py

@@ -83,6 +83,7 @@
 DSORT_ABORT = "abort"
 # AuthN
 URL_PATH_AUTHN_USERS = "users"
+URL_PATH_AUTHN_CLUSTERS = "clusters"
 
 # Bucket providers
 # See api/apc/provider.go

python/aistore/sdk/request_client.py

@@ -109,6 +109,21 @@ def token(self):
         """
         return self._token
 
+    @token.setter
+    def token(self, token: str):
+        """
+        Set the token for Authorization.
+
+        Args:
+            token (str): Token for Authorization. Must be a non-empty string.
+
+        Raises:
+            ValueError: If the provided token is empty.
+        """
+        if not token:
+            raise ValueError("Token must be a non-empty string.")
+        self._token = token
+
     def request_deserialize(
         self, method: str, path: str, res_model: Type[T], **kwargs
     ) -> T: