Initial commit for RedAirsoft VK to TG and MAX poster
This commit is contained in:
+204
@@ -0,0 +1,204 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional
|
||||
import aiosqlite
|
||||
from loguru import logger
|
||||
from .config import settings
|
||||
|
||||
|
||||
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)
|
||||
|
||||
async def init(self) -> None:
|
||||
logger.info("Initializing database at {}", self.db_path)
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("PRAGMA journal_mode=WAL;")
|
||||
await db.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS posts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
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(vk_owner_id, vk_post_id)
|
||||
);
|
||||
"""
|
||||
)
|
||||
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 db.commit()
|
||||
|
||||
async def has_any_posts(self, owner_id: int) -> bool:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
cursor = await db.execute(
|
||||
"SELECT 1 FROM posts WHERE vk_owner_id = ? LIMIT 1",
|
||||
(owner_id,),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
return bool(row)
|
||||
|
||||
async def mark_post_skipped(
|
||||
self,
|
||||
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 aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"""
|
||||
INSERT INTO posts (vk_owner_id, vk_post_id, posted_at, text, raw_json, tg_status, max_status, tg_error, max_error)
|
||||
VALUES (?, ?, ?, ?, ?, 'skipped', 'skipped', ?, ?)
|
||||
ON CONFLICT(vk_owner_id, vk_post_id) DO UPDATE SET
|
||||
tg_status = 'skipped',
|
||||
max_status = 'skipped';
|
||||
""",
|
||||
(owner_id, post_id, posted_at, text, raw_json, reason, reason),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def is_post_processed(self, owner_id: int, post_id: int) -> bool:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
cursor = await db.execute(
|
||||
"SELECT id, tg_status, max_status FROM posts WHERE vk_owner_id = ? AND vk_post_id = ?",
|
||||
(owner_id, post_id),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
if not row:
|
||||
return False
|
||||
# If already published on both or marked skipped, it's processed
|
||||
return bool(
|
||||
row["tg_status"] in ("published", "skipped")
|
||||
and row["max_status"] in ("published", "skipped")
|
||||
)
|
||||
|
||||
async def get_post(self, owner_id: int, post_id: int) -> Optional[dict[str, Any]]:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
cursor = await db.execute(
|
||||
"SELECT * FROM posts WHERE vk_owner_id = ? AND vk_post_id = ?",
|
||||
(owner_id, post_id),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
async def save_or_update_post(
|
||||
self,
|
||||
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 aiosqlite.connect(self.db_path) as db:
|
||||
cursor = await db.execute(
|
||||
"""
|
||||
INSERT INTO posts (vk_owner_id, vk_post_id, posted_at, text, raw_json)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON CONFLICT(vk_owner_id, vk_post_id) DO UPDATE SET
|
||||
text = excluded.text,
|
||||
raw_json = excluded.raw_json
|
||||
RETURNING 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 aiosqlite.connect(self.db_path) 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 aiosqlite.connect(self.db_path) 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,
|
||||
found_count: int,
|
||||
tg_count: int,
|
||||
max_count: int,
|
||||
status: str = "ok",
|
||||
error: Optional[str] = None,
|
||||
) -> None:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"""
|
||||
INSERT INTO publication_runs (found_count, published_tg_count, published_max_count, status, error)
|
||||
VALUES (?, ?, ?, ?, ?);
|
||||
""",
|
||||
(found_count, tg_count, max_count, status, error),
|
||||
)
|
||||
await db.commit()
|
||||
Reference in New Issue
Block a user