From b3faa34a030941d4b489fdd4bafe5cc1bdf84d68 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Tue, 10 Feb 2026 16:12:59 +0800 Subject: [PATCH] add notification logic for backend --- api/controllers/console/__init__.py | 2 + api/controllers/console/admin.py | 116 +++++++++++++++++++++++- api/controllers/console/notification.py | 26 ++++++ api/services/billing_service.py | 32 +++++++ 4 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 api/controllers/console/notification.py diff --git a/api/controllers/console/__init__.py b/api/controllers/console/__init__.py index fdc9aabc83..55e2669583 100644 --- a/api/controllers/console/__init__.py +++ b/api/controllers/console/__init__.py @@ -38,6 +38,7 @@ from . import ( extension, feature, init_validate, + notification, ping, setup, spec, @@ -182,6 +183,7 @@ __all__ = [ "model_config", "model_providers", "models", + "notification", "oauth", "oauth_server", "ops_trace", diff --git a/api/controllers/console/admin.py b/api/controllers/console/admin.py index 03b602f6e8..7761ec9b92 100644 --- a/api/controllers/console/admin.py +++ b/api/controllers/console/admin.py @@ -1,3 +1,6 @@ +import csv +import io +import re from collections.abc import Callable from functools import wraps from typing import ParamSpec, TypeVar @@ -6,7 +9,7 @@ from flask import request from flask_restx import Resource from pydantic import BaseModel, Field, field_validator from sqlalchemy import select -from werkzeug.exceptions import NotFound, Unauthorized +from werkzeug.exceptions import BadRequest, NotFound, Unauthorized from configs import dify_config from constants.languages import supported_language @@ -16,6 +19,7 @@ from core.db.session_factory import session_factory from extensions.ext_database import db from libs.token import extract_access_token from models.model import App, ExporleBanner, InstalledApp, RecommendedApp, TrialApp +from services.billing_service import BillingService P = ParamSpec("P") R = TypeVar("R") @@ -277,3 +281,113 @@ class DeleteExploreBannerApi(Resource): db.session.commit() return {"result": "success"}, 204 + + +class SaveNotificationContentPayload(BaseModel): + content: str = Field(...) + + +class SaveNotificationUserPayload(BaseModel): + user_email: list[str] = Field(...) + + +console_ns.schema_model( + SaveNotificationContentPayload.__name__, + SaveNotificationContentPayload.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0), +) + +console_ns.schema_model( + SaveNotificationUserPayload.__name__, + SaveNotificationUserPayload.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0), +) + + +@console_ns.route("/admin/save_notification_content") +class SaveNotificationContentApi(Resource): + @console_ns.doc("save_notification_content") + @console_ns.doc(description="Save a notification content") + @console_ns.expect(console_ns.models[SaveNotificationContentPayload.__name__]) + @console_ns.response(200, "Notification content saved successfully") + @only_edition_cloud + @admin_required + def post(self): + payload = SaveNotificationContentPayload.model_validate(console_ns.payload) + BillingService.save_notification_content(payload.content) + return {"result": "success"}, 200 + + +@console_ns.route("/admin/save_notification_user") +class SaveNotificationUserApi(Resource): + @console_ns.doc("save_notification_user") + @console_ns.doc(description="Save notification users via JSON body or file upload. " + "JSON: {\"user_email\": [\"a@example.com\", ...]}. " + "File: multipart/form-data with a 'file' field (CSV or TXT, one email per line).") + @console_ns.response(200, "Notification users saved successfully") + @only_edition_cloud + @admin_required + def post(self): + # Determine input mode: file upload or JSON body + if "file" in request.files: + emails = self._parse_emails_from_file() + else: + payload = SaveNotificationUserPayload.model_validate(console_ns.payload) + emails = payload.user_email + + if not emails: + raise BadRequest("No valid email addresses provided.") + + # Use batch API for bulk insert (chunks of 1000 per request to billing service) + result = BillingService.save_notification_users_batch(emails) + + return { + "result": "success", + "total": len(emails), + "succeeded": result["succeeded"], + "failed_chunks": result["failed_chunks"], + }, 200 + + @staticmethod + def _parse_emails_from_file() -> list[str]: + """Parse email addresses from an uploaded CSV or TXT file.""" + file = request.files["file"] + + if not file.filename: + raise BadRequest("Uploaded file has no filename.") + + filename_lower = file.filename.lower() + if not filename_lower.endswith((".csv", ".txt")): + raise BadRequest("Invalid file type. Only CSV (.csv) and TXT (.txt) files are allowed.") + + # Read file content + try: + content = file.read().decode("utf-8") + except UnicodeDecodeError: + try: + file.seek(0) + content = file.read().decode("gbk") + except UnicodeDecodeError: + raise BadRequest("Unable to decode the file. Please use UTF-8 or GBK encoding.") + + emails: list[str] = [] + if filename_lower.endswith(".csv"): + reader = csv.reader(io.StringIO(content)) + for row in reader: + for cell in row: + cell = cell.strip() + emails.append(cell) + else: + # TXT file: one email per line + for line in content.splitlines(): + line = line.strip() + emails.append(line) + + # Deduplicate while preserving order + seen: set[str] = set() + unique_emails: list[str] = [] + for email in emails: + email_lower = email.lower() + if email_lower not in seen: + seen.add(email_lower) + unique_emails.append(email) + + return unique_emails \ No newline at end of file diff --git a/api/controllers/console/notification.py b/api/controllers/console/notification.py new file mode 100644 index 0000000000..dd3feba974 --- /dev/null +++ b/api/controllers/console/notification.py @@ -0,0 +1,26 @@ +from flask_restx import Resource + +from controllers.console import console_ns +from controllers.console.wraps import account_initialization_required, only_edition_cloud, setup_required +from libs.login import current_account_with_tenant, login_required +from services.billing_service import BillingService + + +@console_ns.route("/notification") +class NotificationApi(Resource): + @console_ns.doc("get_notification") + @console_ns.doc(description="Get notification for the current user") + @console_ns.doc( + responses={ + 200: "Success", + 401: "Unauthorized", + } + ) + @setup_required + @login_required + @account_initialization_required + @only_edition_cloud + def get(self): + current_user, _ = current_account_with_tenant() + notification = BillingService.read_notification(current_user.email) + return notification diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 946b8cdfdb..de3a56e1f9 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -393,3 +393,35 @@ class BillingService: for item in data: tenant_whitelist.append(item["tenant_id"]) return tenant_whitelist + + @classmethod + def read_notification(cls, user_email: str): + params = {"user_email": user_email} + return cls._send_request("GET", "/notification/read", params=params) + + @classmethod + def save_notification_user(cls, user_email: str): + json = {"user_email": user_email} + return cls._send_request("POST", "/notification/new-notification-user", json=json) + + @classmethod + def save_notification_users_batch(cls, user_emails: list[str]) -> dict: + """Batch save notification users in chunks of 1000.""" + chunk_size = 1000 + total_succeeded = 0 + failed_chunks: list[dict] = [] + + for i in range(0, len(user_emails), chunk_size): + chunk = user_emails[i : i + chunk_size] + try: + resp = cls._send_request("POST", "/notification/batch-notification-users", json={"user_emails": chunk}) + total_succeeded += resp.get("count", len(chunk)) + except Exception as e: + failed_chunks.append({"offset": i, "count": len(chunk), "error": str(e)}) + + return {"succeeded": total_succeeded, "failed_chunks": failed_chunks} + + @classmethod + def save_notification_content(cls, content: str): + json = {"content": content} + return cls._send_request("POST", "/notification/new-notification", json=json) \ No newline at end of file