azureml-registry-tools 0.1.0a2__tar.gz → 0.1.0a3__tar.gz

{azureml_registry_tools-0.1.0a2/azureml_registry_tools.egg-info → azureml_registry_tools-0.1.0a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: azureml-registry-tools
-Version: 0.1.0a2
+Version: 0.1.0a3
 Summary: AzureML Registry tools and CLI
 Author: Microsoft Corp
 License: https://aka.ms/azureml-sdk-license
@@ -8,6 +8,7 @@ Requires-Python: >=3.9,<3.14
 License-File: LICENSE.txt
 Requires-Dist: azure-identity<2.0
 Requires-Dist: ruamel-yaml<0.19,>=0.17.21
+Requires-Dist: diskcache~=5.6
 Requires-Dist: azure-ai-ml<2.0
 Requires-Dist: azureml-assets<2.0
 Dynamic: author

{azureml_registry_tools-0.1.0a2 → azureml_registry_tools-0.1.0a3}/azureml/registry/_rest_client/arm_client.py

@@ -1,16 +1,14 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-from azure.identity import DefaultAzureCredential
-from .base_rest_client import BaseRestClient
-import time
+from .base_rest_client import BaseAzureRestClient
 from json.decoder import JSONDecodeError
 
 DEFAULT_API_VERSION = "2025-04-01"  # Default API version for Azure Resource Manager
 
 
-class ArmClient(BaseRestClient):
-    """Simple Azure Resource Manager (ARM) client leveraging BaseRestClient for GET, PATCH, PUT, and DELETE operations.
+class ArmClient(BaseAzureRestClient):
+    """Simple Azure Resource Manager (ARM) client leveraging BaseAzureRestClient for GET, PATCH, PUT, and DELETE operations.
 
     Handles authentication via Bearer token (pass as api_key) and supports standard ARM resource operations.
     """
@@ -25,27 +23,8 @@ class ArmClient(BaseRestClient):
             backoff_factor (int): Backoff factor for retries.
         """
         base_url = "https://management.azure.com"
-        self._credential = None
-        self._token_expires_on = None
-        if api_key is None:
-            # Use DefaultAzureCredential for authentication if no API key is provided
-            self._credential = DefaultAzureCredential()
-            token = self._credential.get_token("https://management.azure.com/.default")
-            api_key = token.token
-            self._token_expires_on = token.expires_on
         super().__init__(base_url, api_key=api_key, max_retries=max_retries, backoff_factor=backoff_factor)
 
-    def _refresh_api_key_if_needed(self) -> None:
-        """Refresh the API key if using DefaultAzureCredential and the token is close to expiration."""
-        if self._credential is not None:
-            now = int(time.time())
-            # Refresh if less than 10 minutes (600 seconds) left
-            if not self._token_expires_on or self._token_expires_on - now < 600:
-                token = self._credential.get_token("https://management.azure.com/.default")
-                self.api_key = token.token
-                self.session.headers.update({'Authorization': f'Bearer {self.api_key}'})
-                self._token_expires_on = token.expires_on
-
     def get_resource(self, resource_id, api_version=DEFAULT_API_VERSION, **kwargs) -> object:
         """
         Get an ARM resource by its resource ID.

{azureml_registry_tools-0.1.0a2 → azureml_registry_tools-0.1.0a3}/azureml/registry/_rest_client/base_rest_client.py

@@ -4,6 +4,11 @@
 import requests
 import time
 import logging
+from azure.identity import DefaultAzureCredential
+from diskcache import Cache
+import tempfile
+
+ENABLE_CACHE = True  # Set to False to disable diskcache usage for API key
 
 
 class BaseRestClient:
@@ -136,3 +141,61 @@ class BaseRestClient:
             requests.Response: The HTTP response object.
         """
         return self._request_with_retry('DELETE', url, **kwargs)
+
+
+class BaseAzureRestClient(BaseRestClient):
+    """Base Azure REST client that handles Azure authentication and token refresh, with diskcache for API key."""
+
+    if ENABLE_CACHE:
+        _cache_dir = tempfile.gettempdir() + '/.azureml_registry_token_cache'
+        _cache = Cache(_cache_dir)
+    else:
+        _cache = None
+
+    def __init__(self, base_url: str, api_key: str = None, max_retries: int = 5, backoff_factor: int = 1):
+        """
+        Initialize the BaseAzureRestClient.
+
+        Args:
+            base_url (str): The base URL for the REST API.
+            api_key (str, optional): Bearer token for authentication. If None, uses DefaultAzureCredential.
+            max_retries (int): Maximum number of retries for failed requests.
+            backoff_factor (int): Backoff factor for retry delays.
+        """
+        self._credential = None
+        self._token_expires_on = None
+        cache_key = f"azureml_api_key_{base_url}"
+        cache_expiry_key = f"azureml_api_key_expiry_{base_url}"
+        cached_token = None
+        cached_expiry = None
+        now = int(time.time())
+        if ENABLE_CACHE and self._cache is not None:
+            cached_token = self._cache.get(cache_key)
+            cached_expiry = self._cache.get(cache_expiry_key)
+        if api_key is None and cached_token and cached_expiry and cached_expiry > now:
+            api_key = cached_token
+            self._token_expires_on = cached_expiry
+        super().__init__(base_url, api_key=api_key, max_retries=max_retries, backoff_factor=backoff_factor)
+        # Only after super().__init__ is self.session available
+        if api_key is None:
+            self._credential = DefaultAzureCredential()
+            self._refresh_api_key_if_needed()
+            # Ensure self.api_key is set for future use
+            api_key = self.api_key
+
+    def _refresh_api_key_if_needed(self) -> None:
+        if self._credential is None:
+            return
+        now = int(time.time())
+        if not getattr(self, '_token_expires_on', None) or self._token_expires_on - now < 600:
+            token = self._credential.get_token("https://management.azure.com/.default")
+            self.api_key = token.token
+            self.session.headers.update({'Authorization': f'Bearer {self.api_key}'})
+            self._token_expires_on = token.expires_on
+            # Store in diskcache if enabled
+            if ENABLE_CACHE and self._cache is not None:
+                base_url = getattr(self, 'base_url', 'default')
+                cache_key = f"azureml_api_key_{base_url}"
+                cache_expiry_key = f"azureml_api_key_expiry_{base_url}"
+                self._cache.set(cache_key, self.api_key, expire=(self._token_expires_on - now))
+                self._cache.set(cache_expiry_key, self._token_expires_on, expire=(self._token_expires_on - now))

{azureml_registry_tools-0.1.0a2 → azureml_registry_tools-0.1.0a3}/azureml/registry/_rest_client/registry_management_client.py

@@ -1,18 +1,16 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-from azure.identity import DefaultAzureCredential
-from .base_rest_client import BaseRestClient
-import time
+from .base_rest_client import BaseAzureRestClient
 
 
-class RegistryManagementClient(BaseRestClient):
+class RegistryManagementClient(BaseAzureRestClient):
     """Python client for RegistrySyndicationManifestController (excluding S2S APIs).
 
     Handles authentication, token refresh, and provides methods for manifest and registry management.
     """
 
-    def __init__(self, registry_name: str, primary_region: str = "eastus2euap", api_key: str = None, max_retries: int = 5, backoff_factor: int = 1) -> None:
+    def __init__(self, registry_name: str, primary_region: str = None, api_key: str = None, max_retries: int = 5, backoff_factor: int = 1) -> None:
         """
         Initialize the RegistryManagementClient.
 
@@ -23,29 +21,26 @@ class RegistryManagementClient(BaseRestClient):
             max_retries (int): Maximum number of retries for failed requests.
             backoff_factor (int): Backoff factor for retry delays.
         """
+        if primary_region is None:
+            # Resolve the primary region if not provided
+            primary_region = self.resolve_registry_primary_region(registry_name)
         base_url = f"https://{primary_region}.api.azureml.ms"
-        self._credential = None
-        self._token_expires_on = None
-        if api_key is None:
-            # Use DefaultAzureCredential for authentication if no API key is provided
-            self._credential = DefaultAzureCredential()
-            token = self._credential.get_token("https://management.azure.com/.default")
-            api_key = token.token
-            self._token_expires_on = token.expires_on
         super().__init__(base_url, api_key=api_key, max_retries=max_retries, backoff_factor=backoff_factor)
         self.registry_name = registry_name
 
-    def _refresh_api_key_if_needed(self) -> None:
-        """Refresh the API key if using DefaultAzureCredential and the token is close to expiration."""
-        # Only refresh if using DefaultAzureCredential
-        if self._credential is not None:
-            now = int(time.time())
-            # Refresh if less than 10 minutes (600 seconds) left
-            if not self._token_expires_on or self._token_expires_on - now < 600:
-                token = self._credential.get_token("https://management.azure.com/.default")
-                self.api_key = token.token
-                self.session.headers.update({'Authorization': f'Bearer {self.api_key}'})
-                self._token_expires_on = token.expires_on
+    @staticmethod
+    def resolve_registry_primary_region(registry_name: str) -> str:
+        """
+        Resolve the primary region for the given registry name.
+
+        Args:
+            registry_name (str): The name of the AzureML registry.
+
+        Returns:
+            str: The primary region for the registry.
+        """
+        discovery = RegistryManagementClient(registry_name, primary_region="eastus").discovery()
+        return discovery.get('primaryRegion', 'eastus')
 
     def create_or_update_manifest(self, manifest_dto: dict) -> dict:
         """
@@ -110,5 +105,15 @@ class RegistryManagementClient(BaseRestClient):
         """
         self._refresh_api_key_if_needed()
         url = f"{self.base_url}/registrymanagement/v1.0/registries/{self.registry_name}/discovery"
-        response = self.get(url)
-        return response.json()
+        try:
+            response = self.get(url)
+            return response.json()
+        except Exception as ex:
+            # Special handling for HTTP errors with status_code attribute
+            if hasattr(ex, 'response') and ex.response is not None:
+                status_code = ex.response.status_code
+                if status_code == 403:
+                    raise RuntimeError(f"Received 403 Forbidden. This may indicate that the registry '{self.registry_name}' does not exist or you do not have access.")
+                else:
+                    raise RuntimeError(f"Failed to get discovery information: {status_code} {ex.response.text}")
+            raise RuntimeError(f"Error occurred while trying to get discovery information: {ex}")

{azureml_registry_tools-0.1.0a2 → azureml_registry_tools-0.1.0a3/azureml_registry_tools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: azureml-registry-tools
-Version: 0.1.0a2
+Version: 0.1.0a3
 Summary: AzureML Registry tools and CLI
 Author: Microsoft Corp
 License: https://aka.ms/azureml-sdk-license
@@ -8,6 +8,7 @@ Requires-Python: >=3.9,<3.14
 License-File: LICENSE.txt
 Requires-Dist: azure-identity<2.0
 Requires-Dist: ruamel-yaml<0.19,>=0.17.21
+Requires-Dist: diskcache~=5.6
 Requires-Dist: azure-ai-ml<2.0
 Requires-Dist: azureml-assets<2.0
 Dynamic: author

{azureml_registry_tools-0.1.0a2 → azureml_registry_tools-0.1.0a3}/azureml_registry_tools.egg-info/requires.txt

@@ -1,4 +1,5 @@
 azure-identity<2.0
 ruamel-yaml<0.19,>=0.17.21
+diskcache~=5.6
 azure-ai-ml<2.0
 azureml-assets<2.0

{azureml_registry_tools-0.1.0a2 → azureml_registry_tools-0.1.0a3}/setup.py

@@ -9,6 +9,7 @@ from setuptools import setup, find_packages
 DEPENDENCIES = [
     "azure-identity<2.0",
     "ruamel-yaml>=0.17.21,<0.19",
+    "diskcache~=5.6",
     "azure-ai-ml<2.0",
     "azureml-assets<2.0"
 ]
@@ -17,7 +18,7 @@ exclude_list = ["*.tests"]
 
 setup(
     name='azureml-registry-tools',
-    version="0.1.0a2",
+    version="0.1.0a3",
     description='AzureML Registry tools and CLI',
     author='Microsoft Corp',
     license="https://aka.ms/azureml-sdk-license",