azureml-registry-tools 0.1.0a2__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

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/_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/_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.dist-info → azureml_registry_tools-0.1.0a3.dist-info}/METADATA

@@ -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.dist-info → azureml_registry_tools-0.1.0a3.dist-info}/RECORD

@@ -4,9 +4,9 @@ azureml/registry/_cli/__init__.py,sha256=dJ020ynrregpUaNGu8xVLLmX3cC9DAjaBMP1qnl
 azureml/registry/_cli/registry_syndication_cli.py,sha256=87hxSu4iwFOdwfukSGqc0fjJEpSpQhDRlthiCGCDwQ8,10003
 azureml/registry/_cli/repo2registry_cli.py,sha256=4CDv4tYWLTjVgdXIcoKA9iu_gOv_Z_xn05wSANkdmdg,5378
 azureml/registry/_rest_client/__init__.py,sha256=Eh0vB8AVlwkZlHLO4DCGHyyfRYYgeTeVjiUAX_WujG0,219
-azureml/registry/_rest_client/arm_client.py,sha256=dgm9nDnmGI1NIZaopXCIXmViM8sSAipZETfuOWnQRa0,5638
-azureml/registry/_rest_client/base_rest_client.py,sha256=cXFYvYgSYSveEeY511VIlDpmXLOY9Jat_E-Zs1ehMyA,5164
-azureml/registry/_rest_client/registry_management_client.py,sha256=JGBohCVpL8XR4rNuScOKt7tNQnjw203pqlTSLKVjYgM,4833
+azureml/registry/_rest_client/arm_client.py,sha256=KVSj0V1bN-vCUV3fBbIdZDpWTLKYT6qdWIZtDjx2la4,4516
+azureml/registry/_rest_client/base_rest_client.py,sha256=JcmhckfR-ZdUtEh0EeQIKfe63Zjowa1fLzmVeWpeLno,8173
+azureml/registry/_rest_client/registry_management_client.py,sha256=P2aijKHdHvsSZ507-06Rq_n9pNJ_xci1oQKDGBJjm-I,4986
 azureml/registry/mgmt/__init__.py,sha256=LMhqcEC8ItmmpKZljElGXH-6olHlT3SLl0dJU01OvuM,226
 azureml/registry/mgmt/create_manifest.py,sha256=N9wRmjAKO09A3utN_lCUsM_Ufpj7PL0SJz-XHPHWuyM,9528
 azureml/registry/mgmt/syndication_manifest.py,sha256=ma1_vUEHHtkRcgiFMQanlXZajvKodynylxfblTHuleE,6195
@@ -15,9 +15,9 @@ azureml/registry/tools/config.py,sha256=tjPaoBsWtPXBL8Ww1hcJtsr2SuIjPKt79dR8iovc
 azureml/registry/tools/create_or_update_assets.py,sha256=NexSY22bpN2eHPqZZ21eDIMerwq3SzVk_4fBV5_GTBY,14787
 azureml/registry/tools/registry_utils.py,sha256=zgYlCiOONtQJ4yZ9wg8tKVoE8dh6rrjB8hYBGhpV9-0,1403
 azureml/registry/tools/repo2registry_config.py,sha256=eXp_tU8Jyi30g8xGf7wbpLgKEPpieohBANKxMSLzq7s,4873
-azureml_registry_tools-0.1.0a2.dist-info/licenses/LICENSE.txt,sha256=n20rxwp7_NGrrShv9Qvcs90sjI1l3Pkt3m-5OPCWzgs,845
-azureml_registry_tools-0.1.0a2.dist-info/METADATA,sha256=SJpiD2SsYXqJi_DUSLk-mquFZgbnb54Sd_mrP6H-GBI,491
-azureml_registry_tools-0.1.0a2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-azureml_registry_tools-0.1.0a2.dist-info/entry_points.txt,sha256=iRUkAeQidMnO6RQzpLqMUBTcyYtNzAfSin9WnSdVGLw,147
-azureml_registry_tools-0.1.0a2.dist-info/top_level.txt,sha256=ZOeEa0TAXo6i5wOjwBoqfIGEuxOcKuscGgNSpizqREY,8
-azureml_registry_tools-0.1.0a2.dist-info/RECORD,,
+azureml_registry_tools-0.1.0a3.dist-info/licenses/LICENSE.txt,sha256=n20rxwp7_NGrrShv9Qvcs90sjI1l3Pkt3m-5OPCWzgs,845
+azureml_registry_tools-0.1.0a3.dist-info/METADATA,sha256=0UQUDwCZO5K7-EnWwtUFqrSHJeXX2DU14PUwuoA2df0,521
+azureml_registry_tools-0.1.0a3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+azureml_registry_tools-0.1.0a3.dist-info/entry_points.txt,sha256=iRUkAeQidMnO6RQzpLqMUBTcyYtNzAfSin9WnSdVGLw,147
+azureml_registry_tools-0.1.0a3.dist-info/top_level.txt,sha256=ZOeEa0TAXo6i5wOjwBoqfIGEuxOcKuscGgNSpizqREY,8
+azureml_registry_tools-0.1.0a3.dist-info/RECORD,,