telegram-helper-bot/scripts/add_is_anonymous_column.py

#!/usr/bin/env python3
"""
Скрипт миграции для добавления колонки is_anonymous в таблицу post_from_telegram_suggest.
Для существующих записей определяет is_anonymous на основе текста или устанавливает NULL.
"""
import argparse
import asyncio
import os
import sys
from pathlib import Path

project_root = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(project_root))

import aiosqlite

from logs.custom_logger import logger
from helper_bot.utils.helper_func import determine_anonymity

DEFAULT_DB_PATH = "database/tg-bot-database.db"


def _column_exists(rows: list, name: str) -> bool:
    """PRAGMA table_info returns (cid, name, type, notnull, dflt_value, pk)."""
    for row in rows:
        if row[1] == name:
            return True
    return False


async def main(db_path: str) -> None:
    db_path = os.path.abspath(db_path)
    if not os.path.exists(db_path):
        logger.error("База данных не найдена: %s", db_path)
        print(f"Ошибка: база данных не найдена: {db_path}")
        return

    async with aiosqlite.connect(db_path) as conn:
        await conn.execute("PRAGMA foreign_keys = ON")

        # Проверяем наличие колонки is_anonymous
        cursor = await conn.execute(
            "PRAGMA table_info(post_from_telegram_suggest)"
        )
        rows = await cursor.fetchall()
        await cursor.close()

        if not _column_exists(rows, "is_anonymous"):
            logger.info("Добавление колонки is_anonymous в post_from_telegram_suggest")
            await conn.execute(
                "ALTER TABLE post_from_telegram_suggest "
                "ADD COLUMN is_anonymous INTEGER"
            )
            await conn.commit()
            print("Колонка is_anonymous добавлена.")
        else:
            print("Колонка is_anonymous уже существует.")

        # Получаем все записи с текстом для обновления
        cursor = await conn.execute(
            "SELECT message_id, text FROM post_from_telegram_suggest WHERE text IS NOT NULL"
        )
        posts = await cursor.fetchall()
        await cursor.close()

        updated_count = 0
        null_count = 0

        # Обновляем каждую запись
        for message_id, text in posts:
            try:
                # Определяем is_anonymous на основе текста
                # Если текст пустой или None, устанавливаем NULL (legacy)
                if not text or not text.strip():
                    is_anonymous = None
                else:
                    is_anonymous = determine_anonymity(text)
                
                # Преобразуем bool в int для SQLite (True -> 1, False -> 0, None -> None)
                is_anonymous_int = None if is_anonymous is None else (1 if is_anonymous else 0)
                
                await conn.execute(
                    "UPDATE post_from_telegram_suggest SET is_anonymous = ? WHERE message_id = ?",
                    (is_anonymous_int, message_id)
                )
                
                if is_anonymous is not None:
                    updated_count += 1
                else:
                    null_count += 1
                    
            except Exception as e:
                logger.error(f"Ошибка при обработке поста message_id={message_id}: {e}")
                # В случае ошибки устанавливаем NULL
                await conn.execute(
                    "UPDATE post_from_telegram_suggest SET is_anonymous = NULL WHERE message_id = ?",
                    (message_id,)
                )
                null_count += 1

        # Обновляем записи без текста (устанавливаем NULL)
        cursor = await conn.execute(
            "SELECT COUNT(*) FROM post_from_telegram_suggest WHERE text IS NULL"
        )
        row = await cursor.fetchone()
        posts_without_text = row[0] if row else 0
        await cursor.close()

        if posts_without_text > 0:
            await conn.execute(
                "UPDATE post_from_telegram_suggest SET is_anonymous = NULL WHERE text IS NULL"
            )
            null_count += posts_without_text

        await conn.commit()

        total_updated = updated_count + null_count
        logger.info(
            f"Миграция завершена. Обновлено записей: {total_updated} "
            f"(определено: {updated_count}, установлено NULL: {null_count})"
        )
        print(f"Миграция завершена.")
        print(f"Обновлено записей: {total_updated}")
        print(f"  - Определено is_anonymous: {updated_count}")
        print(f"  - Установлено NULL: {null_count}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Добавление колонки is_anonymous в post_from_telegram_suggest"
    )
    parser.add_argument(
        "--db",
        default=os.environ.get("DB_PATH", DEFAULT_DB_PATH),
        help="Путь к БД (или DB_PATH)",
    )
    args = parser.parse_args()
    asyncio.run(main(args.db))