partner-core/fill_db.py

307 lines
13 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import random
from uuid import uuid4
from sqlmodel import Session
from sql_models import TgAgent, Ref, Sale, Account, Company, AgentTransaction, PartnerTransaction, CompanyBalance, AgentBalance, IntegrationToken
from sqlalchemy import text
from datetime import datetime, timedelta
from hashlib import sha256
from helpers_bff import AUTH_DB_ENGINE, get_password_hash
# Константа: список 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 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 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()
# 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)
# 3. Refs (минимум 22 на агента)
refs = []
desc_count = len(ALL_DESCRIPTIONS)
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)
ref = Ref(
tg_agent_id=tg_agent.id,
ref=ref_val,
description=desc_val,
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)
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)
crediting = round(cost * (company.agent_commission / 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,
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()