def test_should_handle_possible_collisions(self): valid_key1 = '*2atest*2a' valid_key2 = '*test*' escaped1 = CosmosDbKeyEscape.sanitize_key(valid_key1) escaped2 = CosmosDbKeyEscape.sanitize_key(valid_key2) assert escaped1 != escaped2, f'{escaped1} should be different that {escaped2}'
def test_should_create_sufficiently_different_truncated_keys_of_similar_origin( self): # create 2 very similar extra long key where the difference will definitely be trimmed off by truncate function long_key = 'x' * 300 + "1" long_key2 = 'x' * 300 + "2" fixed = CosmosDbKeyEscape.sanitize_key(long_key) fixed2 = CosmosDbKeyEscape.sanitize_key(long_key2) assert len( fixed) != fixed2, 'key truncation failed to create unique key'
def test_should_properly_truncate_keys_with_special_chars(self): # create a short key long_key = "*" * 300 fixed = CosmosDbKeyEscape.sanitize_key(long_key) assert len(fixed) <= 255, "long key with special char was truncated improperly" # create a short key short_key = "#" * 16 fixed2 = CosmosDbKeyEscape.sanitize_key(short_key) assert ( len(fixed2) <= 255 ), "short key with special char was truncated improperly"
def test_should_truncate_longer_keys(self): # create an extra long key # limit is 255 long_key = 'x' * 300 fixed = CosmosDbKeyEscape.sanitize_key(long_key) assert len(fixed) <= 255, 'long key was not properly truncated'
def __init__(self, config: CosmosDbPartitionedConfig): """Create the storage object. :param config: """ super(CosmosDbPartitionedStorage, self).__init__() self.config = config self.client = None self.database = None self.container = None self.compatability_mode_partition_key = False # Lock used for synchronizing container creation self.__lock = Lock() if config.key_suffix is None: config.key_suffix = "" if not config.key_suffix.__eq__(""): if config.compatibility_mode: raise Exception( "compatibilityMode cannot be true while using a keySuffix." ) suffix_escaped = CosmosDbKeyEscape.sanitize_key(config.key_suffix) if not suffix_escaped.__eq__(config.key_suffix): raise Exception( f"Cannot use invalid Row Key characters: {config.key_suffix} in keySuffix." )
async def delete(self, keys: List[str]): """Remove storeitems from storage. :param keys: :return: """ await self.initialize() for key in keys: escaped_key = CosmosDbKeyEscape.sanitize_key( key, self.config.key_suffix, self.config.compatibility_mode ) try: self.client.DeleteItem( document_link=self.__item_link(escaped_key), options=self.__get_partition_key(escaped_key), ) except cosmos_errors.HTTPFailure as err: if ( err.status_code == cosmos_errors.http_constants.StatusCodes.NOT_FOUND ): continue raise err except Exception as err: raise err
async def write(self, changes: Dict[str, object]): """Save storeitems to storage. :param changes: :return: """ if changes is None: raise Exception("Changes are required when writing") if not changes: return await self.initialize() for (key, change) in changes.items(): e_tag = None if isinstance(change, dict): e_tag = change.get("e_tag", None) elif hasattr(change, "e_tag"): e_tag = change.e_tag doc = { "id": CosmosDbKeyEscape.sanitize_key( key, self.config.key_suffix, self.config.compatibility_mode ), "realId": key, "document": self.__create_dict(change), } if e_tag == "": raise Exception("cosmosdb_storage.write(): etag missing") access_condition = { "accessCondition": {"type": "IfMatch", "condition": e_tag} } options = ( access_condition if e_tag != "*" and e_tag and e_tag != "" else None ) try: self.client.UpsertItem( database_or_Container_link=self.__container_link, document=doc, options=options, ) except cosmos_errors.HTTPFailure as err: raise err except Exception as err: raise err
async def read(self, keys: List[str]) -> Dict[str, object]: """Read storeitems from storage. :param keys: :return dict: """ if not keys: raise Exception("Keys are required when reading") await self.initialize() store_items = {} for key in keys: try: escaped_key = CosmosDbKeyEscape.sanitize_key( key, self.config.key_suffix, self.config.compatibility_mode ) read_item_response = self.client.ReadItem( self.__item_link(escaped_key), self.__get_partition_key(escaped_key) ) document_store_item = read_item_response if document_store_item: store_items[document_store_item["realId"]] = self.__create_si( document_store_item ) # When an item is not found a CosmosException is thrown, but we want to # return an empty collection so in this instance we catch and do not rethrow. # Throw for any other exception. except cosmos_errors.HTTPFailure as err: if ( err.status_code == cosmos_errors.http_constants.StatusCodes.NOT_FOUND ): continue raise err except Exception as err: raise err return store_items
def test_should_escape_illegal_characters_case_3(self): # Ascii code of "\" is "5c" sanitized_key = CosmosDbKeyEscape.sanitize_key('\\test\\') assert sanitized_key == '*92test*92'
def test_should_not_truncate_short_key(self): # create a short key short_key = 'x' * 16 fixed2 = CosmosDbKeyEscape.sanitize_key(short_key) assert len(fixed2) == 16, 'short key was truncated improperly'
def test_should_not_change_a_valid_key(self): valid_key = 'Abc12345' sanitized_key = CosmosDbKeyEscape.sanitize_key(valid_key) assert valid_key == sanitized_key, f'{valid_key} should be equal to {sanitized_key}'
def test_should_not_change_a_valid_key(self): valid_key = "Abc12345" sanitized_key = CosmosDbKeyEscape.sanitize_key(valid_key) assert (valid_key == sanitized_key ), f"{valid_key} should be equal to {sanitized_key}"
def test_should_escape_illegal_characters_compound_key(self): # Check a compound key compoundsanitized_key = CosmosDbKeyEscape.sanitize_key('?#/') assert compoundsanitized_key, '*3f*23*2f'
def test_should_escape_illegal_characters_case_5(self): # Ascii code of "*" is "2a". sanitized_key = CosmosDbKeyEscape.sanitize_key('*test*') assert sanitized_key == '*42test*42'
def test_should_escape_illegal_characters_case_4(self): # Ascii code of "#" is "23" sanitized_key = CosmosDbKeyEscape.sanitize_key('#test#') assert sanitized_key == '*35test*35'
def test_should_escape_illegal_characters_case_1(self): # Ascii code of "?" is "3f" sanitized_key = CosmosDbKeyEscape.sanitize_key('?test?') assert sanitized_key == '*63test*63'
def test_should_escape_illegal_characters_case_2(self): # Ascii code of "/" is "2f" sanitized_key = CosmosDbKeyEscape.sanitize_key('/test/') assert sanitized_key == '*47test*47'