Change extra data update function

Signed-off-by: Manank Patni <[email protected]>

registrydao/handlers/on_flush.py

@@ -17,12 +17,12 @@ async def on_flush(
     flush: Transaction[FlushParameter, RegistryStorage],
 ) -> None:
     try:
-        non_flushed_or_executed_keys = list(map(extract_key, flush.data.storage['proposal_key_list_sort_by_level']))
+        non_flushed_or_executed_keys = list(map(extract_key, flush.data.storage['proposals']))
         dao_address = flush.data.target_address
         dao = await models.DAO.get(address=dao_address)
 
         await update_ledger(dao_address, flush.data.diffs)
-        await update_extra(dao_address, flush.data.diffs)
+        await update_extra(dao_address, flush.data.storage['extra']['handler_storage'])
 
         dao.guardian = flush.data.storage["guardian"]
         await dao.save()

registrydao/hooks/on_restart.py

@@ -12,6 +12,7 @@ async def on_restart(
         await models.ProposalStatus.get_or_create(description="rejected_and_flushed")
         await models.DAOType.get_or_create(name='treasury')
         await models.DAOType.get_or_create(name='registry')
+        await models.DAOType.get_or_create(name='lambda')
     except Exception as e:
         print("Error in on_restart")
         print(e)

registrydao/utils/extra.py

@@ -1,38 +1,36 @@
-from typing import List
+from typing import Dict, List
 import registrydao.models as models
 
-async def update_extra(dao_address: str, diffs: List):
+async def update_extra(dao_address: str, storage: Dict):
 
-    if diffs == None:
-        return
+    for key in storage:
+        dao = await models.DAO.get(address=dao_address).prefetch_related("type")
 
-    for diff in diffs:
-        if diff['path'] == 'extra':
-            dao = await models.DAO.get(address=dao_address).prefetch_related("type")
+        if dao.type.name == 'treasury':
+            dao_extra = await models.TreasuryExtra.get(dao=dao)
+        elif dao.type.name == 'registry':
+            dao_extra = await models.RegistryExtra.get(dao=dao)
+        elif dao.type.name == 'lambda':
+            dao_extra = await models.LambdaExtra.get(dao=dao)
+
+        diff_key = key
+        diff_value = storage[key]
+
+        if diff_key == 'frozen_extra_value':
+            dao_extra.frozen_extra_value = diff_value
+        if diff_key == 'frozen_scale_value':
+            dao_extra.frozen_scale_value = diff_value
+        if diff_key == 'slash_division_value':
+            dao_extra.slash_division_value = diff_value
+        if diff_key == 'min_xtz_amount':
+            dao_extra.min_xtz_amount = diff_value
+        if diff_key == 'max_xtz_amount':
+            dao_extra.max_xtz_amount = diff_value
+        if diff_key == 'slash_scale_value':
+            dao_extra.slash_scale_value = diff_value
+        if diff_key == 'registry':
+            dao_extra.registry = diff_value
+        if diff_key == 'registry_affected':
+            dao_extra.registry_affected = diff_value
 
-            if dao.type.name == 'treasury':
-                dao_extra = await models.TreasuryExtra.get(dao=dao)
-            elif dao.type.name == 'registry':
-                dao_extra = await models.RegistryExtra.get(dao=dao)
-
-            diff_key = diff['content']['key']
-            diff_value = diff['content']['value']
-
-            if diff_key == 'frozen_extra_value':
-                dao_extra.frozen_extra_value = diff_value
-            if diff_key == 'frozen_scale_value':
-                dao_extra.frozen_scale_value = diff_value
-            if diff_key == 'slash_division_value':
-                dao_extra.slash_division_value = diff_value
-            if diff_key == 'min_xtz_amount':
-                dao_extra.min_xtz_amount = diff_value
-            if diff_key == 'max_xtz_amount':
-                dao_extra.max_xtz_amount = diff_value
-            if diff_key == 'slash_scale_value':
-                dao_extra.slash_scale_value = diff_value
-            if diff_key == 'registry':
-                dao_extra.registry = diff_value
-            if diff_key == 'registry_affected':
-                dao_extra.registry_affected = diff_value
-
-            await dao_extra.save()
+        await dao_extra.save()