282 lines
11 KiB
Python
282 lines
11 KiB
Python
import asyncpg
|
|
from typing import List, Optional
|
|
from db.models import SourceChannel, TargetChannel, Post, Setting
|
|
from db.database import get_db_pool
|
|
|
|
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;")
|
|
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) -> 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)
|
|
VALUES ($1, $2, $3, $4)
|
|
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,
|
|
)
|
|
return row["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;")
|
|
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 & Deduplication ---
|
|
async def find_duplicate_post_by_hash(self, content_hash: str) -> Optional[Post]:
|
|
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 ORDER BY id ASC LIMIT 1;",
|
|
content_hash,
|
|
)
|
|
return Post(**dict(row)) if row else None
|
|
|
|
async def find_candidate_posts_by_tags(self, tags: List[str], exclude_post_id: Optional[int] = None, hours_lookback: int = 72, limit: int = 5) -> List[Post]:
|
|
if not tags:
|
|
return []
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
query = """
|
|
SELECT * FROM posts
|
|
WHERE tags && $1::text[]
|
|
AND created_at >= NOW() - ($2 || ' hours')::interval
|
|
AND ($3::bigint IS NULL OR id != $3::bigint)
|
|
ORDER BY created_at DESC
|
|
LIMIT $4;
|
|
"""
|
|
rows = await conn.fetch(query, tags, str(hours_lookback), exclude_post_id, limit)
|
|
return [Post(**dict(r)) for r in rows]
|
|
|
|
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_ai')
|
|
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_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 [Post(**dict(r)) for r in rows]
|
|
|
|
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 Post(**dict(row)) if row else None
|
|
|
|
async def update_post_tags(self, post_id: int, tags: List[str], subject: str) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"UPDATE posts SET tags = $1, subject = $2 WHERE id = $3;",
|
|
tags, subject, post_id,
|
|
)
|
|
|
|
async def update_post_duplicate_status(
|
|
self,
|
|
post_id: int,
|
|
is_duplicate: bool,
|
|
duplicate_of_id: Optional[int] = None,
|
|
similarity_reason: Optional[str] = None,
|
|
) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
UPDATE posts
|
|
SET is_duplicate = $1, duplicate_of_id = $2, similarity_reason = $3
|
|
WHERE id = $4;
|
|
""",
|
|
is_duplicate, duplicate_of_id, similarity_reason, post_id,
|
|
)
|
|
|
|
async def update_ai_result(
|
|
self,
|
|
post_id: int,
|
|
subject: str,
|
|
ai_text: str,
|
|
tags: List[str],
|
|
suggested_target_id: Optional[int] = None,
|
|
is_duplicate: bool = False,
|
|
duplicate_of_id: Optional[int] = None,
|
|
similarity_reason: Optional[str] = None,
|
|
) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
UPDATE posts
|
|
SET subject = $1, ai_text = $2, tags = $3, suggested_target_id = $4,
|
|
is_duplicate = $5, duplicate_of_id = $6, similarity_reason = $7,
|
|
status = 'pending_review'
|
|
WHERE id = $8;
|
|
""",
|
|
subject, ai_text, tags, suggested_target_id, is_duplicate, duplicate_of_id, similarity_reason, post_id,
|
|
)
|
|
|
|
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 approve_post(self, post_id: int, target_channel_id: int) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
UPDATE posts
|
|
SET target_channel_id = $1, status = 'approved'
|
|
WHERE id = $2;
|
|
""",
|
|
target_channel_id, 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 mark_post_published(self, post_id: int) -> None:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"UPDATE posts SET status = 'published', published_at = CURRENT_TIMESTAMP WHERE id = $1;",
|
|
post_id,
|
|
)
|
|
|
|
async def get_next_approved_post_for_target(self, target_id: int) -> Optional[Post]:
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
SELECT * FROM posts
|
|
WHERE target_channel_id = $1 AND status = 'approved'
|
|
ORDER BY id ASC
|
|
LIMIT 1;
|
|
""",
|
|
target_id,
|
|
)
|
|
return Post(**dict(row)) if row else None
|
|
|
|
# --- 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:
|
|
row = await conn.fetchrow("SELECT value FROM settings WHERE key = $1;", key)
|
|
return row["value"] if row 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,
|
|
)
|