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'
