332 lines
14 KiB
Python
332 lines
14 KiB
Python
import random
|
||
from uuid import uuid4
|
||
from sqlmodel import Session
|
||
from sql_models import TgAgent, Ref, Sale, Account, Company, AgentTransaction, PartnerTransaction, CompanyBalance, AgentBalance, IntegrationToken, SaleCategory
|
||
from sqlalchemy import text
|
||
from datetime import datetime, timedelta
|
||
from hashlib import sha256
|
||
from helpers_bff import AUTH_DB_ENGINE, get_password_hash
|
||
import string
|
||
|
||
|
||
# Константа: список user_ids
|
||
USER_IDS = [1001, 256844410, 426261192, 564746427]
|
||
|
||
# Примеры телефонов
|
||
PHONES = [
|
||
"+79991234567",
|
||
"+79997654321",
|
||
"+79993456789",
|
||
"+79992345678",
|
||
"+79994561234"
|
||
]
|
||
|
||
# Примеры описаний
|
||
DESCRIPTIONS = [
|
||
"Партнёр по рекламе",
|
||
"Блогер",
|
||
"Тестовая ссылка",
|
||
"Промо акция",
|
||
"VIP клиент",
|
||
"Реклама в Telegram от 01.05.2025",
|
||
"Пост в Instagram от 15.04.2025",
|
||
"Рассылка в ВК от 24.02.2025",
|
||
"Промо в Facebook от 10.03.2025",
|
||
"Пост в Twitter от 05.03.2025",
|
||
"Реклама в YouTube от 12.04.2025",
|
||
"Промо в TikTok от 20.04.2025",
|
||
"Пост в LinkedIn от 18.03.2025",
|
||
"Реклама в Одноклассниках от 22.03.2025",
|
||
"Промо в Pinterest от 30.03.2025",
|
||
"Пост в Reddit от 02.04.2025",
|
||
"Реклама в Dzen от 11.04.2025",
|
||
"Промо в Medium от 08.04.2025",
|
||
"Пост в WhatsApp от 14.04.2025",
|
||
"Реклама в Discord от 21.04.2025"
|
||
]
|
||
|
||
# Примеры имен и логинов
|
||
NAMES = [
|
||
"Иван Иванов",
|
||
"Петр Петров",
|
||
"Сергей Сергеев",
|
||
"Анна Смирнова"
|
||
]
|
||
LOGINS = [
|
||
"ivanov1001",
|
||
"petrov256",
|
||
"serg426",
|
||
"anna564"
|
||
]
|
||
|
||
# --- Загрузка mock-данных ---
|
||
|
||
ALL_DESCRIPTIONS = DESCRIPTIONS
|
||
|
||
# ---
|
||
def get_date_list(days=7):
|
||
today = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0)
|
||
return [today - timedelta(days=i) for i in range(days, -1, -1)]
|
||
|
||
def fill_db():
|
||
date_list = get_date_list(7) # 8 дней: от недели назад до сегодня
|
||
with Session(AUTH_DB_ENGINE) as session:
|
||
# Очистка таблиц
|
||
# session.execute(text("DELETE FROM promocode"))
|
||
session.execute(text("DELETE FROM sale"))
|
||
session.execute(text("DELETE FROM ref"))
|
||
session.execute(text("DELETE FROM tgagent"))
|
||
session.execute(text("DELETE FROM account"))
|
||
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 salecategory"))
|
||
session.execute(text("DELETE FROM company"))
|
||
session.commit()
|
||
# 0. Company
|
||
company = Company(
|
||
name="RE: Premium",
|
||
commission=10.0,
|
||
agent_commission=15.0,
|
||
key="re-premium-key",
|
||
)
|
||
session.add(company)
|
||
session.commit()
|
||
session.refresh(company)
|
||
|
||
# 0.1 IntegrationTokens
|
||
for _ in range(3): # Создаем 3 токена для каждой компании
|
||
new_token_value = str(uuid4()) # Генерируем уникальный токен
|
||
token_hash = sha256(new_token_value.encode()).hexdigest() # Хешируем токен для хранения
|
||
masked_token = new_token_value[:5] + "***********************" + new_token_value[-4:] # Генерируем замаскированный токен
|
||
|
||
integration_token = IntegrationToken(
|
||
description=random.choice(DESCRIPTIONS), # Используем существующие описания
|
||
token_hash=token_hash,
|
||
masked_token=masked_token,
|
||
company_id=company.id,
|
||
use_dttm=random.choice(date_list) if random.random() < 0.7 else None # Пример: 70% токенов будут иметь дату использования
|
||
)
|
||
session.add(integration_token)
|
||
session.commit()
|
||
|
||
# 0.2 SaleCategory
|
||
sale_categories = [
|
||
SaleCategory(category="basic", description="Базовая продажа", perc=10.0, company_id=company.id),
|
||
SaleCategory(category="premium", description="Премиум продажа", perc=20.0, company_id=company.id),
|
||
SaleCategory(category="vip", description="VIP продажа", perc=30.0, company_id=company.id),
|
||
]
|
||
for cat in sale_categories:
|
||
session.add(cat)
|
||
session.commit()
|
||
for cat in sale_categories:
|
||
session.refresh(cat)
|
||
|
||
# 1. Accounts
|
||
accounts = []
|
||
for i in range(4):
|
||
name_parts = NAMES[i % len(NAMES)].split()
|
||
first_name = name_parts[0]
|
||
surname = name_parts[1] if len(name_parts) > 1 else 'Тестов'
|
||
acc = Account(
|
||
login=f"user{i+1}",
|
||
password_hash=get_password_hash("password123"),
|
||
firstName=first_name,
|
||
surname=surname,
|
||
phone=PHONES[i % len(PHONES)],
|
||
email=f"user{i+1}@example.com",
|
||
company_id=company.id
|
||
)
|
||
session.add(acc)
|
||
accounts.append(acc)
|
||
session.commit()
|
||
for acc in accounts:
|
||
session.refresh(acc)
|
||
# 2. TgAgents
|
||
tg_agents = []
|
||
for i, tg_agent_id in enumerate(USER_IDS):
|
||
dt = random.choice(date_list)
|
||
hash_value = sha256(f"{tg_agent_id}sold".encode()).hexdigest()
|
||
tg_agent = TgAgent(
|
||
tg_id=tg_agent_id,
|
||
chat_id=tg_agent_id, # chat_id совпадает с tg_id
|
||
phone=PHONES[i % len(PHONES)],
|
||
name=NAMES[i % len(NAMES)],
|
||
login=LOGINS[i % len(LOGINS)],
|
||
company_id=company.id,
|
||
create_dttm=dt,
|
||
update_dttm=dt,
|
||
hash=hash_value
|
||
)
|
||
session.add(tg_agent)
|
||
tg_agents.append(tg_agent)
|
||
session.commit()
|
||
for tg_agent in tg_agents:
|
||
session.refresh(tg_agent)
|
||
# Отладка: количество агентов
|
||
agent_count = session.execute(text("SELECT COUNT(*) FROM tgagent")).scalar()
|
||
print(f'Агентов в базе: {agent_count}')
|
||
# 3. Refs (минимум 22 на агента)
|
||
refs = []
|
||
desc_count = len(ALL_DESCRIPTIONS)
|
||
alphabet = string.ascii_letters + string.digits + "!@#$%^&*"
|
||
for tg_agent in tg_agents:
|
||
ref_count = random.randint(22, int(22 * 1.25)) # от 22 до 27
|
||
for j in range(ref_count):
|
||
ref_val = str(uuid4())
|
||
desc_val = ALL_DESCRIPTIONS[(j % desc_count)]
|
||
dt = random.choice(date_list)
|
||
promocode = ''.join(random.choices(alphabet, k=8))
|
||
# Проверяем уникальность промокода среди уже созданных рефов
|
||
while any(r.promocode == promocode for r in refs):
|
||
promocode = ''.join(random.choices(alphabet, k=8))
|
||
ref = Ref(
|
||
tg_agent_id=tg_agent.id,
|
||
ref=ref_val,
|
||
description=desc_val,
|
||
promocode=promocode,
|
||
create_dttm=dt,
|
||
update_dttm=dt
|
||
)
|
||
session.add(ref)
|
||
refs.append(ref)
|
||
session.commit()
|
||
for ref in refs:
|
||
session.refresh(ref)
|
||
# 4. Sales (минимум 20 на каждый ref)
|
||
all_categories = session.query(SaleCategory).filter_by(company_id=company.id).all()
|
||
for ref in refs:
|
||
sale_count = random.randint(20, int(20 * 1.25)) # от 20 до 25
|
||
for _ in range(sale_count):
|
||
cost = round(random.uniform(100, 1000), 2)
|
||
sale_category = random.choice(all_categories)
|
||
crediting = round(cost * (sale_category.perc / 100.0), 2)
|
||
# Генерируем случайную дату и время в пределах последних 7 дней
|
||
end_dttm = datetime.utcnow()
|
||
start_dttm = end_dttm - timedelta(days=7)
|
||
time_diff = end_dttm - start_dttm
|
||
random_seconds = random.uniform(0, time_diff.total_seconds())
|
||
sale_dttm = start_dttm + timedelta(seconds=random_seconds)
|
||
sale = Sale(
|
||
cost=cost,
|
||
crediting=crediting,
|
||
ref=ref.id,
|
||
sale_id=str(uuid4()),
|
||
company_id=company.id,
|
||
category=sale_category.id,
|
||
sale_dttm=sale_dttm,
|
||
create_dttm=sale_dttm, # create_dttm также будет случайным в этом диапазоне
|
||
update_dttm=sale_dttm # update_dttm также будет случайным в этом диапазоне
|
||
)
|
||
session.add(sale)
|
||
session.commit()
|
||
# 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:
|
||
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', 'new']
|
||
PARTNER_TRANSACTION_TYPES = ['deposit', 'agent_payout', 'service_fee']
|
||
PARTNER_TRANSACTION_STATUSES = ['process', 'done', 'error', 'new']
|
||
|
||
waiting_transactions_to_ensure = 7
|
||
waiting_transactions_count = 0
|
||
|
||
for tg_agent in tg_agents:
|
||
# Генерируем несколько групп транзакций для каждого агента
|
||
for _ in range(random.randint(3, 6)): # От 3 до 6 групп на агента
|
||
transaction_group_id = uuid4()
|
||
dt = random.choice(date_list)
|
||
agent_trans_amount = round(random.uniform(500, 3000), 2)
|
||
|
||
if waiting_transactions_count < waiting_transactions_to_ensure:
|
||
agent_trans_status = 'waiting'
|
||
waiting_transactions_count += 1
|
||
else:
|
||
agent_trans_status = random.choice(AGENT_TRANSACTION_STATUSES)
|
||
|
||
# Создаем AgentTransaction
|
||
agent_transaction = AgentTransaction(
|
||
tg_agent_id=tg_agent.id,
|
||
amount=agent_trans_amount,
|
||
status=agent_trans_status,
|
||
transaction_group=transaction_group_id,
|
||
create_dttm=dt,
|
||
update_dttm=dt
|
||
)
|
||
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("База успешно заполнена!")
|
||
|
||
if __name__ == "__main__":
|
||
fill_db() |