Merge pull request #472 from BoostryJP/feature/#471

feat: accumulate locked balance in TokenHolder collection batch

app/model/db/token_holders.py

@@ -56,9 +56,12 @@ class TokenHolder(Base):
     account_address = Column(String(42), primary_key=True)
     # Amounts(including balance/pending_transfer/exchange_balance/exchange_commitment)
     hold_balance = Column(BigInteger)
+    # Amounts(locked)
+    locked_balance = Column(BigInteger)
 
     def json(self):
         return {
             "account_address": self.account_address,
-            "hold_balance": self.hold_balance
+            "hold_balance": self.hold_balance,
+            "locked_balance": self.locked_balance
         }

app/model/schema/token_holders.py

@@ -86,11 +86,18 @@ class ListAllTokenHolderCollectionsResponse(BaseModel):
     collections: List[RetrieveTokenHolderCollectionResponse]
 
 
+class TokenHoldersCollectionHolder(BaseModel):
+    account_address: str = Field(description="Account address of token holder.")
+    hold_balance: int = Field(description="Amount of balance."
+                                          "This includes balance/pending_transfer/exchange_balance/exchange_commitment.")
+    locked_balance: int = Field(description="Amount of locked balance.")
+
+
 class RetrieveTokenHoldersListResponse(BaseModel):
     """Retrieve Token Holders List schema (RESPONSE)"""
 
     status: TokenHolderBatchStatus
-    holders: List[Dict[str, Union[int, str]]]
+    holders: List[TokenHoldersCollectionHolder]
 
     class Config:
         schema_extra = {
@@ -99,7 +106,8 @@ class Config:
                 "holders": [
                     {
                         "account_address": "0x85a8b8887a4bD76859751b10C8aC8EC5f3aA1bDB",
-                        "hold_balance": 30000
+                        "hold_balance": 30000,
+                        "locked_balance": 0
                     }
                 ],
             }

batch/indexer_token_holders.py

@@ -53,14 +53,16 @@ class BalanceBook:
         def __init__(self):
             self.pages = {}
 
-        def store(self, account_address: str, amount: int = 0):
+        def store(self, account_address: str, amount: int = 0, locked: int = 0):
             if account_address not in self.pages:
                 token_holder = TokenHolder()
-                token_holder.hold_balance = amount
+                token_holder.hold_balance = 0 + amount
                 token_holder.account_address = account_address
+                token_holder.locked_balance = 0 + locked
                 self.pages[account_address] = token_holder
             else:
                 self.pages[account_address].hold_balance += amount
+                self.pages[account_address].locked_balance += locked
 
     target: Optional[TokenHoldersList]
     balance_book: BalanceBook
@@ -186,6 +188,7 @@ def __process_all(self, db_session: Session, block_from: int, block_to: int):
         self.__process_transfer(block_from, block_to)
         self.__process_issue(block_from, block_to)
         self.__process_redeem(block_from, block_to)
