Добавлены новые модели для агентских и партнерских транзакций, а также балансов компаний и агентов. Обновлено заполнение базы данных с учетом новых таблиц и логики транзакций. Изменены функции для работы с транзакциями, включая фильтрацию по статусам и датам. Улучшены комментарии для ясности кода.
This commit is contained in:
Redsandyg
114
fill_db.py
114
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("База успешно заполнена!")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user