use async partition service

osdu_az/identity/az_identity.py

-from azure.identity import DefaultAzureCredential
-from azure.keyvault.secrets import SecretClient
+import logging
+
+from azure.identity.aio import DefaultAzureCredential
+from azure.keyvault.secrets.aio import SecretClient
 
 from osdu_az import conf
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class AzureIdentity:
     default_credential = None
     default_scope = ''
+    _secret_client = None
+
+    @classmethod
+    def get_default_credential(cls):
+        if not cls.default_credential:
+            cls.default_credential = DefaultAzureCredential()
+        return cls.default_credential
+
+    @classmethod
+    async def close_credentials(cls):
+        """ This cause to gracefully dispose credentials if any. Next calls will then initialize a new one """
+        _LOGGER.info(f"Closing cached Credentials")
+        credentials_to_close, cls.default_credential = cls.default_credential, None  # swap
+        if credentials_to_close is not None:
+            await credentials_to_close.close()
+
+    @classmethod
+    async def get_access_token(cls):
+        credential = cls.get_default_credential()
+        scope = await cls._get_scope()
+        access_token = await credential.get_token(scope)
+        return access_token.token
+
+    @classmethod
+    async def _get_scope(cls):
+        if not cls.default_scope:
+            cls.default_scope = await cls._get_resource_id()
+        return cls.default_scope
+
+    @classmethod
+    async def _get_resource_id(cls) -> str:
+        return await cls.get_secret('aad-client-id')
 
-    @staticmethod
-    def get_default_credential():
-        if not AzureIdentity.default_credential:
-            AzureIdentity.default_credential = DefaultAzureCredential()
-        return AzureIdentity.default_credential
-
-    @staticmethod
-    def get_access_token():
-        credential = AzureIdentity.get_default_credential()
-        return credential.get_token(AzureIdentity.get_scope()).token
-
-    @staticmethod
-    def get_scope():
-        if not AzureIdentity.default_scope:
-            AzureIdentity.default_scope = AzureIdentity.get_resource_id()
-        return AzureIdentity.default_scope
-
-    @staticmethod
-    def get_resource_id() -> str:
-        secret_client = SecretClient(conf.keyvault_url, AzureIdentity.get_default_credential())
-        secret = secret_client.get_secret('aad-client-id')
+    @classmethod
+    async def get_secret(cls, name) -> str:
+        if cls._secret_client is None:
+            cls._secret_client = SecretClient(conf.keyvault_url, cls.get_default_credential())
+        secret = await cls._secret_client.get_secret(name)
         return secret.value

osdu_az/opendes/partition_client.py

 from typing import Optional
-
+import logging
 import aiohttp
 
 from osdu_az import conf
@@ -9,6 +9,8 @@ from osdu_az.partition.partition_info import PartitionInfo
 
 error_message = 'Failed to retrieve data partition'
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class PartitionClient:
 
@@ -18,7 +20,7 @@ class PartitionClient:
 
     @staticmethod
     async def get_partition(data_partition_id: str) -> Optional[PartitionInfo]:
-        access_token = AzureIdentity.get_access_token()
+        access_token = await AzureIdentity.get_access_token()
 
         headers = {
             'authorization': f'Bearer {access_token}'

osdu_az/partition/partition_info.py

@@ -2,8 +2,6 @@ import logging
 from typing import Optional
 from datetime import datetime
 
-from azure.keyvault.secrets import SecretClient
-
 from osdu_az import conf
 from osdu_az.identity.az_identity import AzureIdentity
 
@@ -14,9 +12,8 @@ _LOGGER = logging.getLogger(__name__)
 class PartitionInfo:
     def __init__(self, partition_properties: dict = None):
         self._partition_properties = partition_properties
-        self._secret_client = None
 
-    def get_value(self, property_name: str) -> Optional[str]:
+    async def get_value(self, property_name: str) -> Optional[str]:
         partition_property = self._partition_properties.get(property_name)
         if not partition_property:
             return None
@@ -25,17 +22,13 @@ class PartitionInfo:
             return partition_property['value']
 
         if 'secret' not in partition_property:
-            partition_property['secret'] = self._get_secret(partition_property['value']).value
+            partition_property['secret'] = await self._get_secret(partition_property['value'])
         return partition_property['secret']
 
-    def _get_secret(self, key: str):
+    async def _get_secret(self, key: str) -> str:
         ts = datetime.utcnow()
-        secret = self._get_secret_client().get_secret(key)
+        secret = await AzureIdentity.get_secret(key)
         _LOGGER.info(f'PartitionInfo get secret took {(datetime.utcnow() - ts).total_seconds()} ms')
         return secret
 
-    def _get_secret_client(self):
-        if not self._secret_client:
-            self._secret_client = SecretClient(conf.keyvault_url, AzureIdentity.get_default_credential())
-        return self._secret_client
 

osdu_az/partition/partition_service.py

@@ -41,7 +41,7 @@ class PartitionService:
     async def get_storage_account_name(data_partition_id: str):
         partition_info = await PartitionService.get_partition(data_partition_id)
         if partition_info:
-            return partition_info.get_value(STORAGE_ACCOUNT_NAME)
+            return await partition_info.get_value(STORAGE_ACCOUNT_NAME)
 
     @staticmethod
     def _partition_client():