244 lines
9.9 KiB
Python
244 lines
9.9 KiB
Python
import json
|
|
import asyncpg
|
|
from typing import List, Optional, Dict, Any
|
|
from db.models import SourceChannel, TargetChannel, Post, Setting
|
|
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):
|
|
try:
|
|
data["published_to"] = json.loads(data["published_to"])
|
|
except Exception:
|
|
data["published_to"] = []
|
|
elif data.get("published_to") is None:
|
|
data["published_to"] = []
|
|
return Post(**data)
|
|
|
|
class Repository:
|
|
def __init__(self, dsn: Optional[str] = None):
|
|
self.dsn = dsn
|
|
|
|
async def _get_pool(self) -> asyncpg.Pool:
|
|
if self.dsn:
|
|
return await get_db_pool(self.dsn)
|
|
return await get_db_pool()
|
|
|
|
# --- Source Channels ---
|
|
async def add_source(self, channel_id: int, title: Optional[str] = None, username: Optional[str] = None) -> int:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO sources (channel_id, title, username)
|
|
VALUES ($1, $2, $3)
|
|
ON CONFLICT(channel_id) DO UPDATE SET
|
|
title = EXCLUDED.title,
|
|
username = EXCLUDED.username,
|
|
is_active = TRUE
|
|
RETURNING id;
|
|
""",
|
|
channel_id, title, username,
|
|
)
|
|
return row["id"]
|
|
|
|
async def get_active_sources(self) -> List[SourceChannel]:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
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
|
|
|
|
# --- Target Channels ---
|
|
async def add_target(
|
|
self,
|
|
channel_id: int,
|
|
title: Optional[str] = None,
|
|
username: Optional[str] = None,
|
|
post_interval_min: int = 30,
|
|
personality: str = "",
|
|
custom_footer: str = ""
|
|
) -> int:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO targets (channel_id, title, username, post_interval_min, personality, custom_footer)
|
|
VALUES ($1, $2, $3, $4, $5, $6)
|
|
ON CONFLICT(channel_id) DO UPDATE SET
|
|
title = EXCLUDED.title,
|
|
username = EXCLUDED.username,
|
|
post_interval_min = EXCLUDED.post_interval_min,
|
|
is_active = TRUE
|
|
RETURNING id;
|
|
""",
|
|
channel_id, title, username, post_interval_min, personality, custom_footer,
|
|
)
|
|
return row["id"]
|
|
|
|
async def update_target_personality(self, target_id: int, personality: str, custom_footer: Optional[str] = None) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
if custom_footer is not None:
|
|
await conn.execute(
|
|
"UPDATE targets SET personality = $1, custom_footer = $2 WHERE id = $3;",
|
|
personality, custom_footer, target_id
|
|
)
|
|
else:
|
|
await conn.execute(
|
|
"UPDATE targets SET personality = $1 WHERE id = $2;",
|
|
personality, target_id
|
|
)
|
|
|
|
async def update_target_footer(self, target_id: int, custom_footer: str) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"UPDATE targets SET custom_footer = $1 WHERE id = $2;",
|
|
custom_footer, target_id
|
|
)
|
|
|
|
async def update_target_schedule(
|
|
self,
|
|
target_id: int,
|
|
post_interval_min: Optional[int] = None,
|
|
sleep_start_hour: Optional[int] = None,
|
|
sleep_end_hour: Optional[int] = None,
|
|
is_sleep_enabled: Optional[bool] = None
|
|
) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
target = await self.get_target_by_id(target_id)
|
|
if not target:
|
|
return
|
|
new_interval = post_interval_min if post_interval_min is not None else target.post_interval_min
|
|
new_start = sleep_start_hour if sleep_start_hour is not None else target.sleep_start_hour
|
|
new_end = sleep_end_hour if sleep_end_hour is not None else target.sleep_end_hour
|
|
new_enabled = is_sleep_enabled if is_sleep_enabled is not None else target.is_sleep_enabled
|
|
|
|
await conn.execute(
|
|
"""
|
|
UPDATE targets
|
|
SET post_interval_min = $1, sleep_start_hour = $2, sleep_end_hour = $3, is_sleep_enabled = $4
|
|
WHERE id = $5;
|
|
""",
|
|
new_interval, new_start, new_end, new_enabled, target_id
|
|
)
|
|
|
|
async def get_active_targets(self) -> List[TargetChannel]:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
rows = await conn.fetch("SELECT * FROM targets WHERE is_active = TRUE ORDER BY id ASC;")
|
|
return [TargetChannel(**dict(r)) for r in rows]
|
|
|
|
async def get_target_by_id(self, target_id: int) -> Optional[TargetChannel]:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
row = await conn.fetchrow("SELECT * FROM targets WHERE id = $1;", target_id)
|
|
return TargetChannel(**dict(row)) if row else None
|
|
|
|
async def update_target_last_post(self, target_id: int) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"UPDATE targets SET last_post_time = CURRENT_TIMESTAMP WHERE id = $1;",
|
|
target_id,
|
|
)
|
|
|
|
# --- Posts & Multi-Channel Dispatch ---
|
|
async def create_raw_post(
|
|
self,
|
|
source_channel_id: int,
|
|
source_message_id: int,
|
|
raw_text: Optional[str],
|
|
media_path: Optional[str] = None,
|
|
media_type: Optional[str] = None,
|
|
content_hash: Optional[str] = None,
|
|
is_duplicate: bool = False,
|
|
duplicate_of_id: Optional[int] = None,
|
|
similarity_reason: Optional[str] = None,
|
|
) -> Optional[int]:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
try:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO posts (
|
|
source_channel_id, source_message_id, raw_text, media_path,
|
|
media_type, content_hash, is_duplicate, duplicate_of_id, similarity_reason, status
|
|
)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 'pending_review')
|
|
RETURNING id;
|
|
""",
|
|
source_channel_id,
|
|
source_message_id,
|
|
raw_text,
|
|
media_path,
|
|
media_type,
|
|
content_hash,
|
|
is_duplicate,
|
|
duplicate_of_id,
|
|
similarity_reason,
|
|
)
|
|
return row["id"] if row else None
|
|
except asyncpg.UniqueViolationError:
|
|
return None
|
|
|
|
async def get_post_by_id(self, post_id: int) -> Optional[Post]:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
row = await conn.fetchrow("SELECT * FROM posts WHERE id = $1;", post_id)
|
|
return _parse_post_row(row) if row else None
|
|
|
|
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:
|
|
rows = await conn.fetch(
|
|
"SELECT * FROM posts WHERE status = $1 ORDER BY id ASC LIMIT $2;",
|
|
status, 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:
|
|
await conn.execute(
|
|
"UPDATE posts SET review_message_id = $1 WHERE id = $2;",
|
|
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 reject_post(self, post_id: int) -> None:
|
|
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 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)
|