Добавлены новые модели для агентских и партнерских транзакций, а также балансов компаний и агентов. Обновлено заполнение базы данных с учетом новых таблиц и логики транзакций. Изменены функции для работы с транзакциями, включая фильтрацию по статусам и датам. Улучшены комментарии для ясности кода.

2025-06-06 14:17:47 +03:00 · 2025-06-06 14:17:47 +03:00 · 8c6fadb180
commit 8c6fadb180
parent 161e0b3ec4
2 changed files with 252 additions and 70 deletions
--- a/fill_db.py
+++ b/fill_db.py
@ -1,7 +1,7 @@
 import random
 from uuid import uuid4
 from sqlmodel import Session
-from main import AUTH_DB_ENGINE, TgAgent, Ref, Sale, Transaction, Account, Company
+from main import AUTH_DB_ENGINE, TgAgent, Ref, Sale, Account, Company, AgentTransaction, PartnerTransaction, CompanyBalance, AgentBalance
 from sqlalchemy import text
 from datetime import datetime, timedelta
 from hashlib import sha256
@ -80,7 +80,10 @@ def fill_db():
        session.execute(text("DELETE FROM ref"))
        session.execute(text("DELETE FROM tgagent"))
        session.execute(text("DELETE FROM account"))
-        session.execute(text('DELETE FROM "transaction"'))
+        session.execute(text('DELETE FROM "agent_transactions"'))
+        session.execute(text('DELETE FROM "partner_transactions"'))
+        session.execute(text('DELETE FROM "company_balances"'))
+        session.execute(text('DELETE FROM "agent_balances"'))
        session.execute(text("DELETE FROM company"))
        session.commit()
        # 0. Company
@ -172,29 +175,102 @@ def fill_db():
                )
                session.add(sale)
        session.commit()
-        # 5. Transactions (только withdrawal на агента)
-        TRANSACTION_STATUSES = ['process', 'done', 'error', 'waiting']
+        # 5. Заполнение новых таблиц
+        # 5.1 CompanyBalance
+        company_balance = CompanyBalance(
+            company_id=company.id,
+            available_balance=round(random.uniform(10000, 50000), 2),
+            pending_balance=round(random.uniform(1000, 10000), 2),
+            updated_dttm=datetime.utcnow()
+        )
+        session.add(company_balance)
+        session.commit()
+        session.refresh(company_balance)
+
+        # 5.2 AgentBalances
+        agent_balances = []
        for tg_agent in tg_agents:
-            withdrawal_count = random.randint(5, int(5 * 1.25))  # от 5 до 6
-            used_statuses = set()
-            for i in range(withdrawal_count):
+            dt = random.choice(date_list)
+            agent_balance = AgentBalance(
+                tg_agent_id=tg_agent.id,
+                available_balance=round(random.uniform(100, 5000), 2),
+                frozen_balance=round(random.uniform(0, 1000), 2),
+                updated_dttm=dt
+            )
+            session.add(agent_balance)
+            agent_balances.append(agent_balance)
+        session.commit()
+        for balance in agent_balances:
+            session.refresh(balance)
+
+        # 5.3 AgentTransactions and PartnerTransactions
+        AGENT_TRANSACTION_STATUSES = ['waiting', 'process', 'done', 'reject', 'error']
+        PARTNER_TRANSACTION_TYPES = ['deposit', 'agent_payout', 'service_fee']
+        PARTNER_TRANSACTION_STATUSES = ['process', 'done', 'error']
+
+        for tg_agent in tg_agents:
+            # Генерируем несколько групп транзакций для каждого агента
+            for _ in range(random.randint(3, 6)): # От 3 до 6 групп на агента
+                transaction_group_id = uuid4()
                dt = random.choice(date_list)
