205 lines
7.5 KiB
Python
205 lines
7.5 KiB
Python
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()
|