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()