feat: balance and reward

database/__init__.py

@@ -1 +1,4 @@
-from database.models import *
+from flask_sqlalchemy import SQLAlchemy
+
+db = SQLAlchemy()
+from .models import *

database/enums.py

@@ -8,3 +8,9 @@ class AssemblyState(IntEnum):
     ALL_PRODUCTS_ASSEMBLED = 2,
     CONFIRMED = 3,
     ENDED = 4
+
+
+@unique
+class BalanceTransactionType(IntEnum):
+    TOP_UP = 0
+    WITHDRAW = 1

database/models/__init__.py

@@ -0,0 +1,2 @@
+from .basic import *
+from .balance import *

database/models/balance.py

@@ -0,0 +1,40 @@
+from sqlalchemy import Connection, func
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import Mapper
+
+import database
+from database import db
+
+
+class BalanceTransaction(db.Model):
+    __tablename__ = 'balance_transactions'
+
+    id = db.Column(db.Integer, primary_key=True, comment='ID транзакции')
+    type = db.Column(db.Integer, nullable=False, comment='Тип транзакции')
+    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
+    user = db.relationship('User', backref='transactions')
+    amount = db.Column(db.Double, nullable=False, comment='Сумма транзакции')
+    description = db.Column(db.Text, nullable=False, comment='Описание')
+    json_data = db.Column(JSONB, comment='JSON данные')
+    created_at = db.Column(db.DateTime, nullable=False, comment='Дата и время создания транзакции')
+
+
+def recalculate_user_balance(user_id: int):
+    user_balance = (database.BalanceTransaction.query.filter_by(user_id=user_id)
+                    .group_by(database.BalanceTransaction.user_id)
+                    .with_entities(func.sum(database.BalanceTransaction.amount))
+                    .scalar()) or 0
+
+    (db.session.query(database.User)
+     .filter_by(id=user_id)
+     .update({'balance': user_balance}, synchronize_session='fetch'))
+
+
+@db.event.listens_for(BalanceTransaction, 'after_insert')
+def after_balance_transaction_insert(mapper: Mapper, connection: Connection, target: BalanceTransaction):
+    recalculate_user_balance(target.user_id)
+
+
+@db.event.listens_for(BalanceTransaction, 'after_delete')
+def after_balance_transaction_delete(mapper: Mapper, connection: Connection, target: BalanceTransaction):
+    recalculate_user_balance(target.user_id)

database/models/basic.py

@@ -1,6 +1,4 @@
-from flask_sqlalchemy import SQLAlchemy
-
-db = SQLAlchemy()
+from database import db
 
 
 class User(db.Model):
@@ -14,6 +12,7 @@ class User(db.Model):
     is_admin = db.Column(db.Boolean, nullable=False, default=False, server_default='0', comment='Админ ли юзверь')
 
     city_id = db.Column(db.Integer, nullable=False, default='1', server_default='1', comment='ID страны')
+    balance = db.Column(db.Double, nullable=False, default=0, server_default='0', comment='Баланс пользователя')
 
 
 class Assembly(db.Model):