refactor: redesign notification module

.dockerignore

@@ -38,4 +38,5 @@ config
 .pytest_cache
 test
 .env
-test.py
+test.py
+/backend/src/test

.github/workflows/build.yml

@@ -20,8 +20,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          if [ -f backend/requirements.txt ]; then pip install -r backend/requirements.txt; fi
-          pip install pytest
+          if [ -f backend/requirements-dev.txt ]; then pip install -r backend/requirements-dev.txt; fi
       - name: Test
         working-directory: ./backend/src
         run: |

.gitignore

@@ -212,8 +212,3 @@ dev-dist
 
 # vitepress
 /docs/.vitepress/cache/
-
-
-# test file
-test.*
-test_*

backend/requirements-dev.txt

@@ -2,4 +2,10 @@
 ruff
 black
 pre-commit
-pytest
+
+# testing
+pytest          # testing framework
+pytest-httpx    # httpx mock testing
+pytest-asyncio  # asyncio testing
+pytest-cov      # testing coverage
+pytest-mock     # mock wrapper for pytest

backend/requirements.txt

@@ -27,3 +27,6 @@ sqlmodel==0.0.8
 sse-starlette==1.6.5
 semver==3.0.1
 openai==0.28.1
+tenacity==8.2.3
+litequeue==0.7
+httpx==0.25.0

backend/src/module/api/__init__.py

@@ -4,6 +4,7 @@
 from .bangumi import router as bangumi_router
 from .config import router as config_router
 from .log import router as log_router
+from .notification import router as notification_router
 from .program import router as program_router
 from .rss import router as rss_router
 from .search import router as search_router
@@ -19,3 +20,4 @@
 v1.include_router(config_router)
 v1.include_router(rss_router)
 v1.include_router(search_router)
+v1.include_router(notification_router)

backend/src/module/api/notification.py

