create_database.py

import asyncio
import logging
import os
import time

logger = logging.getLogger(__name__)

try:
    from data.database import execute_query
except ImportError:
    print('Ошибка: Не удалось импортировать "execute_query".')
    print('Убедитесь, что вы запускаете скрипт из корневой папки проекта "YandexMusicBot".')
    exit(1)

# --- Основные таблицы ---

SQL_CREATE_USERS_TABLE = '''
CREATE TABLE IF NOT EXISTS users (
    user_id          INTEGER PRIMARY KEY,
    balance          REAL    DEFAULT (0) 
                             NOT NULL,
    lang             TEXT    NOT NULL
                             DEFAULT ru,
    ts_registration  INTEGER NOT NULL,
    ts_premium_until INTEGER,
    premium_notified INTEGER NOT NULL DEFAULT (0),
    action_count     INTEGER NOT NULL
                             DEFAULT (0),
    referrer_id      INTEGER REFERENCES users (user_id) ON DELETE SET NULL
                                                        ON UPDATE CASCADE,
    daily_usage_count INTEGER NOT NULL DEFAULT (0),
    recognize_search_count INTEGER NOT NULL DEFAULT (0),
    last_usage_ts    INTEGER NOT NULL DEFAULT (0)
);
'''

SQL_CREATE_PAYMENTS_TABLE = '''
CREATE TABLE IF NOT EXISTS payments (
    payment_id      INTEGER PRIMARY KEY,
    user_id         INTEGER REFERENCES users (user_id) ON DELETE SET NULL
                                                       ON UPDATE CASCADE
                            NOT NULL,
    amount          REAL    NOT NULL
                            DEFAULT (0),
    days            INTEGER NOT NULL
                            DEFAULT (1),
    screenshot_path TEXT,
    status          INTEGER NOT NULL
                            DEFAULT (0),
    created_ts      INTEGER NOT NULL,
    processed_ts    INTEGER
);
'''

SQL_CREATE_TRACKS_TABLE = '''
CREATE TABLE IF NOT EXISTS tracks (
    track_id    INTEGER PRIMARY KEY,
    source_id   TEXT    NOT NULL,
    source_type TEXT    NOT NULL
                        DEFAULT youtube,
    title       TEXT    NOT NULL,
    artist      TEXT,
    duration    INTEGER NOT NULL,
    UNIQUE (
        source_id,
        source_type
    )
    ON CONFLICT IGNORE
);
'''

SQL_CREATE_PLAY_HISTORY_TABLE = '''
CREATE TABLE IF NOT EXISTS play_history (
    history_id INTEGER PRIMARY KEY,
    user_id    INTEGER REFERENCES users (user_id) ON DELETE CASCADE
                                                  ON UPDATE CASCADE
                       NOT NULL,
    track_id   INTEGER NOT NULL
                       REFERENCES tracks (track_id) ON DELETE CASCADE
                                                    ON UPDATE CASCADE,
    played_ts  INTEGER NOT NULL
);
'''

SQL_CREATE_FAVORITES_TABLE = '''
CREATE TABLE IF NOT EXISTS favorites (
    user_id  INTEGER REFERENCES users (user_id) ON DELETE CASCADE
                                                ON UPDATE CASCADE,
    track_id INTEGER REFERENCES tracks (track_id) ON DELETE CASCADE
                                                  ON UPDATE CASCADE,
    added_ts INTEGER NOT NULL,
    PRIMARY KEY (
        user_id,
        track_id
    )
    ON CONFLICT IGNORE
);
'''

SQL_CREATE_BOT_STATUS_TABLE = '''
CREATE TABLE IF NOT EXISTS bot_status (
    status_id  INTEGER PRIMARY KEY,
    status     INTEGER NOT NULL
                       DEFAULT (1),
    updated_ts INTEGER NOT NULL
);
'''

SQL_CREATE_ADS_TABLE = '''
CREATE TABLE IF NOT EXISTS ads (
    ads_id       INTEGER PRIMARY KEY AUTOINCREMENT,
    content_path TEXT,
    caption      TEXT    NOT NULL
);
'''

# --- Индексы ---

SQL_CREATE_INDEXES = [
    "CREATE INDEX IF NOT EXISTS idx_users_referrer_id ON users (referrer_id);",
    "CREATE INDEX IF NOT EXISTS idx_users_ts_premium_until ON users (ts_premium_until);",

    "CREATE INDEX IF NOT EXISTS idx_payments_user_id ON payments (user_id);",
    "CREATE INDEX IF NOT EXISTS idx_payments_status ON payments (status);",

    "CREATE INDEX IF NOT EXISTS idx_play_history_user_id_played_ts ON play_history (user_id, played_ts DESC);",
    "CREATE INDEX IF NOT EXISTS idx_play_history_track_id ON play_history (track_id);",

    "CREATE INDEX IF NOT EXISTS idx_favorites_user_id_added_ts ON favorites (user_id, added_ts DESC);"
]


async def setup_database():
    """
    Выполняет полную настройку базы данных.
    Используется режим транзакций.
    """
    logger.info('[*] Настройка базы данных...')
    db_folder = 'data/database'

    if not os.path.exists(db_folder):
        os.makedirs(db_folder)
        logger.info(f'[+] Папка "{db_folder}" создана.')

    all_queries = [
        (SQL_CREATE_USERS_TABLE, ()),
        (SQL_CREATE_PAYMENTS_TABLE, ()),
        (SQL_CREATE_TRACKS_TABLE, ()),
        (SQL_CREATE_PLAY_HISTORY_TABLE, ()),
        (SQL_CREATE_FAVORITES_TABLE, ()),
        (SQL_CREATE_BOT_STATUS_TABLE, ()),
        (SQL_CREATE_ADS_TABLE, ()),
    ]
    for index_query in SQL_CREATE_INDEXES:
        all_queries.append((index_query, ()))

    success = await execute_query(all_queries)

    if success:
        logger.info('[+] База данных успешно собрана и настроена!')

        await execute_query(
            '''
            INSERT OR IGNORE INTO bot_status (status_id, status, updated_ts)
            VALUES (?, ?, ?)
            ''', (1, 1, int(time.time()))
        )
        logger.info('[+] Начальный статус бота был установлен на 1 (активен).')
    else:
        logger.info('[!] Произошла нежданная ошибка при настройке базы данных!')


if __name__ == '__main__':
    asyncio.run(setup_database())