Add staked field to Vote Model to support unstake from all proposals (#…

…24) Signed-off-by: Manank Patni <[email protected]> Signed-off-by: Manank Patni <[email protected]>
dOrgTech · Jan 25, 2023 · e3c1ca1 · e3c1ca1
1 parent b184317
commit e3c1ca1
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 11 deletions.
diff --git a/registrydao/handlers/on_unstake_vote.py b/registrydao/handlers/on_unstake_vote.py
@@ -16,7 +16,22 @@ async def on_unstake_vote(
 ) -> None:
     try:
         dao_address = unstake_vote.data.target_address
+        proposal_key = unstake_vote.parameter.__root__[0]
+
+        dao = await models.DAO.get(address=dao_address)
+        proposal = await models.Proposal.get(key=proposal_key, dao=dao)
+
+        voter = await models.Holder.get_or_create(address=unstake_vote.data.sender_address)
+
         await update_ledger(dao_address, unstake_vote.data.diffs)
+
+        votes = await models.Vote.filter(proposal=proposal, voter=voter[0])
+
+        for vote in votes:
+            vote.staked = False
+            await vote.save()
+
     except Exception as e:
-        print("Error in on_unstake_vote: " + str(unstake_vote.data.target_address))
-        print(e)
+        print("Error in on_unstake_vote: " +
+              str(unstake_vote.data.target_address))
+        print(e)
diff --git a/registrydao/handlers/on_vote.py b/registrydao/handlers/on_vote.py
@@ -33,16 +33,17 @@ async def on_vote(
             support=support,
             voter=voter[0],
             defaults={
-                'amount':amount
-            }
+                'amount': amount
+            },
+            staked=True
         )
 
         if support:
             proposal.upvotes = float(proposal.upvotes) + float(amount)
         else:
             proposal.downvotes = float(proposal.downvotes) + float(amount)
-        
+
         await proposal.save()
     except Exception as e:
         print("Error in on_vote: " + str(vote.data.target_address))
-        print(e)
+        print(e)
diff --git a/registrydao/models.py b/registrydao/models.py
@@ -87,6 +87,7 @@ class RegistryExtra(Model):
     class Meta:
         table = 'registry_extra'
 
+
 class TreasuryExtra(Model):
     id = fields.IntField(pk=True)
     dao: fields.ForeignKeyRelation[DAOType] = fields.ForeignKeyField(
@@ -102,6 +103,7 @@ class TreasuryExtra(Model):
     class Meta:
         table = 'treasury_extra'
 
+
 class LambdaExtra(Model):
     id = fields.IntField(pk=True)
     dao: fields.ForeignKeyRelation[DAOType] = fields.ForeignKeyField(
@@ -120,6 +122,7 @@ class LambdaExtra(Model):
     class Meta:
         table = 'lambda_extra'
 
+
 class Holder(Model):
     id = fields.IntField(pk=True)
     address = fields.CharField(36, unique=True)
@@ -148,6 +151,7 @@ class Meta:
         table = 'ledger'
         unique_together = (("dao", "holder"),)
 
+
 class ProposalStatus(Model):
     id = fields.IntField(pk=True)
     description = fields.CharField(36)
@@ -162,8 +166,8 @@ class Proposal(Model):
     dao: fields.ForeignKeyRelation[DAO] = fields.ForeignKeyField(
         "models.DAO"
     )
-    hash=fields.CharField(128)
-    key=fields.CharField(128)
+    hash = fields.CharField(128)
+    key = fields.CharField(128)
     upvotes = fields.DecimalField(54, 18)
     downvotes = fields.DecimalField(54, 18)
     start_level = fields.IntField()
@@ -172,9 +176,9 @@ class Proposal(Model):
     proposer: fields.ForeignKeyRelation[Holder] = fields.ForeignKeyField(
         "models.Holder"
     )
-    voting_stage_num=fields.CharField(50)
-    proposer_frozen_token=fields.CharField(50)
-    quorum_threshold=fields.DecimalField(54, 18)
+    voting_stage_num = fields.CharField(50)
+    proposer_frozen_token = fields.CharField(50)
+    quorum_threshold = fields.DecimalField(54, 18)
     votes: fields.ReverseRelation["Vote"]
     status_updates: fields.ReverseRelation["ProposalStatusUpdates"]
 
@@ -188,6 +192,7 @@ class Vote(Model):
         "models.Proposal"
     )
     amount = fields.DecimalField(54, 18)
+    staked = fields.BooleanField()
     support = fields.BooleanField()
     voter: fields.ForeignKeyRelation[Holder] = fields.ForeignKeyField(
         "models.Holder"
@@ -196,6 +201,7 @@ class Vote(Model):
     class Meta:
         table = 'votes'
 
+
 class Transfer(Model):
     id = fields.IntField(pk=True)
     timestamp = fields.DatetimeField()
@@ -208,6 +214,7 @@ class Transfer(Model):
     from_address = fields.CharField(36)
     hash = fields.CharField(128)
 
+
 class ProposalStatusUpdates(Model):
     id = fields.IntField(pk=True)
     timestamp = fields.DatetimeField()