1 files changed, 79 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_utils/azure_resource_utils.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_utils/azure_resource_utils.py
new file mode 100644
index 00000000..c038b7d0
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_utils/azure_resource_utils.py
@@ -0,0 +1,79 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+from typing import Dict, List, Optional
+
+from azure.core.credentials import TokenCredential
+
+from .._vendor.azure_resources import ResourceManagementClient
+
+
+def get_resources_from_subscriptions(
+    strQuery: str, credential: TokenCredential, subscription_list: Optional[List[str]] = None
+):
+    # If a subscription list is passed in, use it. Otherwise, get all subscriptions
+    subsList = []
+    if subscription_list is not None:
+        subsList = subscription_list
+    else:
+        try:
+            from azure.mgmt.resource import SubscriptionClient
+        except ImportError as e:
+            raise ImportError("azure-mgmt-resource is required to get all accessible subscriptions") from e
+
+        subsClient = SubscriptionClient(credential)
+        for sub in subsClient.subscriptions.list():
+            subsList.append(sub.as_dict().get("subscription_id"))
+
+    try:
+        import azure.mgmt.resourcegraph as arg
+    except ImportError as e:
+        raise ImportError("azure-mgmt-resourcegraph is required query resources from subscriptions") from e
+
+    # Create Azure Resource Graph client and set options
+    argClient = arg.ResourceGraphClient(credential)
+    argQueryOptions = arg.models.QueryRequestOptions(result_format="objectArray")
+
+    # Create query
+    argQuery = arg.models.QueryRequest(subscriptions=subsList, query=strQuery, options=argQueryOptions)
+
+    # Allowing API version to be set is yet to be released by azure-mgmt-resourcegraph,
+    # hence the commented out code below. This is the API version Studio UX is using.
+    # return argClient.resources(argQuery, api_version="2021-03-01")
+
+    return argClient.resources(argQuery)
+
+
+def get_virtual_clusters_from_subscriptions(
+    credential: TokenCredential, subscription_list: Optional[List[str]] = None
+) -> List[Dict]:
+    # cspell:ignore tolower
+    strQuery = """resources
+    | where type == 'microsoft.machinelearningservices/virtualclusters'
+    | order by tolower(name) asc
+    | project id, subscriptionId, resourceGroup, name, location, tags, type"""
+
+    return get_resources_from_subscriptions(strQuery, credential, subscription_list).data
+
+
+def get_generic_resource_by_id(
+    arm_id: str, credential: TokenCredential, subscription_id: str, api_version: Optional[str] = None
+) -> Dict:
+    resource_client = ResourceManagementClient(credential, subscription_id)
+    generic_resource = resource_client.resources.get_by_id(arm_id, api_version)
+
+    return generic_resource.as_dict()
+
+
+def get_virtual_cluster_by_name(
+    name: str, resource_group: str, subscription_id: str, credential: TokenCredential
+) -> Dict:
+    arm_id = (
+        f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}"
+        f"/providers/Microsoft.MachineLearningServices/virtualClusters/{name}"
+    )
+
+    # This is the API version Studio UX is using.
+    return get_generic_resource_by_id(arm_id, credential, subscription_id, api_version="2021-03-01-preview")