partition_client.py 1.23 KB
Newer Older
Luc Yriarte's avatar
Luc Yriarte committed
1
from typing import Optional
Yannick's avatar
Yannick committed
2
import logging
Luc Yriarte's avatar
Luc Yriarte committed
3
4
5
6
7
8
9
10
11
import aiohttp

from osdu_az import conf
from osdu_az.exceptions.data_access_error import DataAccessError
from osdu_az.identity.az_identity import AzureIdentity
from osdu_az.partition.partition_info import PartitionInfo

error_message = 'Failed to retrieve data partition'

Yannick's avatar
Yannick committed
12
13
_LOGGER = logging.getLogger(__name__)

Luc Yriarte's avatar
Luc Yriarte committed
14
15
16
17
18
19
20
21
22

class PartitionClient:

    @staticmethod
    def partition_endpoint() -> str:
        return conf.partition_api_url + '/v1/partitions/'

    @staticmethod
    async def get_partition(data_partition_id: str) -> Optional[PartitionInfo]:
Yannick's avatar
Yannick committed
23
        access_token = AzureIdentity.get_access_token()
Luc Yriarte's avatar
Luc Yriarte committed
24
25
26
27
28
29
30
31
32
33
34
35
36

        headers = {
            'authorization': f'Bearer {access_token}'
        }

        async with aiohttp.ClientSession(headers=headers) as session:
            async with session.get(PartitionClient.partition_endpoint() + data_partition_id) as response:
                if response.status == 200:
                    partition_properties = await response.json()
                    return PartitionInfo(partition_properties)

                raise DataAccessError(message=f'{error_message} {data_partition_id}. {response.reason}.',
                                      status_code=response.status)