database.py - master - ARCHAEUS/archaeus22-bot

database.py
347 строк12 KB
ARCHAEUS
Update: database.py
08 авг 2026, 01:49
08 авг 2026, 01:494c34afc
# ============================================================
# РАБОТА С БАЗОЙ ДАННЫХ (SQLite) ДЛЯ БОТА «АРХЕЙ22» (ARCHAEUS22)
# ============================================================

import sqlite3
from datetime import datetime
from config import DATABASE_PATH, PROMOCODES


# ============================================================
# 1. ПОДКЛЮЧЕНИЕ К БАЗЕ ДАННЫХ
# ============================================================

def get_db():
    """Возвращает соединение с базой данных."""
    return sqlite3.connect(DATABASE_PATH)


# ============================================================
# 2. ИНИЦИАЛИЗАЦИЯ ТАБЛИЦ
# ============================================================

def init_db():
    """Создаёт все необходимые таблицы, если их нет."""
    conn = get_db()
    c = conn.cursor()

    # Таблица пользователей
    c.execute('''
        CREATE TABLE IF NOT EXISTS users (
            user_id INTEGER PRIMARY KEY,
            full_name TEXT,
            birth_date TEXT,
            birth_time TEXT,
            birth_location TEXT,
            birth_lat REAL,
            birth_lon REAL,
            consent_given INTEGER DEFAULT 0,
            consent_date TEXT,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP
        )
    ''')

    # Таблица использования бесплатных услуг
    c.execute('''
        CREATE TABLE IF NOT EXISTS free_usage (
            user_id INTEGER,
            service_type TEXT,
            used_at TEXT DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (user_id, service_type)
        )
    ''')

    # Таблица подписок
    c.execute('''
        CREATE TABLE IF NOT EXISTS subscriptions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER,
            sub_type TEXT,
            start_date TEXT,
            end_date TEXT,
            auto_renew INTEGER DEFAULT 0,
            UNIQUE(user_id, sub_type)
        )
    ''')

    # Таблица промокодов
    c.execute('''
        CREATE TABLE IF NOT EXISTS promocodes (
            code TEXT PRIMARY KEY,
            discount INTEGER,
            max_uses INTEGER,
            used_count INTEGER DEFAULT 0,
            expires TEXT
        )
    ''')

    # Таблица использования промокодов
    c.execute('''
        CREATE TABLE IF NOT EXISTS promo_usage (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER,
            code TEXT,
            used_at TEXT DEFAULT CURRENT_TIMESTAMP
        )
    ''')

    # Таблица реферальных связей
    c.execute('''
        CREATE TABLE IF NOT EXISTS referrals (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            referrer_id INTEGER,
            referred_id INTEGER UNIQUE,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP
        )
    ''')

    # Таблица баланса (бонусы за рефералов)
    c.execute('''
        CREATE TABLE IF NOT EXISTS balances (
            user_id INTEGER PRIMARY KEY,
            balance INTEGER DEFAULT 0
        )
    ''')

    # Добавляем базовые промокоды из config.py, если их нет
    for code, data in PROMOCODES.items():
        c.execute('''
            INSERT OR IGNORE INTO promocodes (code, discount, max_uses, expires)
            VALUES (?, ?, ?, ?)
        ''', (code, data['discount'], data['max_uses'], None))

    conn.commit()
    conn.close()


# ============================================================
# 3. ПОЛЬЗОВАТЕЛИ
# ============================================================

def get_user(user_id: int) -> dict | None:
    """Возвращает данные пользователя по Telegram ID."""
    conn = get_db()
    c = conn.cursor()
    c.execute("SELECT * FROM users WHERE user_id = ?", (user_id,))
    row = c.fetchone()
    conn.close()
    if row:
        return dict(zip(
            ["user_id", "full_name", "birth_date", "birth_time",
             "birth_location", "birth_lat", "birth_lon",
             "consent_given", "consent_date", "created_at"],
            row
        ))
    return None


def save_user(user_id: int, full_name: str, birth_date: str, birth_time: str,
              birth_location: str, birth_lat: float, birth_lon: float):
    """Сохраняет или обновляет данные пользователя."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        INSERT OR REPLACE INTO users
        (user_id, full_name, birth_date, birth_time, birth_location, birth_lat, birth_lon)
        VALUES (?, ?, ?, ?, ?, ?, ?)
    ''', (user_id, full_name, birth_date, birth_time, birth_location, birth_lat, birth_lon))
    conn.commit()
    conn.close()


def update_consent(user_id: int, consent_given: int = 1):
    """Обновляет статус согласия на обработку ПДн."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        UPDATE users SET consent_given = ?, consent_date = ?
        WHERE user_id = ?
    ''', (consent_given, datetime.now().isoformat(), user_id))
    conn.commit()
    conn.close()


# ============================================================
# 4. БЕСПЛАТНЫЕ УСЛУГИ
# ============================================================

def has_free_service(user_id: int, service_type: str) -> bool:
    """Проверяет, использовал ли пользователь бесплатную услугу."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        SELECT 1 FROM free_usage WHERE user_id = ? AND service_type = ?
    ''', (user_id, service_type))
    result = c.fetchone()
    conn.close()
    return bool(result)