@@ -0,0 +1,115 @@
+from fastapi import APIRouter, Body, Depends, HTTPException, Query
+
+from module.conf.config import settings
+from module.models.api import NotificationMessageIds
+from module.notification import Notifier
+from module.notification.base import NotificationContent
+from module.security.api import get_current_user
+
+router = APIRouter(
+    prefix="/notification",
+    tags=["notification"],
+    dependencies=[Depends(get_current_user)],
+)
+
+
+def get_notifier():
+    yield Notifier(
+        service_name=settings.notification.type,
+        config=settings.notification.dict(by_alias=True),
+    )
+
+
+@router.get("/total")
+async def get_total_notification(notifier: Notifier = Depends(get_notifier)):
+    cursor = notifier.q.conn.cursor()
+    stmt = """SELECT COUNT(*) FROM Queue WHERE status=0"""
+
+    try:
+        total = cursor.execute(stmt).fetchone()[0]
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return dict(code=0, msg="success", data=dict(total=total))
+
+
+@router.get("")
+async def get_notification(
+    page: int = Query(1, ge=1),
+    limit: int = Query(20, ge=10, le=20, description="max limit is 20 per page"),
+    notifier: Notifier = Depends(get_notifier),
+):
+    cursor = notifier.q.conn.cursor()
+    stmt = r"""
+    SELECT message_id, data, in_time as datetime, status as has_read
+    FROM Queue
+    WHERE status=0
+    ORDER BY in_time DESC
+    """
+
+    offset = (page - 1) * limit
+    stmt += f"LIMIT {limit} OFFSET {offset}"
+
+    try:
+        rows = cursor.execute(stmt).fetchall()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    if not rows:
+        return dict(code=0, msg="success", data=dict(total=0, messages=[]))
+
+    messages = [
+        dict(message_id=data[0], data=data[1], datetime=data[2]) for data in rows
+    ]
+
+    return dict(
+        code=0, msg="success", data=dict(total=len(messages), messages=messages)
+    )
+
+
+@router.post("/read")
+async def set_notification_read(
+    body: NotificationMessageIds,
+    notifier: Notifier = Depends(get_notifier),
+):
+    try:
+        for message_id in body.message_ids:
+            notifier.q.done(message_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return dict(code=0, msg="success")
+
+
+@router.post("/send", description="send notification only for test")
+async def send_notification(
+    content: NotificationContent, notifier: Notifier = Depends(get_notifier)
+):
+    try:
+        await notifier.asend(content=content)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return dict(code=0, msg="success")
+
+
+@router.get("/clean")
+async def clean_notification(notifier: Notifier = Depends(get_notifier)):
+    try:
+        notifier.q.prune()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return dict(code=0, msg="success")
+
+
+# Note: put this vildcard route at the end of all routes to avoid masking other routes
+@router.get("/get")
+async def get_notification_by_id(
+    message_id: str, notifier: Notifier = Depends(get_notifier)
+):
+    message = notifier.q.get(message_id)
+    if not message:
+        raise HTTPException(status_code=404, detail="message not found")
+
+    return dict(code=0, msg="success", data=message)

backend/src/module/conf/const.py

@@ -1,4 +1,5 @@
 # -*- encoding: utf-8 -*-
+import pathlib
 from urllib.parse import parse_qs, urlparse
 
 DEFAULT_SETTINGS = {
@@ -116,3 +117,6 @@ def _(color: str, *args: str) -> str:
     ENDC = "\033[0m"
     BOLD = "\033[1m"
     UNDERLINE = "\033[4m"
+
+
+ROOT = pathlib.Path(__file__).parent.parent.parent.absolute()

backend/src/module/core/sub_thread.py

@@ -1,10 +1,10 @@
 import threading
-import time
+from concurrent.futures import ThreadPoolExecutor
 
 from module.conf import settings
 from module.downloader import DownloadClient
 from module.manager import Renamer, eps_complete
-from module.notification import PostNotification
+from module.notification import Notifier
 from module.rss import RSSAnalyser, RSSEngine
 
 from .status import ProgramStatus
@@ -58,11 +58,17 @@ def rename_loop(self):
         while not self.stop_event.is_set():
             with Renamer() as renamer:
                 renamed_info = renamer.rename()
+
             if settings.notification.enable:
-                with PostNotification() as notifier:
-                    for info in renamed_info:
-                        notifier.send_msg(info)
-                        time.sleep(2)
+                notifier = Notifier(
+                    settings.notification.type,
+                    config=settings.notification.dict(by_alias=True),
+                )
+                with ThreadPoolExecutor(max_workers=2) as worker:
+                    worker.map(
+                        lambda info: notifier.send(notification=info), renamed_info
+                    )
+
             self.stop_event.wait(settings.program.rename_time)
 
     def rename_start(self):

backend/src/module/models/api.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class RssLink(BaseModel):
@@ -16,3 +16,7 @@ class ChangeConfig(BaseModel):
 
 class ChangeRule(BaseModel):
     rule: dict
+
+
+class NotificationMessageIds(BaseModel):
+    message_ids: list[str] = Field(..., description="message ids to be set read")

backend/src/module/models/config.py

@@ -3,6 +3,8 @@
 
 from pydantic import BaseModel, Field
 
+from module.notification.services import NotificationType
+
 
 class Program(BaseModel):
     rss_time: int = Field(900, description="Sleep time")
@@ -70,9 +72,11 @@ def password(self):
 
 class Notification(BaseModel):
     enable: bool = Field(False, description="Enable notification")
-    type: str = Field("telegram", description="Notification type")
+    type: NotificationType = Field("telegram", description="Notification type")
     token_: str = Field("", alias="token", description="Notification token")
     chat_id_: str = Field("", alias="chat_id", description="Notification chat id")
+    base_url: str = Field("", description="Notification base url")
+    channel: str = Field("", description="Notification channel")
 
     @property
     def token(self):

backend/src/module/notification/__init__.py

@@ -1 +1,105 @@
-from .notification import PostNotification
+import json
+import logging
+from typing import get_args
+
+from litequeue import SQLQueue
+from tenacity import RetryError
+
+from module.models.bangumi import Notification
+from module.notification.base import NotificationContent
+
+from .services import NotificationService, NotificationType, services
+
+logger = logging.getLogger(__name__)
+
+
+class NotifierHandler(logging.Handler):
+    def __init__(self, service_name: str, **kwargs) -> None:
+        notifier_config = kwargs.pop("config", {})
+        self.notifier = Notifier(service_name, config=notifier_config)
+        super().__init__(**kwargs)
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            self.notifier.send(record=record)
+        except Exception as e:
+            logger.error(f"Can't send log record to notifier because: {e}")
+
+
+class Notifier:
+    def __init__(self, service_name: str, **kwargs):
+        assert service_name in get_args(
+            NotificationType
+        ), f"Invalid service name: {service_name}"
+
+        notifier_config = kwargs.pop("config", {})
+        if not notifier_config:
+            raise ValueError("Invalid notifier config")
+
+        self.notifier = services[service_name](**notifier_config)
+
+        from module.conf.const import ROOT
+
+        self._queue = SQLQueue(filename_or_conn=ROOT.joinpath("data", "queue.db"))
+
+    @property
+    def q(self) -> SQLQueue:
+        return self._queue
+
+    def _get_json(self, **kwargs):
+        content: NotificationContent = kwargs.get("content")
+        notification: Notification = kwargs.get("notification")
+        record: logging.LogRecord = kwargs.get("record")
+
+        if notification:
+            return notification.json()
+
+        if record:
+            args = dict(
+                name=record.name,
+                level=record.levelname,
+                pathname=record.pathname,
+                lineno=record.lineno,
+                msg=record.msg,
+            )
+            return json.dumps(args)
+
+        if content:
+            return content.json()
+
+        raise ValueError(f"Invalid input data: {kwargs}")
+
+    async def asend(self, **kwargs):
+        data = self._get_json(**kwargs)
+        try:
+            await self.notifier.asend(**kwargs)
+            self.q.put(data)
+        except RetryError as e:
+            e.reraise()
+        except Exception as e:
+            logger.warning(f"Failed to send notification: {e}")
+
+    def send(self, **kwargs) -> bool:
+        data = self._get_json(**kwargs)
+        try:
+            self.notifier.send(**kwargs)
+            self.q.put(data)
+        except RetryError as e:
+            e.reraise()
+        except Exception as e:
+            logger.warning(f"Failed to send notification: {e}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+__all__ = [
+    "Notifier",
+    "NotifierHandler",
+    "NotificationService",
+    "NotificationType",
+    "services",
+]