|
10 | 10 | from typing import Dict, List |
11 | 11 | from threading import Semaphore |
12 | 12 | import json |
| 13 | +from jsonpickle.pickler import Pickler |
| 14 | +from jsonpickle.unpickler import Unpickler |
13 | 15 | import azure.cosmos.cosmos_client as cosmos_client # pylint: disable=no-name-in-module,import-error |
14 | 16 | import azure.cosmos.errors as cosmos_errors # pylint: disable=no-name-in-module,import-error |
15 | | -from botbuilder.core.storage import Storage, StoreItem |
| 17 | +from botbuilder.core.storage import Storage |
16 | 18 |
|
17 | 19 |
|
18 | 20 | class CosmosDbConfig: |
@@ -144,15 +146,15 @@ async def read(self, keys: List[str]) -> Dict[str, object]: |
144 | 146 | results = list( |
145 | 147 | self.client.QueryItems(self.__container_link, query, options) |
146 | 148 | ) |
147 | | - # return a dict with a key and a StoreItem |
| 149 | + # return a dict with a key and an object |
148 | 150 | return {r.get("realId"): self.__create_si(r) for r in results} |
149 | 151 |
|
150 | 152 | # No keys passed in, no result to return. |
151 | 153 | return {} |
152 | 154 | except TypeError as error: |
153 | 155 | raise error |
154 | 156 |
|
155 | | - async def write(self, changes: Dict[str, StoreItem]): |
| 157 | + async def write(self, changes: Dict[str, object]): |
156 | 158 | """Save storeitems to storage. |
157 | 159 |
|
158 | 160 | :param changes: |
@@ -224,36 +226,38 @@ async def delete(self, keys: List[str]): |
224 | 226 | except TypeError as error: |
225 | 227 | raise error |
226 | 228 |
|
227 | | - def __create_si(self, result) -> StoreItem: |
228 | | - """Create a StoreItem from a result out of CosmosDB. |
| 229 | + def __create_si(self, result) -> object: |
| 230 | + """Create an object from a result out of CosmosDB. |
229 | 231 |
|
230 | 232 | :param result: |
231 | | - :return StoreItem: |
| 233 | + :return object: |
232 | 234 | """ |
233 | 235 | # get the document item from the result and turn into a dict |
234 | 236 | doc = result.get("document") |
235 | | - # readd the e_tag from Cosmos |
| 237 | + # read the e_tag from Cosmos |
236 | 238 | if result.get("_etag"): |
237 | 239 | doc["e_tag"] = result["_etag"] |
238 | | - # create and return the StoreItem |
239 | | - return StoreItem(**doc) |
240 | 240 |
|
241 | | - def __create_dict(self, store_item: StoreItem) -> Dict: |
242 | | - """Return the dict of a StoreItem. |
| 241 | + result_obj = Unpickler().restore(doc) |
| 242 | + |
| 243 | + # create and return the object |
| 244 | + return result_obj |
| 245 | + |
| 246 | + def __create_dict(self, store_item: object) -> Dict: |
| 247 | + """Return the dict of an object. |
243 | 248 |
|
244 | 249 | This eliminates non_magic attributes and the e_tag. |
245 | 250 |
|
246 | 251 | :param store_item: |
247 | 252 | :return dict: |
248 | 253 | """ |
249 | 254 | # read the content |
250 | | - non_magic_attr = [ |
251 | | - attr |
252 | | - for attr in dir(store_item) |
253 | | - if not attr.startswith("_") or attr.__eq__("e_tag") |
254 | | - ] |
| 255 | + json_dict = Pickler().flatten(store_item) |
| 256 | + if "e_tag" in json_dict: |
| 257 | + del json_dict["e_tag"] |
| 258 | + |
255 | 259 | # loop through attributes and write and return a dict |
256 | | - return {attr: getattr(store_item, attr) for attr in non_magic_attr} |
| 260 | + return json_dict |
257 | 261 |
|
258 | 262 | def __item_link(self, identifier) -> str: |
259 | 263 | """Return the item link of a item in the container. |
|
0 commit comments