fix calcualte balance

Author: olegsta

app/__init__.py

@@ -22,15 +22,6 @@
 sqlite3.register_adapter(decimal.Decimal, lambda x: str(x))
 sqlite3.register_converter("DECTEXT", lambda x: decimal.Decimal(x.decode()))
 
-def init_settings_table(db):
-    create_table_sql = """
-    CREATE TABLE IF NOT EXISTS settings (
-        name TEXT PRIMARY KEY,
-        value TEXT
-    );
-    """
-    db.query_db2(create_table_sql)
-
 def create_app():
 
     from flask.config import Config
@@ -55,7 +46,6 @@ def __dir__(self):
     from . import db
 
     db.init_app(app)
-    init_settings_table(db)
     key_type = "only_read" if app.config.READ_MODE else "onetime"
 
     rows = db.query_db2(f'SELECT public FROM keys WHERE type = "{key_type}"')

app/block_scanner.py

@@ -93,6 +93,12 @@ def main_account(self):
             "public"
         ]
 
+    @functools.cached_property
+    def read_mode_account(self):
+        return query_db2('select * from keys where type = "only_read_finished" ', one=True)[
+            "public"
+        ]
+
     def get_last_seen_block_num(self) -> int:
         row = query_db2(
             'SELECT value FROM settings WHERE name = "last_seen_block_num"', one=True

app/tasks.py

@@ -12,6 +12,7 @@
 from celery.schedules import crontab
 from tronpy.keys import PrivateKey
 from tronpy.tron import current_timestamp
+from tronpy import Tron
 import tronpy.exceptions
 import requests
 from sqlmodel import Session, select
@@ -263,11 +264,10 @@ def precision_of(symbol):
                 .functions.decimals()
             )
 
-        query = 'SELECT public FROM keys WHERE type = ?'
-        accounts_onetime = [row["public"] for row in query_db(query, ("onetime",))]
-        accounts_read_mode = [row["public"] for row in query_db(query, ("only_read",))]
-        accounts = accounts_read_mode if config.READ_MODE else accounts_onetime
-
+        accounts = [
+            row["public"]
+            for row in query_db('SELECT public FROM keys WHERE type = "only_read_finished"')
+        ]
         for index, account in enumerate(accounts, start=1):
             try:
                 #
@@ -381,9 +381,87 @@ def precision_of(symbol):
                 stats["exception_num"] += 1
     return stats
 
+@celery.task(bind=True)
+@skip_if_running
+def scan_ballance(self, *args, **kwargs):
+    """
+    Scans accounts balances (TRX and TRC20),
+    """
+    from .db import engine
+    from .models import Balance
+    from tronpy import Tron
+    client = Tron()
+    with Session(engine) as session:
+        stats = {
+            "balances": collections.defaultdict(Decimal),
+            "exception_num": 0,
+        }
+
+        accounts = [
+            row["public"]
+            for row in query_db('SELECT public FROM keys WHERE type = "only_read_finished"')
+        ]
+
+        for index, account in enumerate(accounts, start=1):
+            try:
+                # === TRX BALANCE ===
+                trx_balance = client.get_account_balance(account)
+                stats["balances"]["TRX"] += trx_balance
+                logger.debug(f"[TRX] {account} -> {trx_balance} TRX")
+
+                if config.READ_MODE:
+                    acc_balance = session.exec(
+                        select(Balance).where(Balance.account == account, Balance.symbol == "TRX")
+                    ).first()
+                    if acc_balance:
+                        acc_balance.balance = trx_balance
+                    else:
+                        acc_balance = Balance(
+                            account=account,
+                            symbol="TRX",
+                            balance=trx_balance,
+                        )
+                    session.add(acc_balance)
+                    session.commit()
+
+                # === TRC20 TOKENS ===
+                for token in config.get_tokens():
+                    symbol = token.symbol
+                    contract_address = token.contract_address
+                    contract = client.get_contract(contract_address)
+
+                    # balanceOf returns balance in raw token units
+                    raw_balance = contract.functions.balanceOf(account)
+                    trc20_balance = Decimal(raw_balance) / (10 ** token.decimal)
+
+                    stats["balances"][symbol] += trc20_balance
+                    logger.debug(f"[{symbol}] {account} -> {trc20_balance} tokens")
+
+                    if config.READ_MODE:
+                        acc_balance = session.exec(
+                            select(Balance).where(Balance.account == account, Balance.symbol == symbol)
+                        ).first()
+                        if acc_balance:
+                            acc_balance.balance = trc20_balance
+                        else:
+                            acc_balance = Balance(
+                                account=account,
+                                symbol=symbol,
+                                balance=trc20_balance,
+                            )
+                        session.add(acc_balance)
+                        session.commit()
+
+            except Exception as e:
+                stats["exception_num"] += 1
+                logger.warning(f"[ERROR] {account} scan error: {e}")
+
+        return stats
 
 @celery.on_after_configure.connect
 def setup_periodic_tasks(sender, **kwargs):
+    if config.READ_MODE:
+        sender.add_periodic_task(config.BALANCES_RESCAN_PERIOD, scan_ballance.s())
     if config.EXTERNAL_DRAIN_CONFIG:
         from .custom.aml.tasks import sweep_accounts, recheck_transactions
 

app/wallet.py

@@ -3,8 +3,12 @@
 import tronpy.exceptions
 from tronpy.keys import PrivateKey
 
+from sqlalchemy import func
+from sqlmodel import Session, select
+from .models import Balance
 from .config import config
 from .db import query_db2
+from .db import engine
 from .logging import logger
 from .connection_manager import ConnectionManager
 from .wallet_encryption import wallet_encryption
@@ -38,7 +42,14 @@ def get_contract(self, contract_address=None):
 
     @property
     def balance(self):
-        return self.balance_of(self.main_account["public"])
+        if config.READ_MODE:
+            with Session(engine) as session:
+                total_sum = session.exec(select(func.sum(Balance.balance))).one()
+                raw_balance = total_sum if total_sum is not None else Decimal(0)
+                overall_balance = float(round(raw_balance, 6))
+            return overall_balance
+        else:
+            return self.balance_of(self.main_account["public"])
 
     def balance_of(self, address):
         if self.symbol == "TRX":

run.py

@@ -11,11 +11,6 @@
 #
 # Refresh best server
 #
-#
-# Flask
-#
-
-server = app.create_app()
 
 refresh_best_server_thread = threading.Thread(
     daemon=True,
@@ -24,6 +19,11 @@
 )
 refresh_best_server_thread.start()
 
+#
+# Flask
+#
+
+server = app.create_app()
 
 #
 # Block scanner