-                # Гарантируем, что каждый статус будет использован хотя бы раз
-                if len(used_statuses) < len(TRANSACTION_STATUSES):
-                    status = TRANSACTION_STATUSES[len(used_statuses)]
-                    used_statuses.add(status)
-                else:
-                    status = random.choice(TRANSACTION_STATUSES)
-                transaction = Transaction(
-                    transaction_id=str(uuid4()),
-                    sum=round(random.uniform(200, 3000), 2),
+                agent_trans_amount = round(random.uniform(500, 3000), 2)
+                agent_trans_status = random.choice(AGENT_TRANSACTION_STATUSES)
+
+                # Создаем AgentTransaction
+                agent_transaction = AgentTransaction(
                    tg_agent_id=tg_agent.id,
-                    status=status,
-                    company_id=company.id,
+                    amount=agent_trans_amount,
+                    status=agent_trans_status,
+                    transaction_group=transaction_group_id,
                    create_dttm=dt,
                    update_dttm=dt
                )
-                session.add(transaction)
+                session.add(agent_transaction)
+                session.commit()
+                session.refresh(agent_transaction)
+
+                # Создаем соответствующие PartnerTransactions
+                # Для каждой AgentTransaction создаем PartnerTransaction типа 'agent_payout'
+                if agent_trans_status != 'waiting': # Создаем партнерскую транзакцию только если агентская не в статусе 'waiting'
+                     # Добавляем PartnerTransaction для выплаты агенту
+                    partner_payout = PartnerTransaction(
+                        company_id=company.id,
+                        type='agent_payout',
+                        amount=agent_trans_amount,
+                        status=random.choice([s for s in PARTNER_TRANSACTION_STATUSES if s != 'process']) if agent_trans_status in ['done', 'error', 'reject'] else 'process', # Статус зависит от статуса агентской
+                        transaction_group=transaction_group_id,
+                        agent_transaction_id=agent_transaction.id,
+                        create_dttm=dt,
+                        update_dttm=dt
+                    )
+                    session.add(partner_payout)
+
+
+                # Добавляем другие типы PartnerTransactions для разнообразия
+                if random.random() < 0.5: # 50% шанс добавить депозит
+                    partner_deposit = PartnerTransaction(
+                        company_id=company.id,
+                        type='deposit',
+                        amount=round(random.uniform(1000, 10000), 2),
+                        status=random.choice(PARTNER_TRANSACTION_STATUSES),
+                        transaction_group=uuid4(), # Новая группа для независимых транзакций
+                        create_dttm=dt,
+                        update_dttm=dt
+                    )
+                    session.add(partner_deposit)
+
+                if random.random() < 0.3: # 30% шанс добавить комиссию
+                    partner_fee = PartnerTransaction(
+                        company_id=company.id,
+                        type='service_fee',
+                        amount=round(random.uniform(50, 500), 2),
+                        status=random.choice(PARTNER_TRANSACTION_STATUSES),
+                        transaction_group=uuid4(), # Новая группа
+                        create_dttm=dt,
+                        update_dttm=dt
+                    )
+                    session.add(partner_fee)
+
        session.commit()
    print("База успешно заполнена!")

--- a/main.py
+++ b/main.py
@ -1,3 +1,4 @@
+import uuid
 from fastapi import FastAPI, Depends, HTTPException, status, Query, Body, Request
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from sqlmodel import SQLModel, Field, create_engine, Session, select
@ -12,6 +13,7 @@ from hashlib import sha256
 import jwt
 from jwt.exceptions import InvalidTokenError
 from pydantic import BaseModel, EmailStr
+from enum import Enum

 # Конфигурация
 AUTH_DATABASE_ADDRESS = "sqlite:///partner.db"
@ -55,20 +57,48 @@ class Sale(SQLModel, table=True):
    create_dttm: datetime = Field(default_factory=datetime.utcnow)
    update_dttm: datetime = Field(default_factory=datetime.utcnow)

-class Transaction(SQLModel, table=True):
+class AgentTransaction(SQLModel, table=True):
+    __tablename__ = "agent_transactions" # Указываем имя таблицы явно
    id: Optional[int] = Field(default=None, primary_key=True)
-    transaction_id: str = Field(default_factory=lambda: str(uuid4()), index=True, unique=True)
-    sum: float
-    tg_agent_id: int = Field(foreign_key="tgagent.id")
-    status: str  # 'process' || 'done' || 'error' || 'waiting'
-    company_id: int = Field(foreign_key="company.id")
+    tg_agent_id: int = Field(foreign_key="tgagent.id") # ID агента, связь с TgAgent
+    amount: float # Используем float для DECIMAL(15,2)
+    status: str # 'waiting', 'process', 'done', 'reject', 'error'
+    transaction_group: uuid.UUID = Field(default_factory=uuid.uuid4) # UUID для группировки
    create_dttm: datetime = Field(default_factory=datetime.utcnow)
    update_dttm: datetime = Field(default_factory=datetime.utcnow)

+class PartnerTransaction(SQLModel, table=True):
+    __tablename__ = "partner_transactions" # Указываем имя таблицы явно
+    id: Optional[int] = Field(default=None, primary_key=True)
+    company_id: int = Field(foreign_key="company.id") # ID партнера, связь с Company
+    type: str # 'deposit', 'agent_payout', 'service_fee'
+    amount: float # Используем float для DECIMAL(15,2)
+    status: str # 'process', 'done', 'error'
+    transaction_group: uuid.UUID # UUID для группировки, может быть связан с agent_transactions
+    agent_transaction_id: Optional[int] = Field(default=None, foreign_key="agent_transactions.id") # Связь с агентской транзакцией
+    create_dttm: datetime = Field(default_factory=datetime.utcnow)
+    update_dttm: datetime = Field(default_factory=datetime.utcnow)
+
+class CompanyBalance(SQLModel, table=True):
+    __tablename__ = "company_balances" # Указываем имя таблицы явно
+    id: Optional[int] = Field(default=None, primary_key=True)
+    company_id: int = Field(foreign_key="company.id", unique=True) # ID компании, уникальный баланс на компанию
+    available_balance: float = Field(default=0.0) # Используем float для DECIMAL(15,2)
+    pending_balance: float = Field(default=0.0) # Используем float для DECIMAL(15,2)
+    updated_dttm: datetime = Field(default_factory=datetime.utcnow)
+
+class AgentBalance(SQLModel, table=True):
+    __tablename__ = "agent_balances" # Указываем имя таблицы явно
+    id: Optional[int] = Field(default=None, primary_key=True)
+    tg_agent_id: int = Field(foreign_key="tgagent.id", unique=True) # ID агента, уникальный баланс на агента
+    available_balance: float = Field(default=0.0) # Используем float для DECIMAL(15,2)
+    frozen_balance: float = Field(default=0.0) # Используем float для DECIMAL(15,2)
+    updated_dttm: datetime = Field(default_factory=datetime.utcnow)
+
 class Account(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    login: str = Field(index=True, unique=True)
-    password_hash: str  # теперь хранится hash пароля
+    password_hash: str  
    firstName: Optional[str] = None
    surname: Optional[str] = None
    phone: Optional[str] = None
@ -87,6 +117,25 @@ class AccountPasswordChangeRequest(BaseModel):
    currentPassword: str
    newPassword: str

+
+class TransactionStatus(str, Enum): # Определяем Enum для статусов
+    WAITING = 'waiting'
+    PROCESS = 'process'
+    DONE = 'done'
+    REJECT = 'reject'
+    ERROR = 'error'
+
+
+    # Новая модель ответа для агентских транзакций с именем агента
+class AgentTransactionResponse(BaseModel):
+    amount: float
+    status: TransactionStatus # Используем Enum
+    transaction_group: uuid.UUID
+    create_dttm: datetime
+    update_dttm: datetime
+    agent_name: Optional[str] = None # Поле для имени агента
+
+
 # Создание движка базы данных
 AUTH_DB_ENGINE = create_engine(AUTH_DATABASE_ADDRESS, echo=True)
 SQLModel.metadata.create_all(AUTH_DB_ENGINE)
@ -188,12 +237,13 @@ def create_access_token(data: dict, expires_delta: timedelta = None):

@app.post("/token", response_model=Token, tags=["bff"])
 def login_account_for_access_token(
-    login: str = Body(...),
-    password: str = Body(...),
+    # login: str = Body(...),
+    # password: str = Body(...),
+    form_data: OAuth2PasswordRequestForm = Depends(),
    db: Session = Depends(get_db)
 ):
-    account = get_account_by_login(db, login)
-    if not account or not verify_password(password, account.password_hash):
+    account = get_account_by_login(db, form_data.username)
+    if not account or not verify_password(form_data.password, account.password_hash):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect login or password",
@ -251,12 +301,9 @@ def get_stat(current_tg_agent: TgAgent = Depends(get_current_tg_agent), db: Sess
    total_sales = db.exec(select(Sale).where(Sale.ref.in_(ref_ids))).all()
    totalSales = len(total_sales)
    totalIncome = sum(sale.crediting for sale in total_sales)
-    withdrawals = db.exec(
-        select(Transaction).where(
-            Transaction.tg_agent_id == current_tg_agent.id
-        )
-    ).all()
-    availableWithdrawal = totalIncome - sum(t.sum for t in withdrawals)
+    # Заменено получение доступного остатка из AgentBalance
+    agent_balance = db.exec(select(AgentBalance).where(AgentBalance.tg_agent_id == current_tg_agent.id)).first()
+    availableWithdrawal = agent_balance.available_balance if agent_balance else 0.0
    return {
        "totalSales": totalSales,
        "totalIncome": totalIncome,
@ -276,10 +323,9 @@ def get_dashboard_cards(db: Session = Depends(get_db)):
    unique_agents = db.exec(select(TgAgent.tg_id)).all()
    activeReferrals = len(set(unique_agents))

-    # 4. Ожидающие выплаты - разница между суммой всех Sale.crediting и суммой всех Transaction.sum
-    all_transactions = db.exec(select(Transaction)).all()
-    totalTransactions = sum(t.sum for t in all_transactions)
-    pendingPayouts = totalPayouts - totalTransactions
+    # 4. Ожидающие выплаты - сумма AgentTransaction со статусом 'waiting'
+    pending_agent_transactions = db.exec(select(AgentTransaction).where(AgentTransaction.status == 'waiting')).all()
+    pendingPayouts = sum(t.amount for t in pending_agent_transactions)

    # 5. Количество продаж
    totalSales = len(total_revenue)
@ -436,15 +482,18 @@ def get_sales_stat(

@app.get("/billing/cards", tags=["bff"])
 def get_billing_cards(db: Session = Depends(get_db)):
-    # 1. cost - сумма всех Sale.cost
+    # 1. cost - Общий заработок (сумма всех Sale.cost)
    sales = db.exec(select(Sale)).all()
    cost = sum(sale.cost for sale in sales)
-    # 2. crediting - сумма всех Sale.crediting
-    crediting = sum(sale.crediting for sale in sales)
-    # 3. pendingPayouts - разница между crediting и суммой всех Transaction.sum
-    transactions = db.exec(select(Transaction)).all()
-    total_transactions = sum(t.sum for t in transactions)
-    pendingPayouts = crediting - total_transactions
+
+    # 2. crediting - Общие выплаты (сумма PartnerTransaction типа 'agent_payout' со статусом 'done')
+    completed_payouts = db.exec(select(PartnerTransaction).where(PartnerTransaction.type == 'agent_payout').where(PartnerTransaction.status == 'done')).all()
+    crediting = sum(t.amount for t in completed_payouts)
+
+    # 3. pendingPayouts - Доступно к выводу всеми партнерами (сумма всех доступных балансов агентов)
+    agent_balances = db.exec(select(AgentBalance)).all()
+    pendingPayouts = sum(balance.available_balance for balance in agent_balances)
+
    return {
        "cost": cost,
        "crediting": crediting,
@ -457,39 +506,44 @@ def get_billing_payouts_transactions(
    date_start: str = Query(None),
    date_end: str = Query(None),
 ):
-    query = select(Transaction)
+    # Используем AgentTransaction вместо Transaction
+    # Явно выбираем обе модели для корректной распаковки
+    query = select(AgentTransaction, TgAgent).join(TgAgent)
    if date_start:
-        query = query.where(Transaction.create_dttm >= date_start)
+        query = query.where(AgentTransaction.create_dttm >= date_start)
    if date_end:
-        query = query.where(Transaction.create_dttm <= date_end)
-    transactions = db.exec(query).all()
+        query = query.where(AgentTransaction.create_dttm <= date_end)
+    # Заказываем по дате создания
+    query = query.order_by(AgentTransaction.create_dttm.desc())
+
+    # Выполняем запрос и формируем результат
+    results = db.exec(query).all()
    result = []
-    for t in transactions:
-        agent = db.exec(select(TgAgent).where(TgAgent.id == t.tg_agent_id)).first()
+    for agent_trans, agent in results:
        result.append({
-            "id": t.transaction_id,
-            "sum": t.sum,
-            "agent": agent.name if agent else None,
-            "status": t.status,
-            "create_dttm": t.create_dttm,
-            "update_dttm": t.update_dttm,
+            "id": agent_trans.transaction_group, # Используем transaction_group как ID транзакции группы
+            "amount": agent_trans.amount,
+            "agent": agent.name if agent else None, # Имя агента из join
+            "status": agent_trans.status,
+            "create_dttm": agent_trans.create_dttm,
+            "update_dttm": agent_trans.update_dttm,
        })
    return result

@app.get("/billing/chart/stat", tags=["bff"])
 def get_billing_chart_stat(db: Session = Depends(get_db)):
-    # Группируем транзакции по дате (день) и статусу
+    # Группируем агентские транзакции по дате (день) и статусу
    result = db.exec(
        select(
-            func.strftime('%Y-%m-%d', Transaction.create_dttm).label('date'),
-            Transaction.status.label('status'),
-            func.count(Transaction.id).label('count')
+            func.strftime('%Y-%m-%d', AgentTransaction.create_dttm).label('date'),
+            AgentTransaction.status.label('status'),
+            func.count(AgentTransaction.id).label('count')
        ).group_by(
-            func.strftime('%Y-%m-%d', Transaction.create_dttm),
-            Transaction.status
+            func.strftime('%Y-%m-%d', AgentTransaction.create_dttm),
+            AgentTransaction.status
        ).order_by(
-            func.strftime('%Y-%m-%d', Transaction.create_dttm),
-            Transaction.status
+            func.strftime('%Y-%m-%d', AgentTransaction.create_dttm),
+            AgentTransaction.status
        )
    ).all()
    data = [
@ -500,11 +554,12 @@ def get_billing_chart_stat(db: Session = Depends(get_db)):

@app.get("/billing/chart/pie", tags=["bff"])
 def get_billing_chart_pie(db: Session = Depends(get_db)):
+    # Группируем агентские транзакции по статусу
    result = db.exec(
        select(
-            Transaction.status.label('status'),
-            func.count(Transaction.id).label('count')
-        ).group_by(Transaction.status)
+            AgentTransaction.status.label('status'),
+            func.count(AgentTransaction.id).label('count')
+        ).group_by(AgentTransaction.status)
    ).all()
    data = [
        {"status": row.status, "count": row.count}
@ -613,3 +668,54 @@ def change_account_password(
    db.commit()
    db.refresh(current_account)
    return {"msg": "Пароль успешно изменён"}
+
+# --- Новый функционал для агентских транзакций партнера ---
+
+@app.get("/account/agent-transaction", response_model=List[AgentTransactionResponse], tags=["bff"])
+def get_account_agent_transactions(
+    statuses: Optional[List[TransactionStatus]] = Query(None), # Изменено на List[TransactionStatus]
+    date_start: str = Query(None), # Добавлен параметр date_start
+    date_end: str = Query(None), # Добавлен параметр date_end
+    current_account: Account = Depends(get_current_account),
+    db: Session = Depends(get_db)
+):
+    """
+    Возвращает список агентских транзакций для компании текущего пользователя,
+    с возможностью фильтрации по статусу и дате создания.
+    """
+    # Получаем ID компании текущего аккаунта
+    company_id = current_account.company_id
+
+    # Строим базовый запрос: выбрать AgentTransaction и TgAgent, связанные с агентами этой компании
+    query = select(AgentTransaction, TgAgent).join(TgAgent).where(TgAgent.company_id == company_id)
+
+    # Если переданы статусы, добавляем фильтрацию по статусам
+    if statuses:
+        query = query.where(AgentTransaction.status.in_(statuses))
+
+    # Если передана дата начала, добавляем фильтрацию по дате создания >= date_start
+    if date_start:
+        query = query.where(AgentTransaction.create_dttm >= date_start)
+
+    # Если передана дата окончания, добавляем фильтрацию по дате создания <= date_end
+    if date_end:
+        query = query.where(AgentTransaction.create_dttm <= date_end)
+
+    # Выполняем запрос
+    results = db.exec(query).all()
+
+    # Формируем список ответов в формате AgentTransactionResponse
+    agent_transactions_response = []
+    for agent_trans, agent in results:
+        agent_transactions_response.append(
+            AgentTransactionResponse(
+                amount=agent_trans.amount,
+                status=agent_trans.status,
+                transaction_group=agent_trans.transaction_group,
+                create_dttm=agent_trans.create_dttm,
+                update_dttm=agent_trans.update_dttm,
+                agent_name=agent.name # Используем имя агента
+            )
+        )
+
+    return agent_transactions_response