Added get_workspace_status method to management API

CHANGELOG.md

@@ -1,4 +1,7 @@
-## [UNRELEASED] neptune 1.9.2
+## [UNRELEASED] neptune 1.10.0
+
+### Features
+- Added `get_workspace_status()` method to management API ([#1662](https://github.com/neptune-ai/neptune-client/pull/1662))
 
 ### Fixes
 - Restored support for SSL verification exception ([#1661](https://github.com/neptune-ai/neptune-client/pull/1661))

src/neptune/management/__init__.py

@@ -39,6 +39,9 @@
 Usage examples
 --------------
 
+Import management API
+>>> from neptune import management
+
 Getting projects in a workspace as a list:
 >>> projects = management.get_project_list()
 
@@ -95,6 +98,9 @@
 >>> # Move the runs to trash:
 ... management.trash_objects(project=project_name, ids=runs_to_trash)
 
+Get information about a workspace, including storage usage and limits:
+>>> management.get_workspace_status(workspace="ml-team")
+
 ---
 
 See also the API reference in the docs: https://docs.neptune.ai/api/management
@@ -112,6 +118,7 @@
     get_project_service_account_list,
     get_workspace_member_list,
     get_workspace_service_account_list,
+    get_workspace_status,
     invite_to_workspace,
     remove_project_member,
     remove_project_service_account,
@@ -138,6 +145,7 @@
     "remove_project_service_account",
     "get_project_service_account_list",
     "get_workspace_service_account_list",
+    "get_workspace_status",
     "trash_objects",
     "MemberRole",
     "ProjectVisibility",

src/neptune/management/internal/api.py

@@ -31,6 +31,7 @@
     "get_project_service_account_list",
     "get_workspace_service_account_list",
     "trash_objects",
+    "get_workspace_status",
 ]
 
 import os
@@ -1033,3 +1034,66 @@ def clear_trash(
 
     for error in response.result.errors:
         logger.warning(error)
+
+
+def get_workspace_status(workspace: str, *, api_token: Optional[str] = None) -> Dict[str, int]:
+    """Retrieves status information about a Neptune workspace.
+
+    Includes the following:
+
+    - Storage usage and limit
+    - Active project count and limit
+    - Member count
+
+    Args:
+        workspace: Name of the Neptune workspace.
+        api_token: Account's API token.
+            If None, the value of the NEPTUNE_API_TOKEN environment variable is used.
+            Note: To keep your token secure, use the NEPTUNE_API_TOKEN environment variable rather than placing your
+            API token in plain text in your source code.
+
+    Returns:
+        Dictionary with metric name as keys and float values
+
+    Example:
+        >>> from neptune import management
+        >>> management.get_workspace_status(workspace="ml-team")
+        ... {'storageBytesAvailable': 214747451765,
+        ... 'storageBytesLimit': 214748364800,
+        ... 'storageBytesUsed': 913035,
+        ... 'activeProjectsUsage': 1,
+        ... 'activeProjectsLimit': 1,
+        ... 'membersCount': 1}
+
+    You may also want to check the management API reference:
+    https://docs.neptune.ai/api/management/#get_workspace_status
+    """
+    verify_type("workspace", workspace, str)
+    verify_type("api_token", api_token, (str, type(None)))
+
+    backend_client = _get_backend_client(api_token=api_token)
+
+    params = {
+        "organizationIdentifier": workspace,
+        **DEFAULT_REQUEST_KWARGS,
+    }
+
+    try:
+        response = backend_client.api.workspaceStatus(**params).response()
+
+        result = dict()
+        if hasattr(response.result, "storageBytesAvailable"):
+            result["storageBytesAvailable"] = response.result.storageBytesAvailable
+        if hasattr(response.result, "storageBytesLimit"):
+            result["storageBytesLimit"] = response.result.storageBytesLimit
+        if hasattr(response.result, "storageBytesAvailable") and hasattr(response.result, "storageBytesLimit"):
+            result["storageBytesUsed"] = response.result.storageBytesLimit - response.result.storageBytesAvailable
+        if hasattr(response.result, "activeProjectsUsage"):
+            result["activeProjectsUsage"] = response.result.activeProjectsUsage
+        if hasattr(response.result, "activeProjectsLimit"):
+            result["activeProjectsLimit"] = response.result.activeProjectsLimit
+        if hasattr(response.result, "membersCount"):
+            result["membersCount"] = response.result.membersCount
+        return result
+    except HTTPNotFound as e:
+        raise WorkspaceNotFound(workspace=workspace) from e

tests/e2e/management/test_management.py

@@ -41,6 +41,7 @@
     get_project_service_account_list,
     get_workspace_member_list,
     get_workspace_service_account_list,
+    get_workspace_status,
     invite_to_workspace,
     remove_project_member,
     remove_project_service_account,
@@ -414,6 +415,16 @@ def test_invite_to_workspace(self, environment: "Environment"):
                 username=environment.user, workspace="non-existent-workspace", api_token=environment.admin_token
             )
 
+    def test_workspace_status(self, environment: "Environment"):
+        status = get_workspace_status(workspace=environment.workspace, api_token=environment.admin_token)
+
+        assert "storageBytesAvailable" in status
+        assert "storageBytesLimit" in status
+        assert "storageBytesUsed" in status
+        assert status["storageBytesAvailable"] >= 0
+        assert status["storageBytesLimit"] >= 0
+        assert status["storageBytesUsed"] >= 0
+
 
 @pytest.mark.management
 class TestTrashObjects(BaseE2ETest):