partner-core/fill_db.py

import random
from uuid import uuid4
from sqlmodel import Session
from main import AUTH_DB_ENGINE, TgAgent, Ref, Sale, Transaction, Account
from sqlalchemy import text
from datetime import datetime, timedelta
from hashlib import sha256
from passlib.context import CryptContext


# Константа: список 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 

# ---
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

def get_password_hash(password):
    return pwd_context.hash(password)

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.commit()
        # 0. Accounts
        accounts = []
        for i in range(4):
            acc = Account(
                login=f"user{i+1}",
                password_hash=get_password_hash("password123"),  # теперь храним хеш
                name=NAMES[i % len(NAMES)],
                email=f"user{i+1}@example.com",
                balance=round(random.uniform(1000, 10000), 2)
            )
            session.add(acc)
            accounts.append(acc)
        session.commit()
        for acc in accounts:
            session.refresh(acc)
        # 1. 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)],
                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)

        # 2. 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)

        # 3. 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 * random.uniform(0.5, 1.0), 2)
                dt = random.choice(date_list)
                sale = Sale(
                    cost=cost,
                    crediting=crediting,
                    ref=ref.id,
                    sale_id=str(uuid4()),
                    create_dttm=dt,
                    update_dttm=dt
                )
                session.add(sale)
        session.commit()

        # 4. Transactions (только withdrawal на агента)
        TRANSACTION_STATUSES = ['process', 'done', 'error', 'waiting']
        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)
                # Гарантируем, что каждый статус будет использован хотя бы раз
                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),
                    tg_agent_id=tg_agent.id,
                    status=status,
                    create_dttm=dt,
                    update_dttm=dt
                )
                session.add(transaction)
        session.commit()
    print("База успешно заполнена!")

if __name__ == "__main__":
    fill_db()