Fix crash while processing storage consisting of a single big map id (#182)

droserasprout · web-flow · commit e37d4cee4ba9 · 2021-12-11T18:26:32.000Z
* WIP

* WIP

* wip

* Test

* Changelog

* Changelog

* Cleanup

* Tidier tests

* Working tests with plain list storage

* whoooah

* Changelog
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,7 +2,7 @@
 
 Please use [this](https://docs.gitlab.com/ee/development/changelog.html) document as guidelines to keep a changelog.
 
-## 4.0.0-rc2 - [unreleased]
+## 4.0.0-rc2 - 2021-12-11
 
 ### ⚠ Migration
 
@@ -16,10 +16,11 @@ Please use [this](https://docs.gitlab.com/ee/development/changelog.html) documen
 
 * cli: Fixed config not being verified when invoking some commands
 * index: Fixed incorrect log messages, remove duplicate ones.
+* index: Fixed crash while processing storage of some contracts.
 
 ### Performance
 
-* index: Checks performed on each iteration of the main dipdup loop are slightly faster now.
+* index: Checks performed on each iteration of the main DipDup loop are slightly faster now.
 
 ## 4.0.0-rc1 - 2021-12-02
 
diff --git a/Makefile b/Makefile
@@ -1,5 +1,5 @@
 .ONESHELL:
-.PHONY: docs
+.PHONY: test
 .DEFAULT_GOAL: all
 
 DEV=1
diff --git a/src/demo_hic_et_nunc/hooks/on_synchronized.py b/src/demo_hic_et_nunc/hooks/on_synchronized.py
@@ -1,7 +1,7 @@
-
 from dipdup.context import HookContext
 
+
 async def on_synchronized(
     ctx: HookContext,
 ) -> None:
-    await ctx.execute_sql('on_synchronized')
+    await ctx.execute_sql('on_synchronized')
diff --git a/src/demo_quipuswap/hooks/on_synchronized.py b/src/demo_quipuswap/hooks/on_synchronized.py
@@ -1,7 +1,7 @@
-
 from dipdup.context import HookContext
 
+
 async def on_synchronized(
     ctx: HookContext,
 ) -> None:
-    await ctx.execute_sql('on_synchronized')
+    await ctx.execute_sql('on_synchronized')
diff --git a/src/demo_registrydao/hooks/on_synchronized.py b/src/demo_registrydao/hooks/on_synchronized.py
@@ -1,7 +1,7 @@
-
 from dipdup.context import HookContext
 
+
 async def on_synchronized(
     ctx: HookContext,
 ) -> None:
-    await ctx.execute_sql('on_synchronized')
+    await ctx.execute_sql('on_synchronized')
diff --git a/src/demo_tezos_domains/hooks/on_synchronized.py b/src/demo_tezos_domains/hooks/on_synchronized.py
@@ -1,7 +1,7 @@
-
 from dipdup.context import HookContext
 
+
 async def on_synchronized(
     ctx: HookContext,
 ) -> None:
-    await ctx.execute_sql('on_synchronized')
+    await ctx.execute_sql('on_synchronized')
diff --git a/src/demo_tezos_domains_big_map/hooks/on_synchronized.py b/src/demo_tezos_domains_big_map/hooks/on_synchronized.py
@@ -1,7 +1,7 @@
-
 from dipdup.context import HookContext
 
+
 async def on_synchronized(
     ctx: HookContext,
 ) -> None:
-    await ctx.execute_sql('on_synchronized')
+    await ctx.execute_sql('on_synchronized')
diff --git a/src/demo_tzbtc/hooks/on_synchronized.py b/src/demo_tzbtc/hooks/on_synchronized.py
@@ -1,7 +1,7 @@
-
 from dipdup.context import HookContext
 
+
 async def on_synchronized(
     ctx: HookContext,
 ) -> None:
-    await ctx.execute_sql('on_synchronized')
+    await ctx.execute_sql('on_synchronized')
diff --git a/src/demo_tzcolors/hooks/on_synchronized.py b/src/demo_tzcolors/hooks/on_synchronized.py
@@ -1,7 +1,7 @@
-
 from dipdup.context import HookContext
 
+
 async def on_synchronized(
     ctx: HookContext,
 ) -> None:
-    await ctx.execute_sql('on_synchronized')
+    await ctx.execute_sql('on_synchronized')
diff --git a/src/dipdup/config.py b/src/dipdup/config.py
@@ -123,7 +123,7 @@ def valid_immune_tables(cls, v):
 @dataclass
 class HTTPConfig:
     """Advanced configuration of HTTP client
-    
+
     :param cache: Whether to cache responses
     :param retry_count: Number of retries before giving up
     :param retry_sleep: Sleep time between retries
@@ -132,8 +132,9 @@ class HTTPConfig:
     :param ratelimit_period: Time period for rate limiting
     :param connection_limit: Number of simultaneous connections
     :param connection_timeout: Connection timeout
-    :param batch_size: Number of items fetched in a single request 
+    :param batch_size: Number of items fetched in a single request
     """
+
     cache: Optional[bool] = None
     retry_count: Optional[int] = None
     retry_sleep: Optional[float] = None
@@ -259,6 +260,7 @@ class CoinbaseDatasourceConfig(NameMixin):
     :param passphrase: API passphrase
     :param http: HTTP client configuration
     """
+
     kind: Literal['coinbase']
     api_key: Optional[str] = None
     secret_key: Optional[str] = None
@@ -476,6 +478,7 @@ class OperationHandlerOriginationPatternConfig(PatternConfig, StorageTypeMixin):
     :param source: Source contract alias to filter operations with
     :param similar_to: Alias of contract having the same code/signature (depending on `strict` field)
     """
+
     type: Literal['origination'] = 'origination'
     source: Optional[Union[str, ContractConfig]] = None
     similar_to: Optional[Union[str, ContractConfig]] = None
diff --git a/src/dipdup/datasources/tzkt/datasource.py b/src/dipdup/datasources/tzkt/datasource.py
@@ -677,11 +677,6 @@ async def _on_head_message(self, message: List[Dict[str, Any]]) -> None:
     @classmethod
     def convert_operation(cls, operation_json: Dict[str, Any], type_: Optional[str] = None) -> OperationData:
         """Convert raw operation message from WS/REST into dataclass"""
-        storage = operation_json.get('storage')
-        # FIXME: Plain storage, has issues in codegen: KT1CpeSQKdkhWi4pinYcseCFKmDhs5M74BkU
-        if not isinstance(storage, Dict):
-            storage = {}
-
         return OperationData(
             type=type_ or operation_json['type'],
             id=operation_json['id'],
@@ -711,19 +706,14 @@ def convert_operation(cls, operation_json: Dict[str, Any], type_: Optional[str]
             originated_contract_code_hash=operation_json['originatedContract']['codeHash']
             if operation_json.get('originatedContract')
             else None,
-            storage=storage,
-            diffs=operation_json.get('diffs'),
+            storage=operation_json.get('storage'),
+            diffs=operation_json.get('diffs') or (),
         )
 
     @classmethod
     def convert_migration_origination(cls, migration_origination_json: Dict[str, Any]) -> OperationData:
         """Convert raw migration message from REST into dataclass"""
-        storage = migration_origination_json.get('storage')
-        # FIXME: Plain storage, has issues in codegen: KT1CpeSQKdkhWi4pinYcseCFKmDhs5M74BkU
-        if not isinstance(storage, Dict):
-            storage = {}
-
-        fake_operation_data = OperationData(
+        return OperationData(
             type='origination',
             id=migration_origination_json['id'],
             level=migration_origination_json['level'],
@@ -732,8 +722,8 @@ def convert_migration_origination(cls, migration_origination_json: Dict[str, Any
             originated_contract_address=migration_origination_json['account']['address'],
             originated_contract_alias=migration_origination_json['account'].get('alias'),
             amount=migration_origination_json['balanceChange'],
-            storage=storage,
-            diffs=migration_origination_json.get('diffs'),
+            storage=migration_origination_json.get('storage'),
+            diffs=migration_origination_json.get('diffs', ()),
             status='applied',
             has_internals=False,
             hash='[none]',
@@ -742,7 +732,6 @@ def convert_migration_origination(cls, migration_origination_json: Dict[str, Any
             target_address=None,
             initiator_address=None,
         )
-        return fake_operation_data
 
     @classmethod
     def convert_big_map(cls, big_map_json: Dict[str, Any]) -> BigMapData:
diff --git a/src/dipdup/datasources/tzkt/models.py b/src/dipdup/datasources/tzkt/models.py
@@ -0,0 +1,177 @@
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import Type
+from typing import Union
+
+from pydantic.error_wrappers import ValidationError
+from typing_extensions import get_args
+from typing_extensions import get_origin
+
+from dipdup.exceptions import ConfigurationError
+from dipdup.exceptions import InvalidDataError
+from dipdup.models import OperationData
+from dipdup.models import StorageType
+
+# NOTE: typing_extensions introspection is pretty expensive, let's cache it's results
+_is_nested_dict: Dict[Type, bool] = {}
+
+
+def _unwrap(storage: Dict[str, Any]) -> Any:
+    return storage['']
+
+
+def _wrap(storage: Any) -> Dict[str, Any]:
+    return {'': storage}
+
+
+def _is_array(storage_type) -> bool:
+    if not hasattr(storage_type, '__fields__'):
+        # NOTE: Not a dataclass
+        return False
+
+    f = storage_type.__fields__
+    tzkt_array = 'key' in f and 'value' in f
+    # FIXME: fuck it (╯°□°）╯︵ ┻━┻)
+    try:
+        pydantic_array = get_origin(get_args(getattr(f.get('__root__'), 'type_', None))[1]) == list
+    except IndexError:
+        pydantic_array = False
+
+    return tzkt_array or pydantic_array
+
+
+def _apply_bigmap_diffs(
+    storage_dict: Dict[str, Any],
+    bigmap_diffs: Iterable[Dict[str, Any]],
+    bigmap_name: str,
+    is_array: bool,
+) -> None:
+    """Apply big map diffs of specific path to storage"""
+
+    for diff in bigmap_diffs:
+        bigmap_key: Union[int, str]
+
+        # NOTE: Match by bigmap name
+        if diff['path'] == bigmap_name:
+            bigmap_key = bigmap_name.split('.')[-1]
+            bigmap_dict = storage_dict
+        # NOTE: Match by index in plain list storage
+        elif diff['path'].isdigit():
+            bigmap_key = int(diff['path'])
+            bigmap_dict = _unwrap(storage_dict)
+        else:
+            continue
+
+        key = diff['content']['key']
+        if is_array:
+            bigmap_dict[bigmap_key].append(  # type: ignore
+                {
+                    'key': key,
+                    'value': diff['content']['value'],
+                }
+            )
+        else:
+            bigmap_dict[bigmap_key][key] = diff['content']['value']  # type: ignore
+
+
+def _process_storage(
+    storage_dict: Dict[str, Any],
+    storage_type: Type[StorageType],
+    bigmap_diffs: Iterable[Dict[str, Any]],
+    prefix: str = None,
+) -> None:
+    for key, field in storage_type.__fields__.items():
+        if key == '__root__':
+            continue
+
+        if field.alias:
+            key = field.alias
+
+        bigmap_name = key if prefix is None else f'{prefix}.{key}'
+
+        # NOTE: TzKT could return bigmaps as object or as array of key-value objects. We need to guess this from storage.
+        if (value := storage_dict.get(key)) is None:
+            if not field.required:
+                continue
+            raise ConfigurationError(f'Type `{storage_type.__name__}` is invalid: `{key}` field does not exists')
+
+        # NOTE: Pydantic bug? I have no idea how does it work, this workaround is just a guess.
+        # NOTE: `BaseModel.type_` returns incorrect value when annotation is Dict[str, bool], Dict[str, BaseModel], and possibly any other cases.
+        is_complex_type = field.type_ != field.outer_type_
+        is_nested_dict_model = _is_nested_dict.get(field.outer_type_)
+        if is_nested_dict_model is None:
+            try:
+                get_args(field.outer_type_)[1].__fields__
+                is_nested_dict_model = _is_nested_dict[field.outer_type_] = True
+            except Exception:
+                is_nested_dict_model = _is_nested_dict[field.outer_type_] = False
+
+        if is_complex_type and (field.type_ == bool or is_nested_dict_model):
+            field_type = field.outer_type_
+        else:
+            field_type = field.type_
+
+        if field_type not in (int, bool) and isinstance(value, int):
+            is_array = _is_array(field_type)
+            storage_dict[key] = [] if is_array else {}
+            _apply_bigmap_diffs(storage_dict, bigmap_diffs, bigmap_name, is_array)
+
+        elif hasattr(field_type, '__fields__') and isinstance(storage_dict[key], dict):
+            _process_storage(storage_dict[key], field_type, bigmap_diffs, bigmap_name)
+
+
+def _process_plain_storage(
+    storage_dict: Dict[str, Any],
+    storage_type: Type[StorageType],
+    bigmap_diffs: Iterable[Dict[str, Any]],
+) -> None:
+    # NOTE: Plain storage is either an empty model with `Extra.allow` or one with `__root__: list`
+    is_array = _is_array(storage_type)
+    unwrapped_storage = _unwrap(storage_dict)
+
+    # NOTE: Replace bigmap ids with empty structures
+    if isinstance(unwrapped_storage, int):
+        storage_dict[''] = [] if is_array else {}
+
+    elif isinstance(unwrapped_storage, list):
+        for i, item in enumerate(unwrapped_storage):
+            if isinstance(item, int):
+                unwrapped_storage[i] = [] if is_array else {}
+
+    _apply_bigmap_diffs(storage_dict, bigmap_diffs, '', is_array)
+
+
+def deserialize_storage(operation_data: OperationData, storage_type: Type[StorageType]) -> StorageType:
+    """Merge big map diffs and deserialize raw storage into typeclass"""
+    diffs = tuple(d for d in operation_data.diffs if d['action'] in ('add_key', 'update_key'))
+
+    if isinstance(operation_data.storage, dict):
+        _process_storage(
+            storage_dict=operation_data.storage,
+            storage_type=storage_type,
+            bigmap_diffs=diffs,
+            prefix=None,
+        )
+
+    elif isinstance(operation_data.storage, (int, list)):
+        # NOTE: TzKT returns empty string path for storage root, this hack allows to threat it as a regular dict storage
+        operation_data.storage = _wrap(operation_data.storage)
+        _process_plain_storage(
+            storage_dict=operation_data.storage,
+            storage_type=storage_type,
+            bigmap_diffs=diffs,
+        )
+        operation_data.storage = _unwrap(operation_data.storage)
+
+    else:
+        raise RuntimeError('Storage type must be one of `dict`, `list` or `int`')
+
+    try:
+        return storage_type.parse_obj(operation_data.storage)
+    except ValidationError as e:
+        raise InvalidDataError(
+            type_cls=storage_type,
+            data=operation_data.storage,
+            parsed_object=operation_data,
+        ) from e
diff --git a/src/dipdup/index.py b/src/dipdup/index.py
diff --git a/src/dipdup/models.py b/src/dipdup/models.py
diff --git a/tests/integration_tests/test_rollback.py b/tests/integration_tests/test_rollback.py
diff --git a/tests/test_dipdup/test_index.py b/tests/test_dipdup/test_index.py
diff --git a/tests/test_dipdup/test_models.py b/tests/test_dipdup/test_models.py
diff --git a/tests/test_dipdup/types.py b/tests/test_dipdup/types.py
