Spaces:
Build error
Build error
File size: 12,328 Bytes
0827183 |
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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 |
"""Implements a CosmosDB based storage provider using partitioning for a bot.
"""
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
from typing import Dict, List
from threading import Lock
import json
from azure.cosmos import documents, http_constants
from jsonpickle.pickler import Pickler
from jsonpickle.unpickler import Unpickler
import azure.cosmos.cosmos_client as cosmos_client # pylint: disable=no-name-in-module,import-error
import azure.cosmos.errors as cosmos_errors # pylint: disable=no-name-in-module,import-error
from botbuilder.core.storage import Storage
from botbuilder.azure import CosmosDbKeyEscape
class CosmosDbPartitionedConfig:
"""The class for partitioned CosmosDB configuration for the Azure Bot Framework."""
def __init__(
self,
cosmos_db_endpoint: str = None,
auth_key: str = None,
database_id: str = None,
container_id: str = None,
cosmos_client_options: dict = None,
container_throughput: int = 400,
key_suffix: str = "",
compatibility_mode: bool = False,
**kwargs,
):
"""Create the Config object.
:param cosmos_db_endpoint: The CosmosDB endpoint.
:param auth_key: The authentication key for Cosmos DB.
:param database_id: The database identifier for Cosmos DB instance.
:param container_id: The container identifier.
:param cosmos_client_options: The options for the CosmosClient. Currently only supports connection_policy and
consistency_level
:param container_throughput: The throughput set when creating the Container. Defaults to 400.
:param key_suffix: The suffix to be added to every key. The keySuffix must contain only valid ComosDb
key characters. (e.g. not: '\\', '?', '/', '#', '*')
:param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb
max key length of 255.
:return CosmosDbPartitionedConfig:
"""
self.__config_file = kwargs.get("filename")
if self.__config_file:
kwargs = json.load(open(self.__config_file))
self.cosmos_db_endpoint = cosmos_db_endpoint or kwargs.get("cosmos_db_endpoint")
self.auth_key = auth_key or kwargs.get("auth_key")
self.database_id = database_id or kwargs.get("database_id")
self.container_id = container_id or kwargs.get("container_id")
self.cosmos_client_options = cosmos_client_options or kwargs.get(
"cosmos_client_options", {}
)
self.container_throughput = container_throughput or kwargs.get(
"container_throughput"
)
self.key_suffix = key_suffix or kwargs.get("key_suffix")
self.compatibility_mode = compatibility_mode or kwargs.get("compatibility_mode")
class CosmosDbPartitionedStorage(Storage):
"""A CosmosDB based storage provider using partitioning for a bot."""
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 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
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 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 initialize(self):
if not self.container:
if not self.client:
self.client = cosmos_client.CosmosClient(
self.config.cosmos_db_endpoint,
{"masterKey": self.config.auth_key},
self.config.cosmos_client_options.get("connection_policy", None),
self.config.cosmos_client_options.get("consistency_level", None),
)
if not self.database:
with self.__lock:
try:
if not self.database:
self.database = self.client.CreateDatabase(
{"id": self.config.database_id}
)
except cosmos_errors.HTTPFailure:
self.database = self.client.ReadDatabase(
"dbs/" + self.config.database_id
)
self.__get_or_create_container()
def __get_or_create_container(self):
with self.__lock:
container_def = {
"id": self.config.container_id,
"partitionKey": {
"paths": ["/id"],
"kind": documents.PartitionKind.Hash,
},
}
try:
if not self.container:
self.container = self.client.CreateContainer(
"dbs/" + self.database["id"],
container_def,
{"offerThroughput": self.config.container_throughput},
)
except cosmos_errors.HTTPFailure as err:
if err.status_code == http_constants.StatusCodes.CONFLICT:
self.container = self.client.ReadContainer(
"dbs/" + self.database["id"] + "/colls/" + container_def["id"]
)
if "partitionKey" not in self.container:
self.compatability_mode_partition_key = True
else:
paths = self.container["partitionKey"]["paths"]
if "/partitionKey" in paths:
self.compatability_mode_partition_key = True
elif "/id" not in paths:
raise Exception(
f"Custom Partition Key Paths are not supported. {self.config.container_id} "
"has a custom Partition Key Path of {paths[0]}."
)
else:
raise err
def __get_partition_key(self, key: str) -> str:
return None if self.compatability_mode_partition_key else {"partitionKey": key}
@staticmethod
def __create_si(result) -> object:
"""Create an object from a result out of CosmosDB.
:param result:
:return object:
"""
# get the document item from the result and turn into a dict
doc = result.get("document")
# read the e_tag from Cosmos
if result.get("_etag"):
doc["e_tag"] = result["_etag"]
result_obj = Unpickler().restore(doc)
# create and return the object
return result_obj
@staticmethod
def __create_dict(store_item: object) -> Dict:
"""Return the dict of an object.
This eliminates non_magic attributes and the e_tag.
:param store_item:
:return dict:
"""
# read the content
json_dict = Pickler().flatten(store_item)
if "e_tag" in json_dict:
del json_dict["e_tag"]
# loop through attributes and write and return a dict
return json_dict
def __item_link(self, identifier) -> str:
"""Return the item link of a item in the container.
:param identifier:
:return str:
"""
return self.__container_link + "/docs/" + identifier
@property
def __container_link(self) -> str:
"""Return the container link in the database.
:param:
:return str:
"""
return self.__database_link + "/colls/" + self.config.container_id
@property
def __database_link(self) -> str:
"""Return the database link.
:return str:
"""
return "dbs/" + self.config.database_id
|