Commit f10ab48a authored by Yannick's avatar Yannick
Browse files

use async partition service

parent 677d74d8
Pipeline #45016 failed with stage
in 2 minutes and 57 seconds
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
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}'
......
......@@ -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
......@@ -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():
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment