partition_client.py 1.17 KB
Newer Older
Luc Yriarte's avatar
Luc Yriarte committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from typing import Optional

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'


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]:
        access_token = AzureIdentity.get_access_token()

        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)