fix: BI-5901 us deserializer fix (#780)

* dict secrets deserializer fix * type annotation * missing secret key fix * swap deserialization and decryption in us manager * mypy fix * one more us deserializer fix * typo fx
datalens-tech · Jan 13, 2025 · de730b7 · de730b7
1 parent d9fd0b1
commit de730b7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 13 deletions.
diff --git a/lib/dl_core/dl_core/us_manager/us_entry_serializer.py b/lib/dl_core/dl_core/us_manager/us_entry_serializer.py
@@ -2,9 +2,7 @@
 
 import abc
 from collections import ChainMap
-import copy
 from functools import reduce
-import json
 import logging
 from typing import (
     ClassVar,
@@ -122,7 +120,6 @@ def deserialize(
         data_strict: bool = True,
     ) -> USEntry:
         # Assumed that data_pack's dicts was decoupled with initial US response
-        data_pack = copy.deepcopy(data_pack)
         secret_source_addressable = AddressableData(data_pack.secrets)
         raw_addressable = AddressableData(data_pack.data)
 
@@ -131,8 +128,7 @@ def deserialize(
         for secret_key in declared_secret_keys:
             if secret_source_addressable.contains(secret_key):
                 sec_val = secret_source_addressable.pop(secret_key)
-                sec_val_str = json.dumps(sec_val)
-                raw_addressable.set(secret_key, sec_val_str)
+                raw_addressable.set(secret_key, sec_val)
 
         if secret_source_addressable.data:
             LOGGER.warning("Undeclared secrets found")

diff --git a/lib/dl_core/dl_core/us_manager/us_manager.py b/lib/dl_core/dl_core/us_manager/us_manager.py
@@ -2,7 +2,7 @@
 
 from collections import ChainMap
 from contextlib import contextmanager
-import json
+import copy
 import logging
 from typing import (
     ClassVar,
@@ -342,6 +342,11 @@ def _entry_dict_to_obj(self, us_resp: dict, expected_type: Optional[Type[USEntry
                 secrets=us_resp.get("unversionedData"),  # type: ignore  # 2024-01-30 # TODO: Argument "secrets" to "USDataPack" has incompatible type "Any | None"; expected "dict[str, str | EncryptedData | None]"  [arg-type]
             )
 
+            data_pack = copy.deepcopy(data_pack)
+            if data_pack.secrets:
+                for key, secret in data_pack.secrets.items():
+                    data_pack.secrets[key] = self._crypto_controller.decrypt(secret)  # type: ignore # TODO: Argument 1 to "decrypt" of "CryptoController" has incompatible type "str | EncryptedData | None"; expected "EncryptedData | None"  [arg-type]
+
             entry = serializer.deserialize(
                 entry_cls,
                 data_pack,
@@ -350,11 +355,6 @@ def _entry_dict_to_obj(self, us_resp: dict, expected_type: Optional[Type[USEntry
                 common_properties=common_properties,
                 data_strict=False,
             )
-            secret_keys = serializer.get_secret_keys(entry_cls)
-            for key in secret_keys:
-                old_data = serializer.get_data_attr(entry, key)
-                decrypted_data = self._crypto_controller.decrypt(json.loads(old_data)) if old_data is not None else None
-                serializer.set_data_attr(entry, key, decrypted_data)
 
         entry.stored_in_db = True
         entry._us_resp = us_resp
@@ -433,8 +433,6 @@ def _get_entry_save_params(self, entry: USEntry) -> dict:
                 data_pack = USDataPack(data=data_dict)
 
             for key, secret in data_pack.secrets.items():
-                if isinstance(secret, dict):
-                    secret = json.dumps(secret)
                 assert secret is None or isinstance(secret, str)
                 data_pack.secrets[key] = self._crypto_controller.encrypt_with_actual_key(secret)