103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
from typing import Any, Dict, List
|
|
from app.crawler.base import PipelineConfig, PipelineResult
|
|
from app.services.filter_service import dedup_by_hash
|
|
|
|
|
|
class PostCrawlPipeline:
|
|
def __init__(self, db_session, notification_service):
|
|
self.db = db_session
|
|
self.notify = notification_service
|
|
|
|
async def process(self, announcements: List[Dict[str, Any]],
|
|
config: PipelineConfig) -> PipelineResult:
|
|
result = PipelineResult()
|
|
|
|
if not announcements:
|
|
return result
|
|
|
|
# 1. Dedup
|
|
if config.dedup_enabled:
|
|
announcements = dedup_by_hash(announcements)
|
|
|
|
# 2. Store to database
|
|
stored = await self._save_to_db(announcements)
|
|
result.stored = stored
|
|
|
|
to_notify = announcements
|
|
|
|
# 3. Filter
|
|
if config.filter_enabled and config.keywords:
|
|
before = len(to_notify)
|
|
to_notify = [a for a in to_notify
|
|
if self._match_keywords(a, config.keywords)]
|
|
result.filtered = before - len(to_notify)
|
|
|
|
# 4. Notify
|
|
if config.notify_mode == "all":
|
|
result.notified = await self._send_notifications(to_notify)
|
|
elif config.notify_mode == "filtered":
|
|
if config.filter_enabled and config.keywords:
|
|
result.notified = await self._send_notifications(to_notify)
|
|
elif not config.filter_enabled:
|
|
result.notified = await self._send_notifications(to_notify)
|
|
|
|
# 5. Mark sent
|
|
if config.mark_sent and result.notified > 0:
|
|
await self._mark_sent(to_notify)
|
|
|
|
return result
|
|
|
|
async def _save_to_db(self, announcements: List[Dict[str, Any]]) -> int:
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
from app.models.announcement import Announcement
|
|
|
|
if not announcements:
|
|
return 0
|
|
|
|
values = [{
|
|
"title": a["title"],
|
|
"publish_date": a["publish_date"],
|
|
"purchase_name": a.get("purchase_name", ""),
|
|
"content_url": a.get("content_url", ""),
|
|
"source_code": a["source_code"],
|
|
"source_name": a["source_name"],
|
|
"announcement_type": a.get("announcement_type", "purchase"),
|
|
"content_hash": a["content_hash"],
|
|
"crawl_mode": a.get("crawl_mode", "auto"),
|
|
"is_new": a.get("is_new", True),
|
|
"is_sent": False,
|
|
"keyword_matched": a.get("keyword_matched", False),
|
|
} for a in announcements]
|
|
|
|
stmt = insert(Announcement).values(values)
|
|
stmt = stmt.on_conflict_do_nothing(index_elements=["content_hash"])
|
|
|
|
result_proxy = await self.db.execute(stmt)
|
|
await self.db.commit()
|
|
return result_proxy.rowcount if result_proxy.rowcount >= 0 else len(values)
|
|
|
|
async def _send_notifications(self, announcements: List[Dict[str, Any]]) -> int:
|
|
return await self.notify.send(announcements)
|
|
|
|
async def _mark_sent(self, announcements: List[Dict[str, Any]]) -> int:
|
|
from app.models.announcement import Announcement
|
|
from sqlalchemy import update
|
|
|
|
hashes = [a["content_hash"] for a in announcements if a.get("content_hash")]
|
|
if not hashes:
|
|
return 0
|
|
|
|
stmt = (
|
|
update(Announcement)
|
|
.where(Announcement.content_hash.in_(hashes))
|
|
.values(is_sent=True)
|
|
)
|
|
result = await self.db.execute(stmt)
|
|
await self.db.commit()
|
|
return result.rowcount
|
|
|
|
@staticmethod
|
|
def _match_keywords(announcement: Dict[str, Any], keywords: List[str]) -> bool:
|
|
text = f"{announcement.get('title', '')} {announcement.get('purchase_name', '')}"
|
|
return any(kw in text for kw in keywords)
|