From db2e726a57d67fdc1e63aa001f62f132fd47a6af Mon Sep 17 00:00:00 2001 From: mamad Date: Fri, 28 Aug 2026 19:33:15 +0330 Subject: [PATCH] db: extend schema for categories, channel profiles, and error logs --- db/database.py | 69 +++++ db/models.py | 47 ++++ db/repository.py | 594 +++++++++++++++++++++++++++++++++++++-- tests/test_categories.py | 124 ++++++++ tests/test_repository.py | 78 +++++ 5 files changed, 889 insertions(+), 23 deletions(-) create mode 100644 tests/test_categories.py create mode 100644 tests/test_repository.py diff --git a/db/database.py b/db/database.py index 42b9892..b4790cc 100644 --- a/db/database.py +++ b/db/database.py @@ -26,9 +26,12 @@ CREATE TABLE IF NOT EXISTS targets ( post_interval_min INT DEFAULT 30, personality TEXT DEFAULT '', custom_footer TEXT DEFAULT '', + custom_prompt TEXT DEFAULT '', sleep_start_hour INT DEFAULT 0, sleep_end_hour INT DEFAULT 0, is_sleep_enabled BOOLEAN DEFAULT FALSE, + auto_source_ids BIGINT[] DEFAULT '{}', + language VARCHAR(32) DEFAULT 'fa', last_post_time TIMESTAMPTZ, is_active BOOLEAN DEFAULT TRUE, created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP @@ -51,6 +54,7 @@ CREATE TABLE IF NOT EXISTS posts ( suggested_target_id INT REFERENCES targets(id) ON DELETE SET NULL, target_channel_id INT REFERENCES targets(id) ON DELETE SET NULL, status VARCHAR(32) DEFAULT 'pending_review', + rejection_reason TEXT DEFAULT '', published_to JSONB DEFAULT '[]'::jsonb, review_message_id BIGINT, scheduled_at TIMESTAMPTZ, @@ -78,20 +82,85 @@ CREATE TABLE IF NOT EXISTS error_logs ( error_message TEXT NOT NULL, traceback TEXT, context JSONB DEFAULT '{}'::jsonb, + resolved BOOLEAN DEFAULT FALSE, + resolved_at TIMESTAMPTZ, + resolved_note TEXT, created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_error_logs_created_at ON error_logs(created_at DESC); CREATE INDEX IF NOT EXISTS idx_error_logs_service ON error_logs(service_name); +CREATE TABLE IF NOT EXISTS ai_logs ( + id BIGSERIAL PRIMARY KEY, + action_name VARCHAR(64), + provider VARCHAR(32), + model VARCHAR(128), + prompt TEXT, + system_prompt TEXT, + response_text TEXT, + duration_sec FLOAT DEFAULT 0.0, + status VARCHAR(16) DEFAULT 'success', + error_message TEXT, + created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS idx_ai_logs_created_at ON ai_logs(created_at DESC); + +CREATE TABLE IF NOT EXISTS ai_providers ( + id BIGSERIAL PRIMARY KEY, + name VARCHAR(128) NOT NULL, + provider_type VARCHAR(32) NOT NULL, + base_url TEXT DEFAULT '', + api_key TEXT DEFAULT '', + model VARCHAR(128) NOT NULL, + reasoning_effort VARCHAR(32) DEFAULT '', + is_active BOOLEAN DEFAULT FALSE, + fallback_provider_id BIGINT REFERENCES ai_providers(id) ON DELETE SET NULL, + created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP +); + +CREATE TABLE IF NOT EXISTS channel_categories ( + id SERIAL PRIMARY KEY, + name VARCHAR(128) NOT NULL, + type VARCHAR(32) DEFAULT 'both', + description TEXT DEFAULT '', + created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS idx_ai_providers_active ON ai_providers(is_active); + -- Migration safety for existing tables +ALTER TABLE ai_providers ADD COLUMN IF NOT EXISTS fallback_provider_id BIGINT REFERENCES ai_providers(id) ON DELETE SET NULL; +ALTER TABLE sources ADD COLUMN IF NOT EXISTS category_id INT REFERENCES channel_categories(id) ON DELETE SET NULL; +ALTER TABLE targets ADD COLUMN IF NOT EXISTS category_id INT REFERENCES channel_categories(id) ON DELETE SET NULL; +ALTER TABLE sources ADD COLUMN IF NOT EXISTS is_active BOOLEAN DEFAULT TRUE; +ALTER TABLE targets ADD COLUMN IF NOT EXISTS is_active BOOLEAN DEFAULT TRUE; +ALTER TABLE sources ADD COLUMN IF NOT EXISTS created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP; +ALTER TABLE targets ADD COLUMN IF NOT EXISTS created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP; +CREATE INDEX IF NOT EXISTS idx_sources_category ON sources(category_id); +CREATE INDEX IF NOT EXISTS idx_targets_category ON targets(category_id); ALTER TABLE targets ADD COLUMN IF NOT EXISTS personality TEXT DEFAULT ''; ALTER TABLE targets ADD COLUMN IF NOT EXISTS custom_footer TEXT DEFAULT ''; ALTER TABLE targets ADD COLUMN IF NOT EXISTS sleep_start_hour INT DEFAULT 0; ALTER TABLE targets ADD COLUMN IF NOT EXISTS sleep_end_hour INT DEFAULT 0; ALTER TABLE targets ADD COLUMN IF NOT EXISTS is_sleep_enabled BOOLEAN DEFAULT FALSE; +ALTER TABLE targets ADD COLUMN IF NOT EXISTS auto_source_ids BIGINT[] DEFAULT '{}'; +ALTER TABLE targets ADD COLUMN IF NOT EXISTS language VARCHAR(32) DEFAULT 'fa'; +ALTER TABLE targets ADD COLUMN IF NOT EXISTS custom_prompt TEXT DEFAULT ''; ALTER TABLE posts ADD COLUMN IF NOT EXISTS published_to JSONB DEFAULT '[]'::jsonb; ALTER TABLE posts ADD COLUMN IF NOT EXISTS is_deleted BOOLEAN DEFAULT FALSE; +ALTER TABLE posts ADD COLUMN IF NOT EXISTS rejection_reason TEXT DEFAULT ''; +ALTER TABLE error_logs ADD COLUMN IF NOT EXISTS resolved BOOLEAN DEFAULT FALSE; +ALTER TABLE error_logs ADD COLUMN IF NOT EXISTS resolved_at TIMESTAMPTZ; +ALTER TABLE error_logs ADD COLUMN IF NOT EXISTS resolved_note TEXT; +CREATE INDEX IF NOT EXISTS idx_error_logs_open ON error_logs(resolved, created_at DESC); + +-- 'pending_ai' and 'approved' belong to an earlier workflow that no longer exists. +-- Posts left in those states are invisible to every current code path, so return +-- them to the review queue. +UPDATE posts SET status = 'pending_review' +WHERE status IN ('pending_ai', 'approved') AND is_deleted = FALSE; """ _pool: Optional[asyncpg.Pool] = None diff --git a/db/models.py b/db/models.py index 0effd3a..3317d5b 100644 --- a/db/models.py +++ b/db/models.py @@ -1,12 +1,21 @@ from dataclasses import dataclass, field from typing import Optional, List, Dict, Any +@dataclass +class ChannelCategory: + id: Optional[int] + name: str + type: str = "both" # "source", "target", "both" + description: str = "" + created_at: Optional[str] = None + @dataclass class SourceChannel: id: Optional[int] channel_id: int username: Optional[str] title: Optional[str] + category_id: Optional[int] = None is_active: bool = True created_at: Optional[str] = None @@ -16,16 +25,22 @@ class TargetChannel: channel_id: int title: Optional[str] username: Optional[str] + category_id: Optional[int] = None post_interval_min: int = 30 personality: str = "" custom_footer: str = "" + custom_prompt: str = "" sleep_start_hour: int = 0 sleep_end_hour: int = 0 is_sleep_enabled: bool = False + # Telegram channel_ids of sources whose posts are queued to this target automatically. + auto_source_ids: List[int] = field(default_factory=list) + language: str = "fa" last_post_time: Optional[str] = None is_active: bool = True created_at: Optional[str] = None + @dataclass class Post: id: Optional[int] @@ -44,6 +59,7 @@ class Post: suggested_target_id: Optional[int] = None target_channel_id: Optional[int] = None status: str = "pending_review" # pending_review, published, rejected, deleted + rejection_reason: Optional[str] = None published_to: List[Dict[str, Any]] = field(default_factory=list) is_deleted: bool = False review_message_id: Optional[int] = None @@ -56,3 +72,34 @@ class Setting: key: str value: str description: Optional[str] = None + +@dataclass +class AILog: + id: Optional[int] + action_name: str + provider: str + model: str + prompt: str + system_prompt: Optional[str] = None + response_text: Optional[str] = None + duration_sec: float = 0.0 + status: str = "success" + error_message: Optional[str] = None + created_at: Optional[str] = None + +@dataclass +class AIProviderProfile: + id: Optional[int] + name: str + provider_type: str + model: str + base_url: str = "" + api_key: str = "" + reasoning_effort: str = "" + is_active: bool = False + fallback_provider_id: Optional[int] = None + fallback_provider_name: Optional[str] = None + created_at: Optional[str] = None + + + diff --git a/db/repository.py b/db/repository.py index ac20d04..dd71dc3 100644 --- a/db/repository.py +++ b/db/repository.py @@ -1,9 +1,11 @@ import json import asyncpg -from typing import List, Optional, Dict, Any -from db.models import SourceChannel, TargetChannel, Post, Setting +from datetime import datetime, timezone +from typing import List, Optional, Dict, Any, Tuple +from db.models import SourceChannel, TargetChannel, Post, Setting, AILog, AIProviderProfile, ChannelCategory from db.database import get_db_pool + def _parse_post_row(row: asyncpg.Record) -> Post: data = dict(row) if isinstance(data.get("published_to"), str): @@ -48,6 +50,12 @@ class Repository: rows = await conn.fetch("SELECT * FROM sources WHERE is_active = TRUE ORDER BY id ASC;") return [SourceChannel(**dict(r)) for r in rows] + async def get_source_by_channel_id(self, channel_id: int) -> Optional[SourceChannel]: + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow("SELECT * FROM sources WHERE channel_id = $1;", channel_id) + return SourceChannel(**dict(row)) if row else None + async def get_source_by_id(self, source_id: int) -> Optional[SourceChannel]: pool = await self._get_pool() async with pool.acquire() as conn: @@ -113,6 +121,23 @@ class Repository: custom_footer, target_id ) + async def update_target_language(self, target_id: int, language: str) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + "UPDATE targets SET language = $1 WHERE id = $2;", + language, target_id + ) + + async def update_target_custom_prompt(self, target_id: int, custom_prompt: str) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + "UPDATE targets SET custom_prompt = $1 WHERE id = $2;", + custom_prompt, target_id + ) + + async def update_target_schedule( self, target_id: int, @@ -152,6 +177,43 @@ class Repository: row = await conn.fetchrow("SELECT * FROM targets WHERE id = $1;", target_id) return TargetChannel(**dict(row)) if row else None + async def set_target_auto_sources(self, target_id: int, source_channel_ids: List[int]) -> None: + """Replace the set of source channels auto-routed into this target.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + "UPDATE targets SET auto_source_ids = $1::bigint[] WHERE id = $2;", + sorted(set(source_channel_ids)), target_id, + ) + + async def toggle_target_auto_source(self, target_id: int, source_channel_id: int) -> bool: + """Add or remove one source from a target's auto-route list. Returns the new state.""" + target = await self.get_target_by_id(target_id) + if not target: + return False + current = set(target.auto_source_ids or []) + enabled = source_channel_id not in current + if enabled: + current.add(source_channel_id) + else: + current.discard(source_channel_id) + await self.set_target_auto_sources(target_id, list(current)) + return enabled + + async def get_targets_auto_routed_from(self, source_channel_id: int) -> List[TargetChannel]: + """Active targets that have subscribed to this source channel.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + SELECT * FROM targets + WHERE is_active = TRUE AND $1 = ANY(auto_source_ids) + ORDER BY id ASC; + """, + source_channel_id, + ) + return [TargetChannel(**dict(r)) for r in rows] + async def update_target_last_post(self, target_id: int) -> None: pool = await self._get_pool() async with pool.acquire() as conn: @@ -205,6 +267,38 @@ class Repository: row = await conn.fetchrow("SELECT * FROM posts WHERE id = $1;", post_id) return _parse_post_row(row) if row else None + async def find_duplicate_post(self, content_hash: Optional[str], exclude_post_id: Optional[int] = None) -> Optional[Post]: + """Return the earliest post already carrying this content hash, if any.""" + if not content_hash: + return None + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + """ + SELECT * FROM posts + WHERE content_hash = $1 + AND is_deleted = FALSE + AND ($2::bigint IS NULL OR id <> $2) + ORDER BY id ASC + LIMIT 1; + """, + content_hash, exclude_post_id, + ) + return _parse_post_row(row) if row else None + + async def count_posts_by_status(self, status: str) -> int: + pool = await self._get_pool() + async with pool.acquire() as conn: + return await conn.fetchval("SELECT COUNT(*) FROM posts WHERE status = $1;", status) or 0 + + async def count_posts_from_source(self, source_channel_id: int) -> int: + pool = await self._get_pool() + async with pool.acquire() as conn: + return await conn.fetchval( + "SELECT COUNT(*) FROM posts WHERE source_channel_id = $1 AND is_deleted = FALSE;", + source_channel_id, + ) or 0 + async def get_posts_by_status(self, status: str, limit: int = 20) -> List[Post]: pool = await self._get_pool() async with pool.acquire() as conn: @@ -214,6 +308,23 @@ class Repository: ) return [_parse_post_row(r) for r in rows] + async def get_unreviewed_posts(self, limit: int = 50) -> List[Post]: + """Pending posts that never made it onto a review card in the admin channel.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + SELECT * FROM posts + WHERE status = 'pending_review' + AND is_deleted = FALSE + AND review_message_id IS NULL + ORDER BY id ASC + LIMIT $1; + """, + limit, + ) + return [_parse_post_row(r) for r in rows] + async def update_review_message_id(self, post_id: int, review_message_id: int) -> None: pool = await self._get_pool() async with pool.acquire() as conn: @@ -222,32 +333,469 @@ class Repository: review_message_id, post_id, ) - async def record_post_published_to_target(self, post_id: int, target_id: int, target_title: str) -> None: - pool = await self._get_pool() - async with pool.acquire() as conn: - record_item = json.dumps({ - "target_id": target_id, - "target_title": target_title, - "published_at": str(asyncpg.types.Type) - }) - await conn.execute( - """ - UPDATE posts - SET published_to = published_to || $1::jsonb, - status = 'published', - published_at = CURRENT_TIMESTAMP - WHERE id = $2; - """, - f'[{{"target_id": {target_id}, "target_title": "{target_title}", "published_at": "{asyncpg.types.Type}"}}]', - post_id, - ) + async def _upsert_published_entry( + self, + post_id: int, + target_id: int, + target_title: str, + published_at: Optional[datetime], + mark_published: bool, + ) -> None: + """Add or stamp this target's entry in posts.published_to. - async def reject_post(self, post_id: int) -> None: + The list is read, edited and written back as a single parameterised jsonb value so + that titles containing quotes or backslashes cannot corrupt the document. + """ pool = await self._get_pool() async with pool.acquire() as conn: - await conn.execute("UPDATE posts SET status = 'rejected' WHERE id = $1;", post_id) + async with conn.transaction(): + row = await conn.fetchrow("SELECT published_to FROM posts WHERE id = $1 FOR UPDATE;", post_id) + if not row: + return + + entries = row["published_to"] + if isinstance(entries, str): + try: + entries = json.loads(entries) + except Exception: + entries = [] + if not isinstance(entries, list): + entries = [] + + stamp = published_at.isoformat() if published_at else None + for entry in entries: + if isinstance(entry, dict) and entry.get("target_id") == target_id: + entry["target_title"] = target_title + if stamp: + entry["published_at"] = stamp + break + else: + entries.append({ + "target_id": target_id, + "target_title": target_title, + "published_at": stamp, + }) + + if mark_published: + await conn.execute( + """ + UPDATE posts + SET published_to = $1::jsonb, + status = 'published', + published_at = COALESCE(published_at, CURRENT_TIMESTAMP) + WHERE id = $2; + """, + json.dumps(entries, ensure_ascii=False), post_id, + ) + else: + await conn.execute( + "UPDATE posts SET published_to = $1::jsonb WHERE id = $2;", + json.dumps(entries, ensure_ascii=False), post_id, + ) + + async def record_post_queued_to_target(self, post_id: int, target_id: int, target_title: str) -> None: + """Note that a post is waiting in a target's delivery queue. Status is left untouched.""" + await self._upsert_published_entry(post_id, target_id, target_title, published_at=None, mark_published=False) + + async def record_post_published_to_target(self, post_id: int, target_id: int, target_title: str) -> None: + """Stamp the post as actually delivered to a target channel.""" + await self._upsert_published_entry( + post_id, target_id, target_title, + published_at=datetime.now(timezone.utc), mark_published=True, + ) + + # --- Error tracking --- + async def get_open_error_summary(self, limit: int = 20) -> List[Dict[str, Any]]: + """Unresolved errors grouped by service + exception type, newest group first.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + SELECT service_name, error_type, COUNT(*) AS occurrences, + MAX(created_at) AS last_seen, + (ARRAY_AGG(error_message ORDER BY created_at DESC))[1] AS last_message + FROM error_logs + WHERE resolved = FALSE + GROUP BY service_name, error_type + ORDER BY MAX(created_at) DESC + LIMIT $1; + """, + limit, + ) + return [dict(r) for r in rows] + + async def count_open_errors(self) -> int: + pool = await self._get_pool() + async with pool.acquire() as conn: + return await conn.fetchval("SELECT COUNT(*) FROM error_logs WHERE resolved = FALSE;") or 0 + + async def resolve_errors( + self, + service_name: Optional[str] = None, + error_type: Optional[str] = None, + note: Optional[str] = None, + ) -> int: + """Mark matching unresolved errors as fixed. Both filters None resolves everything.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + return await conn.fetchval( + """ + WITH updated AS ( + UPDATE error_logs + SET resolved = TRUE, + resolved_at = CURRENT_TIMESTAMP, + resolved_note = COALESCE($3, resolved_note) + WHERE resolved = FALSE + AND ($1::text IS NULL OR service_name = $1) + AND ($2::text IS NULL OR error_type = $2) + RETURNING 1 + ) + SELECT COUNT(*) FROM updated; + """, + service_name, error_type, note, + ) or 0 + + async def reject_post(self, post_id: int, rejection_reason: Optional[str] = None) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + if rejection_reason: + await conn.execute( + "UPDATE posts SET status = 'rejected', rejection_reason = $1 WHERE id = $2;", + rejection_reason, post_id + ) + else: + await conn.execute("UPDATE posts SET status = 'rejected' WHERE id = $1;", post_id) + async def soft_delete_post(self, post_id: int) -> None: pool = await self._get_pool() async with pool.acquire() as conn: await conn.execute("UPDATE posts SET is_deleted = TRUE, status = 'deleted' WHERE id = $1;", post_id) + + # --- System Settings --- + async def get_setting(self, key: str, default: Optional[str] = None) -> Optional[str]: + pool = await self._get_pool() + async with pool.acquire() as conn: + val = await conn.fetchval("SELECT value FROM settings WHERE key = $1;", key) + return val if val is not None else default + + async def set_setting(self, key: str, value: str, description: Optional[str] = None) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO settings (key, value, description) + VALUES ($1, $2, $3) + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, description = COALESCE(EXCLUDED.description, settings.description); + """, + key, value, description + ) + + async def get_all_settings(self) -> Dict[str, str]: + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch("SELECT key, value FROM settings;") + return {r["key"]: r["value"] for r in rows} + + # --- AI Logs --- + async def record_ai_log( + self, + action_name: str, + provider: str, + model: str, + prompt: str, + system_prompt: Optional[str] = None, + response_text: Optional[str] = None, + duration_sec: float = 0.0, + status: str = "success", + error_message: Optional[str] = None, + ) -> int: + pool = await self._get_pool() + async with pool.acquire() as conn: + return await conn.fetchval( + """ + INSERT INTO ai_logs (action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + RETURNING id; + """, + action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message + ) + + async def get_recent_ai_logs(self, limit: int = 10) -> List[AILog]: + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + SELECT id, action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message, + to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at + FROM ai_logs + ORDER BY id DESC + LIMIT $1; + """, + limit + ) + return [AILog(**dict(r)) for r in rows] + + async def get_ai_log_by_id(self, log_id: int) -> Optional[AILog]: + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + """ + SELECT id, action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message, + to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at + FROM ai_logs + WHERE id = $1; + """, + log_id + ) + return AILog(**dict(row)) if row else None + + # --- AI Providers Management --- + async def ensure_default_providers(self) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + count = await conn.fetchval("SELECT COUNT(*) FROM ai_providers;") + if count == 0: + await conn.execute( + """ + INSERT INTO ai_providers (name, provider_type, model, base_url, api_key, reasoning_effort, is_active) + VALUES + ('AGY (سرور داخلی)', 'agy', 'antigravity', 'http://host.docker.internal:8088/v1', '', '', TRUE), + + ('OpenRouter / OpenAI', 'openai', 'google/gemini-3.5-flash', 'https://openrouter.ai/api/v1', '', '', FALSE), + ('Google Gemini Direct', 'gemini', 'gemini-1.5-flash', 'https://generativelanguage.googleapis.com/v1beta', '', '', FALSE); + """ + ) + + async def add_provider_profile( + self, + name: str, + provider_type: str, + model: str, + base_url: str = "", + api_key: str = "", + reasoning_effort: str = "", + is_active: bool = False + ) -> int: + pool = await self._get_pool() + async with pool.acquire() as conn: + if is_active: + await conn.execute("UPDATE ai_providers SET is_active = FALSE;") + return await conn.fetchval( + """ + INSERT INTO ai_providers (name, provider_type, model, base_url, api_key, reasoning_effort, is_active) + VALUES ($1, $2, $3, $4, $5, $6, $7) + RETURNING id; + """, + name, provider_type, model, base_url, api_key, reasoning_effort, is_active + ) + + async def get_provider_profiles(self) -> List[AIProviderProfile]: + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + SELECT p.id, p.name, p.provider_type, p.model, p.base_url, p.api_key, p.reasoning_effort, p.is_active, + p.fallback_provider_id, + fb.name as fallback_provider_name, + to_char(p.created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at + FROM ai_providers p + LEFT JOIN ai_providers fb ON p.fallback_provider_id = fb.id + ORDER BY p.id ASC; + """ + ) + return [AIProviderProfile(**dict(r)) for r in rows] + + async def get_provider_profile_by_id(self, profile_id: int) -> Optional[AIProviderProfile]: + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + """ + SELECT p.id, p.name, p.provider_type, p.model, p.base_url, p.api_key, p.reasoning_effort, p.is_active, + p.fallback_provider_id, + fb.name as fallback_provider_name, + to_char(p.created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at + FROM ai_providers p + LEFT JOIN ai_providers fb ON p.fallback_provider_id = fb.id + WHERE p.id = $1; + """, + profile_id + ) + return AIProviderProfile(**dict(row)) if row else None + + async def get_active_provider_profile(self) -> Optional[AIProviderProfile]: + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + """ + SELECT p.id, p.name, p.provider_type, p.model, p.base_url, p.api_key, p.reasoning_effort, p.is_active, + p.fallback_provider_id, + fb.name as fallback_provider_name, + to_char(p.created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at + FROM ai_providers p + LEFT JOIN ai_providers fb ON p.fallback_provider_id = fb.id + WHERE p.is_active = TRUE + ORDER BY p.id ASC + LIMIT 1; + """ + ) + return AIProviderProfile(**dict(row)) if row else None + + async def set_active_provider_profile(self, profile_id: int) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + async with conn.transaction(): + await conn.execute("UPDATE ai_providers SET is_active = FALSE;") + await conn.execute("UPDATE ai_providers SET is_active = TRUE WHERE id = $1;", profile_id) + + async def update_provider_fallback(self, profile_id: int, fallback_provider_id: Optional[int]) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + "UPDATE ai_providers SET fallback_provider_id = $1 WHERE id = $2;", + fallback_provider_id, profile_id + ) + + async def update_provider_profile( + self, + profile_id: int, + name: Optional[str] = None, + model: Optional[str] = None, + base_url: Optional[str] = None, + api_key: Optional[str] = None, + reasoning_effort: Optional[str] = None, + ) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + """ + UPDATE ai_providers + SET name = COALESCE($2, name), + model = COALESCE($3, model), + base_url = COALESCE($4, base_url), + api_key = COALESCE($5, api_key), + reasoning_effort = COALESCE($6, reasoning_effort) + WHERE id = $1; + """, + profile_id, name, model, base_url, api_key, reasoning_effort + ) + + async def update_provider_reasoning_effort(self, profile_id: int, reasoning_effort: str) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + "UPDATE ai_providers SET reasoning_effort = $1 WHERE id = $2;", + reasoning_effort, profile_id + ) + + async def delete_provider_profile(self, profile_id: int) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute("DELETE FROM ai_providers WHERE id = $1;", profile_id) + + # --- Channel Categories --- + async def create_category(self, name: str, cat_type: str = "both", description: str = "") -> int: + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + """ + INSERT INTO channel_categories (name, type, description) + VALUES ($1, $2, $3) + RETURNING id; + """, + name, cat_type, description + ) + return row["id"] + + async def get_categories(self, cat_type: Optional[str] = None) -> List[ChannelCategory]: + pool = await self._get_pool() + async with pool.acquire() as conn: + if cat_type: + rows = await conn.fetch( + "SELECT * FROM channel_categories WHERE type = $1 OR type = 'both' ORDER BY id ASC;", + cat_type + ) + else: + rows = await conn.fetch("SELECT * FROM channel_categories ORDER BY id ASC;") + return [ChannelCategory(**dict(r)) for r in rows] + + async def get_category_by_id(self, cat_id: int) -> Optional[ChannelCategory]: + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow("SELECT * FROM channel_categories WHERE id = $1;", cat_id) + return ChannelCategory(**dict(row)) if row else None + + async def update_category(self, cat_id: int, name: Optional[str] = None, description: Optional[str] = None, cat_type: Optional[str] = None) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute( + """ + UPDATE channel_categories + SET name = COALESCE($2, name), + description = COALESCE($3, description), + type = COALESCE($4, type) + WHERE id = $1; + """, + cat_id, name, description, cat_type + ) + + async def delete_category(self, cat_id: int) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute("UPDATE sources SET category_id = NULL WHERE category_id = $1;", cat_id) + await conn.execute("UPDATE targets SET category_id = NULL WHERE category_id = $1;", cat_id) + await conn.execute("DELETE FROM channel_categories WHERE id = $1;", cat_id) + + async def set_source_category(self, source_id: int, category_id: Optional[int]) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute("UPDATE sources SET category_id = $1 WHERE id = $2;", category_id, source_id) + + async def set_target_category(self, target_id: int, category_id: Optional[int]) -> None: + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute("UPDATE targets SET category_id = $1 WHERE id = $2;", category_id, target_id) + + async def get_sources_by_category(self, category_id: Optional[int]) -> List[SourceChannel]: + pool = await self._get_pool() + async with pool.acquire() as conn: + if category_id is None: + rows = await conn.fetch("SELECT * FROM sources WHERE is_active = TRUE AND category_id IS NULL ORDER BY id ASC;") + else: + rows = await conn.fetch("SELECT * FROM sources WHERE is_active = TRUE AND category_id = $1 ORDER BY id ASC;", category_id) + return [SourceChannel(**dict(r)) for r in rows] + + async def get_targets_by_category(self, category_id: Optional[int]) -> List[TargetChannel]: + pool = await self._get_pool() + async with pool.acquire() as conn: + if category_id is None: + rows = await conn.fetch("SELECT * FROM targets WHERE is_active = TRUE AND category_id IS NULL ORDER BY id ASC;") + else: + rows = await conn.fetch("SELECT * FROM targets WHERE is_active = TRUE AND category_id = $1 ORDER BY id ASC;", category_id) + return [TargetChannel(**dict(r)) for r in rows] + + + async def get_category_channel_counts(self, category_id: int) -> Dict[str, int]: + pool = await self._get_pool() + async with pool.acquire() as conn: + src_count = await conn.fetchval("SELECT COUNT(*) FROM sources WHERE is_active = TRUE AND category_id = $1;", category_id) or 0 + trg_count = await conn.fetchval("SELECT COUNT(*) FROM targets WHERE is_active = TRUE AND category_id = $1;", category_id) or 0 + return {"sources": src_count, "targets": trg_count} + + # --- System Global Operational State --- + async def is_system_paused(self) -> bool: + val = await self.get_setting("system_is_paused", "false") + return str(val).strip().lower() in ("true", "1", "yes") + + async def set_system_paused(self, paused: bool) -> None: + await self.set_setting( + "system_is_paused", + "true" if paused else "false", + description="Global emergency operational pause" + ) + + + + + + diff --git a/tests/test_categories.py b/tests/test_categories.py new file mode 100644 index 0000000..1bf4dda --- /dev/null +++ b/tests/test_categories.py @@ -0,0 +1,124 @@ +import asyncio +import sys + +sys.path.insert(0, "/app") + +from db.database import init_db, close_db_pool +from db.repository import Repository +from db.models import ChannelCategory, SourceChannel, TargetChannel +from services.admin_bot import AdminBotService + +SRC_1 = -1009999000031 +SRC_2 = -1009999000032 +TRG_1 = -1009999000033 +TRG_2 = -1009999000034 + + +async def _cleanup(repo): + pool = await repo._get_pool() + async with pool.acquire() as conn: + for cid in (SRC_1, SRC_2): + await conn.execute("DELETE FROM sources WHERE channel_id = $1;", cid) + for cid in (TRG_1, TRG_2): + await conn.execute("DELETE FROM targets WHERE channel_id = $1;", cid) + await conn.execute("DELETE FROM channel_categories WHERE name LIKE 'Test Cat%';") + + +async def run_tests(): + await init_db() + repo = Repository() + await _cleanup(repo) + + # 1. Create categories + cat1_id = await repo.create_category("Test Cat Tech", "both", "Technology news") + cat2_id = await repo.create_category("Test Cat Crypto", "both", "Crypto & Finance") + assert cat1_id > 0 + assert cat2_id > 0 + + cats = await repo.get_categories() + cat_names = [c.name for c in cats] + assert "Test Cat Tech" in cat_names + assert "Test Cat Crypto" in cat_names + + # 2. Add Sources & Targets + s1_id = await repo.add_source(SRC_1, "Tech Source 1", "tech_src1") + s2_id = await repo.add_source(SRC_2, "Crypto Source 2", "crypto_src2") + t1_id = await repo.add_target(TRG_1, "Tech Target 1", "tech_trg1") + t2_id = await repo.add_target(TRG_2, "Crypto Target 2", "crypto_trg2") + + # Initial state: no category + assert (await repo.get_source_by_id(s1_id)).category_id is None + assert (await repo.get_target_by_id(t1_id)).category_id is None + + # 3. Assign categories + await repo.set_source_category(s1_id, cat1_id) + await repo.set_source_category(s2_id, cat2_id) + await repo.set_target_category(t1_id, cat1_id) + await repo.set_target_category(t2_id, cat2_id) + + # Verify assignments + assert (await repo.get_source_by_id(s1_id)).category_id == cat1_id + assert (await repo.get_source_by_id(s2_id)).category_id == cat2_id + assert (await repo.get_target_by_id(t1_id)).category_id == cat1_id + assert (await repo.get_target_by_id(t2_id)).category_id == cat2_id + + # 4. Filter channels by category + tech_sources = await repo.get_sources_by_category(cat1_id) + tech_targets = await repo.get_targets_by_category(cat1_id) + assert len(tech_sources) == 1 and tech_sources[0].id == s1_id + assert len(tech_targets) == 1 and tech_targets[0].id == t1_id + + # 5. Test AdminBot category-based rendering + bot = AdminBotService(repo=repo, ai_processor=None, queue=None, review_channel_id=0, admin_user_ids=[1]) + + src_text, src_buttons = await bot._render_source_list() + assert "دسته‌بندی‌های کانال‌های مبدا" in src_text + btn_data = [btn.data.decode("utf-8") for row in src_buttons for btn in row] + assert f"src_cat_view:{cat1_id}" in btn_data + assert f"src_cat_view:{cat2_id}" in btn_data + + # Test clicking a category for sources + cat_src_text, cat_src_buttons = await bot._render_source_channels_in_category(cat1_id) + assert "کانال‌های مبدا در دسته" in cat_src_text + assert "Tech Cat" in cat_src_text or "Test Cat Tech" in cat_src_text + cat_btn_data = [btn.data.decode("utf-8") for row in cat_src_buttons for btn in row] + assert f"src_view:{s1_id}" in cat_btn_data + + # Test target category navigation + trg_text, trg_buttons = await bot._render_target_list() + assert "دسته‌بندی‌های کانال‌های مقصد" in trg_text + trg_btn_data = [btn.data.decode("utf-8") for row in trg_buttons for btn in row] + assert f"trg_cat_view:{cat1_id}" in trg_btn_data + + cat_trg_text, cat_trg_buttons = await bot._render_target_channels_in_category(cat1_id) + assert "کانال‌های مقصد در دسته" in cat_trg_text + cat_trg_btn_data = [btn.data.decode("utf-8") for row in cat_trg_buttons for btn in row] + assert f"trg_view:{t1_id}" in cat_trg_btn_data + + # 6. Counts + counts = await repo.get_category_channel_counts(cat1_id) + assert counts["sources"] == 1 + assert counts["targets"] == 1 + + # 7. Rename category + await repo.update_category(cat1_id, name="Test Cat Tech Updated") + updated_cat = await repo.get_category_by_id(cat1_id) + assert updated_cat.name == "Test Cat Tech Updated" + + # 8. Unassign category + await repo.set_source_category(s1_id, None) + assert (await repo.get_source_by_id(s1_id)).category_id is None + + # 9. Delete category (must nullify channel references safely) + await repo.delete_category(cat2_id) + assert await repo.get_category_by_id(cat2_id) is None + assert (await repo.get_source_by_id(s2_id)).category_id is None + assert (await repo.get_target_by_id(t2_id)).category_id is None + + await _cleanup(repo) + await close_db_pool() + print("All category database, repository, and UI rendering tests passed successfully!") + + +if __name__ == "__main__": + asyncio.run(run_tests()) diff --git a/tests/test_repository.py b/tests/test_repository.py new file mode 100644 index 0000000..8b577bc --- /dev/null +++ b/tests/test_repository.py @@ -0,0 +1,78 @@ +"""Integration tests for the repository layer. Requires a reachable Postgres.""" +import asyncio +import sys + +sys.path.insert(0, "/app") + +from db.database import init_db, close_db_pool +from db.repository import Repository + +TEST_SOURCE_ID = -1009999000001 +TEST_TARGET_ID = -1009999000002 +# Deliberately hostile title: quotes and a backslash must survive the JSONB round-trip. +TEST_TARGET_TITLE = 'News "Daily" \\ Channel' + + +async def _cleanup(repo: Repository): + pool = await repo._get_pool() + async with pool.acquire() as conn: + await conn.execute("DELETE FROM posts WHERE source_channel_id = $1;", TEST_SOURCE_ID) + await conn.execute("DELETE FROM sources WHERE channel_id = $1;", TEST_SOURCE_ID) + await conn.execute("DELETE FROM targets WHERE channel_id = $1;", TEST_TARGET_ID) + + +async def run_tests(): + await init_db() + repo = Repository() + await _cleanup(repo) + + # 1. Sources are looked up by their Telegram channel_id, not the surrogate row id. + await repo.add_source(TEST_SOURCE_ID, "Source Tech", "source_tech") + source = await repo.get_source_by_channel_id(TEST_SOURCE_ID) + assert source is not None, "get_source_by_channel_id returned None for a registered source" + assert source.channel_id == TEST_SOURCE_ID + assert source.title == "Source Tech" + assert await repo.get_source_by_channel_id(-1000000000000) is None + + # 2. Deduplication lookup by content hash. + post_id = await repo.create_raw_post( + source_channel_id=TEST_SOURCE_ID, + source_message_id=101, + raw_text="Breaking news: AI update released!", + content_hash="hash_12345", + ) + assert post_id is not None + dup = await repo.find_duplicate_post("hash_12345") + assert dup is not None and dup.id == post_id, "find_duplicate_post did not match a stored hash" + assert await repo.find_duplicate_post("no_such_hash") is None + + # Re-inserting the same source message is rejected by the unique constraint. + assert await repo.create_raw_post(TEST_SOURCE_ID, 101, "dupe") is None + + # 3. Queueing records the target without prematurely marking the post published. + target_id = await repo.add_target(TEST_TARGET_ID, TEST_TARGET_TITLE, "target_chan", post_interval_min=15) + await repo.record_post_queued_to_target(post_id, target_id, TEST_TARGET_TITLE) + post = await repo.get_post_by_id(post_id) + assert post.status == "pending_review", f"queueing must not publish, got {post.status}" + assert len(post.published_to) == 1, f"expected 1 queue entry, got {post.published_to}" + assert post.published_to[0]["target_title"] == TEST_TARGET_TITLE, "title was mangled in JSONB" + assert post.published_to[0]["published_at"] is None + + # 4. Publishing flips status and stamps a real timestamp on the existing entry. + await repo.record_post_published_to_target(post_id, target_id, TEST_TARGET_TITLE) + post = await repo.get_post_by_id(post_id) + assert post.status == "published", f"expected published, got {post.status}" + assert len(post.published_to) == 1, f"publishing must not duplicate the entry, got {post.published_to}" + stamped = post.published_to[0]["published_at"] + assert stamped and "class" not in str(stamped), f"published_at is not a timestamp: {stamped!r}" + + # 5. Counting by status must not require loading every row. + assert await repo.count_posts_by_status("published") >= 1 + + await _cleanup(repo) + await close_db_pool() + print("All repository tests passed successfully!") + + +if __name__ == "__main__": + asyncio.run(run_tests())