redairsoft_vk_gt_max_sender/src/database.py

from __future__ import annotations
from loguru import logger

import json
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Optional
import aiosqlite
try:
    from .config import settings
except (ImportError, ValueError):
    from config import settings


# Each method opens its own short-lived connection (no pooling). busy_timeout
# means a second connection opened concurrently (e.g. inspecting the DB by
# hand with sqlite3 while the service runs) waits instead of immediately
# failing with "database is locked".
_BUSY_TIMEOUT_MS = 5000

# route_id used to backfill rows written before multi-route support existed,
# and the id load_routes() assigns to the auto-generated legacy single route -
# keeping these in sync means an upgrade from the old single-group deployment
# doesn't lose "already published" history for that group.
_LEGACY_ROUTE_ID = "default"


class Database:
    def __init__(self, db_path: Optional[str] = None) -> None:
        self.db_path = str(settings.db_path if db_path is None else db_path)

    @asynccontextmanager
    async def _connect(self) -> AsyncIterator[aiosqlite.Connection]:
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(f"PRAGMA busy_timeout = {_BUSY_TIMEOUT_MS};")
            yield db

    async def _table_columns(self, db: aiosqlite.Connection, table: str) -> set[str]:
        cursor = await db.execute(f"PRAGMA table_info({table})")
        rows = await cursor.fetchall()
        return {row[1] for row in rows}

    async def _ensure_column(self, db: aiosqlite.Connection, table: str, column: str, ddl: str) -> None:
        cols = await self._table_columns(db, table)
        if column not in cols:
            await db.execute(f"ALTER TABLE {table} ADD COLUMN {ddl}")

    async def _migrate_posts_table(self, db: aiosqlite.Connection) -> None:
        """Adds route_id to a pre-multi-route posts table and rebuilds the UNIQUE
        constraint to (route_id, vk_owner_id, vk_post_id). A plain ALTER TABLE ADD
        COLUMN can't change a table-level UNIQUE constraint in SQLite, so this does
        the standard rename/recreate/copy/drop dance. Existing rows are backfilled
        with _LEGACY_ROUTE_ID, matching the route id routes.load_routes() assigns
        to the auto-generated single route - so "already published" history for
        the pre-existing group survives the upgrade.
        """
        cols = await self._table_columns(db, "posts")
        if "route_id" in cols:
            return

        logger.info("Migrating 'posts' table: adding route_id (backfilled as '{}')", _LEGACY_ROUTE_ID)
        await db.execute("ALTER TABLE posts RENAME TO posts_old;")
        await db.execute(
            """
            CREATE TABLE posts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                route_id TEXT NOT NULL DEFAULT 'default',
                vk_post_id INTEGER NOT NULL,
                vk_owner_id INTEGER NOT NULL,
                posted_at INTEGER,
                text TEXT,
                raw_json TEXT,
                tg_status TEXT DEFAULT 'pending',
                tg_message_ids TEXT,
                tg_url TEXT,
                tg_error TEXT,
                max_status TEXT DEFAULT 'pending',
                max_message_ids TEXT,
                max_url TEXT,
                max_error TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                published_at TIMESTAMP,
                UNIQUE(route_id, vk_owner_id, vk_post_id)
            );
            """
        )
        await db.execute(
            f"""
            INSERT INTO posts (
                id, route_id, vk_post_id, vk_owner_id, posted_at, text, raw_json,
                tg_status, tg_message_ids, tg_url, tg_error,
                max_status, max_message_ids, max_url, max_error,
                created_at, published_at
            )
            SELECT id, '{_LEGACY_ROUTE_ID}', vk_post_id, vk_owner_id, posted_at, text, raw_json,
                   tg_status, tg_message_ids, tg_url, tg_error,
                   max_status, max_message_ids, max_url, max_error,
                   created_at, published_at
            FROM posts_old;
            """
        )
        await db.execute("DROP TABLE posts_old;")
        logger.info("Migration of 'posts' table complete.")

    async def init(self) -> None:
        logger.info("Initializing database at {}", self.db_path)
        async with self._connect() as db:
            await db.execute("PRAGMA journal_mode=WAL;")
            await db.execute(
                """
                CREATE TABLE IF NOT EXISTS posts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    route_id TEXT NOT NULL DEFAULT 'default',
                    vk_post_id INTEGER NOT NULL,
                    vk_owner_id INTEGER NOT NULL,
                    posted_at INTEGER,
                    text TEXT,
                    raw_json TEXT,
                    tg_status TEXT DEFAULT 'pending',
                    tg_message_ids TEXT,
                    tg_url TEXT,
                    tg_error TEXT,
                    max_status TEXT DEFAULT 'pending',
                    max_message_ids TEXT,
                    max_url TEXT,
                    max_error TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    published_at TIMESTAMP,
                    UNIQUE(route_id, vk_owner_id, vk_post_id)
                );
                """
            )
            await self._migrate_posts_table(db)

            await db.execute(
                """
                CREATE TABLE IF NOT EXISTS publication_runs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    checked_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    found_count INTEGER DEFAULT 0,
                    published_tg_count INTEGER DEFAULT 0,
                    published_max_count INTEGER DEFAULT 0,
                    status TEXT DEFAULT 'ok',
                    error TEXT
                );
                """
            )
            await self._ensure_column(db, "publication_runs", "route_id", "route_id TEXT")
            await self._ensure_column(db, "publication_runs", "route_name", "route_name TEXT")
            await db.commit()

    async def has_any_posts(self, route_id: str, owner_id: int) -> bool:
        async with self._connect() as db:
            cursor = await db.execute(
                "SELECT 1 FROM posts WHERE route_id = ? AND vk_owner_id = ? LIMIT 1",
                (route_id, owner_id),
            )
            row = await cursor.fetchone()
            return bool(row)

    async def mark_post_skipped(
        self,
        route_id: str,
        owner_id: int,
        post_id: int,
        posted_at: int,
        text: str,
        raw_data: dict[str, Any],
        reason: str = "bootstrap_initial_skip",
    ) -> None:
        raw_json = json.dumps(raw_data, ensure_ascii=False)
        async with self._connect() as db:
            await db.execute(
                """
                INSERT INTO posts (route_id, vk_owner_id, vk_post_id, posted_at, text, raw_json, tg_status, max_status, tg_error, max_error)
                VALUES (?, ?, ?, ?, ?, ?, 'skipped', 'skipped', ?, ?)
                ON CONFLICT(route_id, vk_owner_id, vk_post_id) DO UPDATE SET
                    tg_status = 'skipped',
                    max_status = 'skipped';
                """,
                (route_id, owner_id, post_id, posted_at, text, raw_json, reason, reason),
            )
            await db.commit()

    async def is_post_processed(self, route_id: str, owner_id: int, post_id: int) -> bool:
        async with self._connect() as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute(
                "SELECT tg_status, max_status FROM posts WHERE route_id = ? AND vk_owner_id = ? AND vk_post_id = ?",
                (route_id, owner_id, post_id),
            )
            row = await cursor.fetchone()
            if not row:
                return False
            return bool(
                row["tg_status"] in ("published", "skipped")
                and row["max_status"] in ("published", "skipped")
            )

    async def get_post(self, route_id: str, owner_id: int, post_id: int) -> Optional[dict[str, Any]]:
        async with self._connect() as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute(
                "SELECT * FROM posts WHERE route_id = ? AND vk_owner_id = ? AND vk_post_id = ?",
                (route_id, owner_id, post_id),
            )
            row = await cursor.fetchone()
            return dict(row) if row else None

    async def save_or_update_post(
        self,
        route_id: str,
        owner_id: int,
        post_id: int,
        posted_at: int,
        text: str,
        raw_data: dict[str, Any],
    ) -> int:
        raw_json = json.dumps(raw_data, ensure_ascii=False)
        async with self._connect() as db:
            cursor = await db.execute(
                """
                INSERT INTO posts (route_id, vk_owner_id, vk_post_id, posted_at, text, raw_json)
                VALUES (?, ?, ?, ?, ?, ?)
                ON CONFLICT(route_id, vk_owner_id, vk_post_id) DO UPDATE SET
                    text = excluded.text,
                    raw_json = excluded.raw_json
                RETURNING id;
                """,
                (route_id, owner_id, post_id, posted_at, text, raw_json),
            )
            row = await cursor.fetchone()
            await db.commit()
            return int(row[0]) if row else 0

    async def update_tg_result(
        self,
        post_db_id: int,
        status: str,
        message_ids: Optional[list[int]] = None,
        url: Optional[str] = None,
        error: Optional[str] = None,
    ) -> None:
        msg_str = ",".join(str(m) for m in message_ids) if message_ids else None
        async with self._connect() as db:
            await db.execute(
                """
                UPDATE posts
                SET tg_status = ?,
                    tg_message_ids = COALESCE(?, tg_message_ids),
                    tg_url = COALESCE(?, tg_url),
                    tg_error = ?,
                    published_at = CASE WHEN ? = 'published' THEN CURRENT_TIMESTAMP ELSE published_at END
                WHERE id = ?;
                """,
                (status, msg_str, url, error, status, post_db_id),
            )
            await db.commit()

    async def update_max_result(
        self,
        post_db_id: int,
        status: str,
        message_ids: Optional[list[str]] = None,
        url: Optional[str] = None,
        error: Optional[str] = None,
    ) -> None:
        msg_str = ",".join(message_ids) if message_ids else None
        async with self._connect() as db:
            await db.execute(
                """
                UPDATE posts
                SET max_status = ?,
                    max_message_ids = COALESCE(?, max_message_ids),
                    max_url = COALESCE(?, max_url),
                    max_error = ?,
                    published_at = CASE WHEN ? = 'published' THEN CURRENT_TIMESTAMP ELSE published_at END
                WHERE id = ?;
                """,
                (status, msg_str, url, error, status, post_db_id),
            )
            await db.commit()

    async def record_run(
        self,
        route_id: str,
        route_name: str,
        found_count: int,
        tg_count: int,
        max_count: int,
        status: str = "ok",
        error: Optional[str] = None,
    ) -> None:
        async with self._connect() as db:
            await db.execute(
                """
                INSERT INTO publication_runs (route_id, route_name, found_count, published_tg_count, published_max_count, status, error)
                VALUES (?, ?, ?, ?, ?, ?, ?);
                """,
                (route_id, route_name, found_count, tg_count, max_count, status, error),
            )
            await db.commit()