db: extend schema for categories, channel profiles, and error logs

This commit is contained in:
mamad
2026-08-28 19:33:15 +03:30
parent 8eb5ca0a3f
commit db2e726a57
5 changed files with 889 additions and 23 deletions
+69
View File
@@ -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
+47
View File
@@ -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
+571 -23
View File
@@ -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"
)
+124
View File
@@ -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())
+78
View File
@@ -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())