def mark_free_service(user_id: int, service_type: str):
    """Отмечает, что пользователь использовал бесплатную услугу."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        INSERT INTO free_usage (user_id, service_type) VALUES (?, ?)
    ''', (user_id, service_type))
    conn.commit()
    conn.close()


# ============================================================
# 5. ПОДПИСКИ
# ============================================================

def add_subscription(user_id: int, sub_type: str, end_date: str):
    """Добавляет или обновляет подписку пользователя."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        INSERT OR REPLACE INTO subscriptions (user_id, sub_type, start_date, end_date)
        VALUES (?, ?, ?, ?)
    ''', (user_id, sub_type, datetime.now().isoformat(), end_date))
    conn.commit()
    conn.close()


def get_subscription(user_id: int, sub_type: str) -> str | None:
    """Возвращает дату окончания подписки, если она активна."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        SELECT end_date FROM subscriptions
        WHERE user_id = ? AND sub_type = ? AND end_date > datetime('now')
    ''', (user_id, sub_type))
    row = c.fetchone()
    conn.close()
    return row[0] if row else None


# ============================================================
# 6. ПРОМОКОДЫ
# ============================================================

def validate_promocode(code: str, user_id: int) -> tuple:
    """
    Проверяет промокод: возвращает (скидка, None) или (None, сообщение об ошибке).
    """
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        SELECT discount, max_uses, used_count, expires FROM promocodes
        WHERE code = ?
    ''', (code.lower(),))
    row = c.fetchone()
    if not row:
        conn.close()
        return None, "Промокод не найден"

    discount, max_uses, used_count, expires = row
    if used_count >= max_uses:
        conn.close()
        return None, "Промокод уже использован максимальное количество раз"

    if expires:
        try:
            if datetime.now().date() > datetime.strptime(expires, "%Y-%m-%d").date():
                conn.close()
                return None, "Срок действия промокода истёк"
        except ValueError:
            pass  # Если дата не парсится — игнорируем

    # Проверяем, не использовал ли пользователь этот промокод ранее
    c.execute('''
        SELECT 1 FROM promo_usage WHERE user_id = ? AND code = ?
    ''', (user_id, code.lower()))
    if c.fetchone():
        conn.close()
        return None, "Вы уже использовали этот промокод"

    conn.close()
    return discount, None


def use_promocode(code: str, user_id: int):
    """Активирует промокод для пользователя (увеличивает счётчик использований)."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        UPDATE promocodes SET used_count = used_count + 1
        WHERE code = ?
    ''', (code.lower(),))
    c.execute('''
        INSERT INTO promo_usage (user_id, code) VALUES (?, ?)
    ''', (user_id, code.lower()))
    conn.commit()
    conn.close()


# ============================================================
# 7. РЕФЕРАЛЬНАЯ СИСТЕМА
# ============================================================

def add_referral(referrer_id: int, referred_id: int):
    """Добавляет реферальную связь."""
    conn = get_db()
    c = conn.cursor()
    try:
        c.execute('''
            INSERT INTO referrals (referrer_id, referred_id) VALUES (?, ?)
        ''', (referrer_id, referred_id))
        conn.commit()
    except sqlite3.IntegrityError:
        # Связь уже существует — игнорируем
        pass
    conn.close()


def get_referrer(user_id: int) -> int | None:
    """Возвращает ID реферера для указанного пользователя."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        SELECT referrer_id FROM referrals WHERE referred_id = ?
    ''', (user_id,))
    row = c.fetchone()
    conn.close()
    return row[0] if row else None


def get_referrals_count(user_id: int) -> int:
    """Возвращает количество приглашённых пользователем друзей."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        SELECT COUNT(*) FROM referrals WHERE referrer_id = ?
    ''', (user_id,))
    row = c.fetchone()
    conn.close()
    return row[0] if row else 0


# ============================================================
# 8. БАЛАНС (БОНУСЫ ЗА РЕФЕРАЛОВ)
# ============================================================

def add_balance(user_id: int, amount: int):
    """Добавляет сумму на баланс пользователя."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        INSERT INTO balances (user_id, balance) VALUES (?, ?)
        ON CONFLICT(user_id) DO UPDATE SET balance = balance + ?
    ''', (user_id, amount, amount))
    conn.commit()
    conn.close()


def get_balance(user_id: int) -> int:
    """Возвращает текущий баланс пользователя."""
    conn = get_db()
    c = conn.cursor()
    c.execute('''
        SELECT balance FROM balances WHERE user_id = ?
    ''', (user_id,))
    row = c.fetchone()
    conn.close()
    return row[0] if row else 0