+        self.__process_lock(block_from, block_to)
         self.__process_unlock(block_from, block_to)
 
         self.__save_holders(
@@ -332,6 +335,33 @@ def __process_redeem(self, block_from: int, block_to: int):
         except Exception:
             raise
 
+    def __process_lock(self, block_from: int, block_to: int):
+        """Process Lock Event
+
+        - The process of updating Hold-Balance data by capturing the following events
+        - `Lock` event on Token contracts
+
+        :param block_from: From block
+        :param block_to: To block
+        :return: None
+        """
+        try:
+            # Get "Lock" events from token contract
+            events = ContractUtils.get_event_logs(
+                contract=self.token_contract,
+                event="Lock",
+                block_from=block_from,
+                block_to=block_to
+            )
+            for event in events:
+                args = event["args"]
+                account_address = args.get("accountAddress", ZERO_ADDRESS)
+                amount = args.get("value")
+                if amount is not None and amount <= sys.maxsize:
+                    self.balance_book.store(account_address=account_address, amount=-amount, locked=+amount)
+        except Exception:
+            raise
+
     def __process_unlock(self, block_from: int, block_to: int):
         """Process Unlock Event
 
@@ -354,9 +384,10 @@ def __process_unlock(self, block_from: int, block_to: int):
                 args = event["args"]
                 account_address = args.get("accountAddress", ZERO_ADDRESS)
                 recipient_address = args.get("recipientAddress", ZERO_ADDRESS)
-                amount = args.get("value", ZERO_ADDRESS)
-                self.balance_book.store(account_address=account_address, amount=-amount)
-                self.balance_book.store(account_address=recipient_address, amount=+amount)
+                amount = args.get("value")
+                if amount is not None and amount <= sys.maxsize:
+                    self.balance_book.store(account_address=account_address, locked=-amount)
+                    self.balance_book.store(account_address=recipient_address, amount=+amount)
         except Exception:
             raise
 
@@ -370,6 +401,7 @@ def __save_holders(
     ):
         for account_address, page in zip(balance_book.pages.keys(), balance_book.pages.values()):
             if page.account_address == token_owner_address:
+                # Skip storing data for token owner
                 continue
             token_holder: TokenHolder = (
                 db_session.query(TokenHolder)
@@ -378,10 +410,10 @@ def __save_holders(
                 .first()
             )
             if token_holder is not None:
-                if page.hold_balance is not None:
-                    token_holder.hold_balance = page.hold_balance
+                token_holder.hold_balance = page.hold_balance
+                token_holder.locked_balance = page.locked_balance
                 db_session.merge(token_holder)
-            elif page.hold_balance is not None and page.hold_balance > 0:
+            elif page.hold_balance > 0 or page.locked_balance > 0:
                 LOG.debug(f"Collection record created : token_address={token_address}, account_address={account_address}")
                 page.holder_list_id = holder_list_id
                 db_session.add(page)

migrations/versions/7d09371c365e_v23_3_0_feature_471.py

@@ -0,0 +1,37 @@
+"""v23.3.0_feature_471
+
+Revision ID: 7d09371c365e
+Revises: 3a0cc5e324b4
+Create Date: 2023-01-31 15:08:30.238839
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+from app.database import get_db_schema
+
+# revision identifiers, used by Alembic.
+revision = '7d09371c365e'
+down_revision = '3a0cc5e324b4'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    # ### end Alembic commands ###
+
+    # NOTE: Manually modified
+    #   1. Added server_default arg to migrate existing record.
+    #   2. Added alter column operation to remove server default setting.
+    op.add_column('token_holder', sa.Column('locked_balance', sa.BigInteger(), server_default='0', nullable=True), schema=get_db_schema())
+    op.alter_column('token_holder', 'locked_balance',
+                    existing_type=sa.BigInteger(),
+                    existing_nullable=True, server_default=None, schema=get_db_schema())
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('token_holder', 'locked_balance', schema=get_db_schema())
+    # ### end Alembic commands ###

tests/test_app_routers_token_holders_{token_address}_collection_{collection_id}_GET.py

@@ -118,6 +118,7 @@ def test_normal_2(self, client, db):
             _token_holder.holder_list_id = _token_holders_list.id
             _token_holder.account_address = config_eth_account(user)["address"]
             _token_holder.hold_balance = 10000 * (i+1)
+            _token_holder.locked_balance = 20000 * (i+1)
             db.add(_token_holder)
             holders.append(_token_holder.json())
 

tests/test_batch_indexer_token_holders.py

@@ -295,7 +295,7 @@ def test_normal_1(
         # user1: 16000 user2: 44000
 
         STContractUtils.lock(token_contract.address, user_address_1, user_pk_1, [issuer_address, 3000, ""])
-        # user1: 16000 user2: 44000
+        # user1: (hold: 13000, locked: 3000) user2: 44000
 
         # Insert collection record with above token and current block number
         list_id = str(uuid.uuid4())
@@ -323,8 +323,10 @@ def test_normal_1(
             .first()
         )
 
-        assert user1_record.hold_balance == 16000
+        assert user1_record.hold_balance == 13000
+        assert user1_record.locked_balance == 3000
         assert user2_record.hold_balance == 44000
+        assert user2_record.locked_balance == 0
 
         assert len(list(db.query(TokenHolder).filter(TokenHolder.holder_list_id == _token_holders_list.id))) == 2
 
@@ -449,7 +451,9 @@ def test_normal_2(
         )
 
         assert user1_record.hold_balance == 10000
+        assert user1_record.locked_balance == 0
         assert user2_record.hold_balance == 20000
+        assert user2_record.locked_balance == 0
 
         assert len(list(db.query(TokenHolder).filter(TokenHolder.holder_list_id == _token_holders_list.id))) == 2
 
@@ -674,7 +678,7 @@ def test_normal_4(
         # user1: 6000 user2: 44000
 
         STContractUtils.lock(token_contract.address, user_address_1, user_pk_1, [issuer_address, 3000, ""])
-        # user1: 16000 user2: 44000
+        # user1: (hold: 3000, locked: 3000) user2: 44000
 
         # Insert collection record with above token and current block number
         list_id = str(uuid.uuid4())
@@ -702,8 +706,10 @@ def test_normal_4(
             .first()
         )
 
-        assert user1_record.hold_balance == 6000
+        assert user1_record.hold_balance == 3000
+        assert user1_record.locked_balance == 3000
         assert user2_record.hold_balance == 44000
+        assert user2_record.locked_balance == 0
 
         assert len(list(db.query(TokenHolder).filter(TokenHolder.holder_list_id == _token_holders_list.id))) == 2
 
@@ -828,7 +834,9 @@ def test_normal_5(
         )
 
         assert user1_record.hold_balance == 10000
+        assert user1_record.locked_balance == 0
         assert user2_record.hold_balance == 20000
+        assert user2_record.locked_balance == 0
 
         assert len(list(db.query(TokenHolder).filter(TokenHolder.holder_list_id == _token_holders_list.id))) == 2