|
| 1 | +import dataclasses |
| 2 | +import uuid |
| 3 | +from typing import Callable |
| 4 | + |
| 5 | +import azure.cosmos.cosmos_client as cosmos_client |
| 6 | +import azure.cosmos.exceptions as exceptions |
| 7 | +from azure.cosmos import ContainerProxy |
| 8 | +from flask import Flask |
| 9 | + |
| 10 | + |
| 11 | +class CosmosDBFacade: |
| 12 | + def __init__(self, app: Flask): # pragma: no cover |
| 13 | + self.app = app |
| 14 | + |
| 15 | + db_uri = app.config.get('DATABASE_URI') |
| 16 | + if db_uri is None: |
| 17 | + app.logger.warn("DATABASE_URI was not found. Looking for alternative variables.") |
| 18 | + account_uri = app.config.get('DATABASE_ACCOUNT_URI') |
| 19 | + if account_uri is None: |
| 20 | + raise EnvironmentError("DATABASE_ACCOUNT_URI is not defined in the environment") |
| 21 | + |
| 22 | + master_key = app.config.get('DATABASE_MASTER_KEY') |
| 23 | + if master_key is None: |
| 24 | + raise EnvironmentError("DATABASE_MASTER_KEY is not defined in the environment") |
| 25 | + |
| 26 | + self.client = cosmos_client.CosmosClient(account_uri, {'masterKey': master_key}, |
| 27 | + user_agent="CosmosDBDotnetQuickstart", |
| 28 | + user_agent_overwrite=True) |
| 29 | + else: |
| 30 | + self.client = cosmos_client.CosmosClient.from_connection_string(db_uri) |
| 31 | + |
| 32 | + db_id = app.config.get('DATABASE_NAME') |
| 33 | + if db_id is None: |
| 34 | + raise EnvironmentError("DATABASE_NAME is not defined in the environment") |
| 35 | + |
| 36 | + self.db = self.client.get_database_client(db_id) |
| 37 | + |
| 38 | + def create_container(self, container_definition: dict): |
| 39 | + try: |
| 40 | + return self.db.create_container(**container_definition) |
| 41 | + |
| 42 | + except exceptions.CosmosResourceExistsError: # pragma: no cover |
| 43 | + self.app.logger.info('Container with id \'{0}\' was found'.format(container_definition["id"])) |
| 44 | + |
| 45 | + def delete_container(self, container_id: str): |
| 46 | + try: |
| 47 | + return self.db.delete_container(container_id) |
| 48 | + |
| 49 | + except exceptions.CosmosHttpResponseError: # pragma: no cover |
| 50 | + self.app.logger.info('Container with id \'{0}\' was not deleted'.format(container_id)) |
| 51 | + |
| 52 | + |
| 53 | +cosmos_helper: CosmosDBFacade = None |
| 54 | + |
| 55 | + |
| 56 | +class CosmosDBModel(): |
| 57 | + def __init__(self, data): |
| 58 | + names = set([f.name for f in dataclasses.fields(self)]) |
| 59 | + for k, v in data.items(): |
| 60 | + if k in names: |
| 61 | + setattr(self, k, v) |
| 62 | + |
| 63 | + |
| 64 | +class CosmosDBRepository: |
| 65 | + def __init__(self, container_id: str, |
| 66 | + mapper: Callable = None, |
| 67 | + custom_cosmos_helper: CosmosDBFacade = None): |
| 68 | + global cosmos_helper |
| 69 | + self.cosmos_helper = custom_cosmos_helper or cosmos_helper |
| 70 | + if self.cosmos_helper is None: # pragma: no cover |
| 71 | + raise ValueError("The cosmos_db module has not been initialized!") |
| 72 | + self.mapper = mapper |
| 73 | + self.container: ContainerProxy = self.cosmos_helper.db.get_container_client(container_id) |
| 74 | + |
| 75 | + @classmethod |
| 76 | + def from_definition(cls, container_definition: dict, |
| 77 | + mapper: Callable = None, |
| 78 | + custom_cosmos_helper: CosmosDBFacade = None): |
| 79 | + return cls(container_definition['id'], mapper, custom_cosmos_helper) |
| 80 | + |
| 81 | + def create(self, data: dict, mapper: Callable = None): |
| 82 | + function_mapper = self.get_mapper_or_dict(mapper) |
| 83 | + return function_mapper(self.container.create_item(body=data)) |
| 84 | + |
| 85 | + def find(self, id: str, partition_key_value, visible_only=True, mapper: Callable = None): |
| 86 | + found_item = self.container.read_item(id, partition_key_value) |
| 87 | + function_mapper = self.get_mapper_or_dict(mapper) |
| 88 | + return function_mapper(self.check_visibility(found_item, visible_only)) |
| 89 | + |
| 90 | + def find_all(self, partition_key_value: str, max_count=None, offset=0, |
| 91 | + visible_only=True, mapper: Callable = None): |
| 92 | + # TODO Use the tenant_id param and change container alias |
| 93 | + max_count = self.get_page_size_or(max_count) |
| 94 | + result = self.container.query_items( |
| 95 | + query=""" |
| 96 | + SELECT * FROM c WHERE c.tenant_id=@tenant_id AND {visibility_condition} |
| 97 | + OFFSET @offset LIMIT @max_count |
| 98 | + """.format(visibility_condition=self.create_sql_condition_for_visibility(visible_only)), |
| 99 | + parameters=[ |
| 100 | + {"name": "@tenant_id", "value": partition_key_value}, |
| 101 | + {"name": "@offset", "value": offset}, |
| 102 | + {"name": "@max_count", "value": max_count}, |
| 103 | + ], |
| 104 | + partition_key=partition_key_value, |
| 105 | + max_item_count=max_count) |
| 106 | + |
| 107 | + function_mapper = self.get_mapper_or_dict(mapper) |
| 108 | + return list(map(function_mapper, result)) |
| 109 | + |
| 110 | + def partial_update(self, id: str, changes: dict, partition_key_value: str, |
| 111 | + visible_only=True, mapper: Callable = None): |
| 112 | + item_data = self.find(id, partition_key_value, visible_only=visible_only) |
| 113 | + item_data.update(changes) |
| 114 | + return self.update(id, item_data, mapper=mapper) |
| 115 | + |
| 116 | + def update(self, id: str, item_data: dict, mapper: Callable = None): |
| 117 | + function_mapper = self.get_mapper_or_dict(mapper) |
| 118 | + return function_mapper(self.container.replace_item(id, body=item_data)) |
| 119 | + |
| 120 | + def delete(self, id: str, partition_key_value: str, mapper: Callable = None): |
| 121 | + return self.partial_update(id, { |
| 122 | + 'deleted': str(uuid.uuid4()) |
| 123 | + }, partition_key_value, visible_only=True, mapper=mapper) |
| 124 | + |
| 125 | + def check_visibility(self, item, throw_not_found_if_deleted): |
| 126 | + if throw_not_found_if_deleted and item.get('deleted') is not None: |
| 127 | + raise exceptions.CosmosResourceNotFoundError(message='Deleted item', |
| 128 | + status_code=404) |
| 129 | + |
| 130 | + return item |
| 131 | + |
| 132 | + def create_sql_condition_for_visibility(self, visible_only: bool, container_name='c') -> str: |
| 133 | + if visible_only: |
| 134 | + # We are considering that `deleted == null` is not a choice |
| 135 | + return 'NOT IS_DEFINED(%s.deleted)' % container_name |
| 136 | + return 'true' |
| 137 | + |
| 138 | + def get_mapper_or_dict(self, alternative_mapper: Callable) -> Callable: |
| 139 | + return alternative_mapper or self.mapper or dict |
| 140 | + |
| 141 | + def get_page_size_or(self, custom_page_size: int) -> int: |
| 142 | + # TODO The default value should be taken from the Azure Feature Manager |
| 143 | + # or any other repository for the settings |
| 144 | + return custom_page_size or 100 |
| 145 | + |
| 146 | + |
| 147 | +def init_app(app: Flask) -> None: |
| 148 | + global cosmos_helper |
| 149 | + cosmos_helper = CosmosDBFacade(app) |
0 commit comments