From e3c1ca1120adaaab7a1450f6caba76d176bd5ca1 Mon Sep 17 00:00:00 2001
From: Manank Patni <manank321@gmail.com>
Date: Wed, 25 Jan 2023 17:49:24 +0530
Subject: [PATCH] Add staked field to Vote Model to support unstake from all
 proposals (#24)

Signed-off-by: Manank Patni <manank321@gmail.com>

Signed-off-by: Manank Patni <manank321@gmail.com>
---
 registrydao/handlers/on_unstake_vote.py | 19 +++++++++++++++++--
 registrydao/handlers/on_vote.py         |  9 +++++----
 registrydao/models.py                   | 17 ++++++++++++-----
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/registrydao/handlers/on_unstake_vote.py b/registrydao/handlers/on_unstake_vote.py
index c965e0b..c766d9a 100644
--- 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)
\ No newline at end of file
+        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
index 8391f5e..0b256bc 100644
--- 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)
\ No newline at end of file
+        print(e)
diff --git a/registrydao/models.py b/registrydao/models.py
index f0d1803..ffebb8b 100644
--- 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()