From 29ab7ce6f12a18c429cdccf939e21c0615a2965e Mon Sep 17 00:00:00 2001
From: Manank Patni <manank321@gmail.com>
Date: Wed, 12 Oct 2022 15:24:52 +0530
Subject: [PATCH] Change extra data update function

Signed-off-by: Manank Patni <manank321@gmail.com>
---
 registrydao/handlers/on_flush.py |  4 +-
 registrydao/hooks/on_restart.py  |  1 +
 registrydao/utils/extra.py       | 64 ++++++++++++++++----------------
 3 files changed, 34 insertions(+), 35 deletions(-)

diff --git a/registrydao/handlers/on_flush.py b/registrydao/handlers/on_flush.py
index f65730d..7c1a055 100644
--- a/registrydao/handlers/on_flush.py
+++ b/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()
diff --git a/registrydao/hooks/on_restart.py b/registrydao/hooks/on_restart.py
index f13ac42..54a3f0e 100644
--- a/registrydao/hooks/on_restart.py
+++ b/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)
\ No newline at end of file
diff --git a/registrydao/utils/extra.py b/registrydao/utils/extra.py
index 85c4852..a7a7e27 100644
--- a/registrydao/utils/extra.py
+++ b/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()
\ No newline at end of file
+        await dao_extra.save()
\ No newline at end of file