mib-new/database.py

import logging
import sqlite3
from contextlib import closing

from config import DB_PATH

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("database")


def create_connection() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def create_tables() -> None:
    with closing(create_connection()) as conn:
        cursor = conn.cursor()

        cursor.execute(
            '''
            CREATE TABLE IF NOT EXISTS file_ids (
                file_key TEXT PRIMARY KEY,
                file_id TEXT
            )
            '''
        )

        cursor.execute(
            '''
            CREATE TABLE IF NOT EXISTS user_data (
                user_id INTEGER,
                chat_id INTEGER,
                count INTEGER NOT NULL DEFAULT 0,
                start_time REAL NOT NULL DEFAULT 0,
                ban_trigger INTEGER NOT NULL DEFAULT 0,
                ban_time REAL NOT NULL DEFAULT 0,
                username TEXT,
                PRIMARY KEY (user_id, chat_id)
            )
            '''
        )

        cursor.execute(
            '''
            CREATE TABLE IF NOT EXISTS user_data_ban (
                user_id INTEGER,
                chat_id INTEGER,
                command TEXT,
                count INTEGER NOT NULL DEFAULT 0,
                ban_time_inf REAL NOT NULL DEFAULT 0,
                PRIMARY KEY (user_id, chat_id, command)
            )
            '''
        )

        conn.commit()


def get_file_id(file_key: str) -> str | None:
    try:
        with closing(create_connection()) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT file_id FROM file_ids WHERE file_key = ?", (file_key,))
            row = cursor.fetchone()
            return row["file_id"] if row and row["file_id"] else None
    except Exception as exc:
        logger.error("Ошибка при получении file_id: %s", exc)
        return None


def save_file_id(file_key: str, file_id: str | None) -> None:
    try:
        with closing(create_connection()) as conn:
            cursor = conn.cursor()
            if file_id is None:
                cursor.execute("DELETE FROM file_ids WHERE file_key = ?", (file_key,))
            else:
                cursor.execute(
                    "REPLACE INTO file_ids (file_key, file_id) VALUES (?, ?)",
                    (file_key, file_id),
                )
            conn.commit()
    except Exception as exc:
        logger.error("Ошибка при сохранении file_id: %s", exc)


def get_user_data(user_id: int, chat_id: int):
    try:
        with closing(create_connection()) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT * FROM user_data WHERE user_id = ? AND chat_id = ?",
                (user_id, chat_id),
            )
            row = cursor.fetchone()
            return tuple(row) if row else None
    except Exception as exc:
        logger.error("Ошибка при получении данных пользователя: %s", exc)
        return None


def save_user_data(
    user_id: int,
    chat_id: int,
    count: int | None = None,
    start_time: float | None = None,
    ban_trigger: bool | None = None,
    ban_time: float | None = None,
    username: str | None = None,
) -> None:
    try:
        with closing(create_connection()) as conn:
            cursor = conn.cursor()
            existing = get_user_data(user_id, chat_id)

            if existing is None:
                cursor.execute(
                    '''
                    INSERT INTO user_data (user_id, chat_id, count, start_time, ban_trigger, ban_time, username)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    ''',
                    (
                        user_id,
                        chat_id,
                        count or 0,
                        start_time or 0,
                        int(bool(ban_trigger)) if ban_trigger is not None else 0,
                        ban_time or 0,
                        username,
                    ),
                )
            else:
                updates = []
                params = []
                if count is not None:
                    updates.append("count = ?")
                    params.append(count)
                if start_time is not None:
                    updates.append("start_time = ?")
                    params.append(start_time)
                if ban_trigger is not None:
                    updates.append("ban_trigger = ?")
                    params.append(int(bool(ban_trigger)))
                if ban_time is not None:
                    updates.append("ban_time = ?")
                    params.append(ban_time)
                if username is not None:
                    updates.append("username = ?")
                    params.append(username)

                if updates:
                    params.extend([user_id, chat_id])
                    cursor.execute(
                        f"UPDATE user_data SET {', '.join(updates)} WHERE user_id = ? AND chat_id = ?",
                        params,
                    )

            conn.commit()
    except Exception as exc:
        logger.error("Ошибка при сохранении данных пользователя: %s", exc)


def get_user_data_ban(user_id: int, chat_id: int, command: str):
    try:
        with closing(create_connection()) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT * FROM user_data_ban WHERE user_id = ? AND chat_id = ? AND command = ?",
                (user_id, chat_id, command),
            )
            row = cursor.fetchone()
            return tuple(row) if row else None
    except Exception as exc:
        logger.error("Ошибка при получении данных о блокировке пользователя: %s", exc)
        return None


def save_user_data_ban(
    user_id: int,
    chat_id: int,
    command: str,
    count: int | None = None,
    ban_time_inf: float | None = None,
) -> None:
    try:
        with closing(create_connection()) as conn:
            cursor = conn.cursor()
            existing = get_user_data_ban(user_id, chat_id, command)

            if existing is None:
                cursor.execute(
                    '''
                    INSERT INTO user_data_ban (user_id, chat_id, command, count, ban_time_inf)
                    VALUES (?, ?, ?, ?, ?)
                    ''',
                    (user_id, chat_id, command, count or 0, ban_time_inf or 0),
                )
            else:
                updates = []
                params = []
                if count is not None:
                    updates.append("count = ?")
                    params.append(count)
                if ban_time_inf is not None:
                    updates.append("ban_time_inf = ?")
                    params.append(ban_time_inf)

                if updates:
                    params.extend([user_id, chat_id, command])
                    cursor.execute(
                        f"UPDATE user_data_ban SET {', '.join(updates)} WHERE user_id = ? AND chat_id = ? AND command = ?",
                        params,
                    )

            conn.commit()
    except Exception as exc:
        logger.error("Ошибка при сохранении данных о блокировке пользователя: %s", exc)


create_tables()