Initial commit

dify_1.9.0/api/schedule/check_upgradable_plugin_task.py

@@ -0,0 +1,49 @@
+import time
+
+import click
+
+import app
+from extensions.ext_database import db
+from models.account import TenantPluginAutoUpgradeStrategy
+from tasks.process_tenant_plugin_autoupgrade_check_task import process_tenant_plugin_autoupgrade_check_task
+
+AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL = 15 * 60  # 15 minutes
+
+
+@app.celery.task(queue="plugin")
+def check_upgradable_plugin_task():
+    click.echo(click.style("Start check upgradable plugin.", fg="green"))
+    start_at = time.perf_counter()
+
+    now_seconds_of_day = time.time() % 86400 - 30  # we assume the tz is UTC
+    click.echo(click.style(f"Now seconds of day: {now_seconds_of_day}", fg="green"))
+
+    strategies = (
+        db.session.query(TenantPluginAutoUpgradeStrategy)
+        .where(
+            TenantPluginAutoUpgradeStrategy.upgrade_time_of_day >= now_seconds_of_day,
+            TenantPluginAutoUpgradeStrategy.upgrade_time_of_day
+            < now_seconds_of_day + AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL,
+            TenantPluginAutoUpgradeStrategy.strategy_setting
+            != TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED,
+        )
+        .all()
+    )
+
+    for strategy in strategies:
+        process_tenant_plugin_autoupgrade_check_task.delay(
+            strategy.tenant_id,
+            strategy.strategy_setting,
+            strategy.upgrade_time_of_day,
+            strategy.upgrade_mode,
+            strategy.exclude_plugins,
+            strategy.include_plugins,
+        )
+
+    end_at = time.perf_counter()
+    click.echo(
+        click.style(
+            f"Checked upgradable plugin success latency: {end_at - start_at}",
+            fg="green",
+        )
+    )

dify_1.9.0/api/schedule/clean_embedding_cache_task.py

@@ -0,0 +1,42 @@
+import datetime
+import time
+
+import click
+from sqlalchemy import text
+from sqlalchemy.exc import SQLAlchemyError
+
+import app
+from configs import dify_config
+from extensions.ext_database import db
+from models.dataset import Embedding
+
+
+@app.celery.task(queue="dataset")
+def clean_embedding_cache_task():
+    click.echo(click.style("Start clean embedding cache.", fg="green"))
+    clean_days = int(dify_config.PLAN_SANDBOX_CLEAN_DAY_SETTING)
+    start_at = time.perf_counter()
+    thirty_days_ago = datetime.datetime.now() - datetime.timedelta(days=clean_days)
+    while True:
+        try:
+            embedding_ids = (
+                db.session.query(Embedding.id)
+                .where(Embedding.created_at < thirty_days_ago)
+                .order_by(Embedding.created_at.desc())
+                .limit(100)
+                .all()
+            )
+            embedding_ids = [embedding_id[0] for embedding_id in embedding_ids]
+        except SQLAlchemyError:
+            raise
+        if embedding_ids:
+            for embedding_id in embedding_ids:
+                db.session.execute(
+                    text("DELETE FROM embeddings WHERE id = :embedding_id"), {"embedding_id": embedding_id}
+                )
+
+            db.session.commit()
+        else:
+            break
+    end_at = time.perf_counter()
+    click.echo(click.style(f"Cleaned embedding cache from db success latency: {end_at - start_at}", fg="green"))

dify_1.9.0/api/schedule/clean_messages.py

@@ -0,0 +1,89 @@
+import datetime
+import logging
+import time
+
+import click
+from sqlalchemy.exc import SQLAlchemyError
+
+import app
+from configs import dify_config
+from extensions.ext_database import db
+from extensions.ext_redis import redis_client
+from models.model import (
+    App,
+    Message,
+    MessageAgentThought,
+    MessageAnnotation,
+    MessageChain,
+    MessageFeedback,
+    MessageFile,
+)
+from models.web import SavedMessage
+from services.feature_service import FeatureService
+
+logger = logging.getLogger(__name__)
+
+
+@app.celery.task(queue="dataset")
+def clean_messages():
+    click.echo(click.style("Start clean messages.", fg="green"))
+    start_at = time.perf_counter()
+    plan_sandbox_clean_message_day = datetime.datetime.now() - datetime.timedelta(
+        days=dify_config.PLAN_SANDBOX_CLEAN_MESSAGE_DAY_SETTING
+    )
+    while True:
+        try:
+            # Main query with join and filter
+            messages = (
+                db.session.query(Message)
+                .where(Message.created_at < plan_sandbox_clean_message_day)
+                .order_by(Message.created_at.desc())
+                .limit(100)
+                .all()
+            )
+
+        except SQLAlchemyError:
+            raise
+        if not messages:
+            break
+        for message in messages:
+            app = db.session.query(App).filter_by(id=message.app_id).first()
+            if not app:
+                logger.warning(
+                    "Expected App record to exist, but none was found, app_id=%s, message_id=%s",
+                    message.app_id,
+                    message.id,
+                )
+                continue
+            features_cache_key = f"features:{app.tenant_id}"
+            plan_cache = redis_client.get(features_cache_key)
+            if plan_cache is None:
+                features = FeatureService.get_features(app.tenant_id)
+                redis_client.setex(features_cache_key, 600, features.billing.subscription.plan)
+                plan = features.billing.subscription.plan
+            else:
+                plan = plan_cache.decode()
+            if plan == "sandbox":
+                # clean related message
+                db.session.query(MessageFeedback).where(MessageFeedback.message_id == message.id).delete(
+                    synchronize_session=False
+                )
+                db.session.query(MessageAnnotation).where(MessageAnnotation.message_id == message.id).delete(
+                    synchronize_session=False
+                )
+                db.session.query(MessageChain).where(MessageChain.message_id == message.id).delete(
+                    synchronize_session=False
+                )
+                db.session.query(MessageAgentThought).where(MessageAgentThought.message_id == message.id).delete(
+                    synchronize_session=False
+                )
+                db.session.query(MessageFile).where(MessageFile.message_id == message.id).delete(
+                    synchronize_session=False
+                )
+                db.session.query(SavedMessage).where(SavedMessage.message_id == message.id).delete(
+                    synchronize_session=False
+                )
+                db.session.query(Message).where(Message.id == message.id).delete()
+                db.session.commit()
+    end_at = time.perf_counter()
+    click.echo(click.style(f"Cleaned messages from db success latency: {end_at - start_at}", fg="green"))

dify_1.9.0/api/schedule/clean_unused_datasets_task.py

@@ -0,0 +1,155 @@
+import datetime
+import time
+from typing import TypedDict
+
+import click
+from sqlalchemy import func, select
+from sqlalchemy.exc import SQLAlchemyError
+
+import app
+from configs import dify_config
+from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
+from extensions.ext_database import db
+from extensions.ext_redis import redis_client
+from models.dataset import Dataset, DatasetAutoDisableLog, DatasetQuery, Document
+from services.feature_service import FeatureService
+
+
+class CleanupConfig(TypedDict):
+    clean_day: datetime.datetime
+    plan_filter: str | None
+    add_logs: bool
+
+
+@app.celery.task(queue="dataset")
+def clean_unused_datasets_task():
+    click.echo(click.style("Start clean unused datasets indexes.", fg="green"))
+    start_at = time.perf_counter()
+
+    # Define cleanup configurations
+    cleanup_configs: list[CleanupConfig] = [
+        {
+            "clean_day": datetime.datetime.now() - datetime.timedelta(days=dify_config.PLAN_SANDBOX_CLEAN_DAY_SETTING),
+            "plan_filter": None,
+            "add_logs": True,
+        },
+        {
+            "clean_day": datetime.datetime.now() - datetime.timedelta(days=dify_config.PLAN_PRO_CLEAN_DAY_SETTING),
+            "plan_filter": "sandbox",
+            "add_logs": False,
+        },
+    ]
+
+    for config in cleanup_configs:
+        clean_day = config["clean_day"]
+        plan_filter = config["plan_filter"]
+        add_logs = config["add_logs"]
+
+        page = 1
+        while True:
+            try:
+                # Subquery for counting new documents
+                document_subquery_new = (
+                    db.session.query(Document.dataset_id, func.count(Document.id).label("document_count"))
+                    .where(
+                        Document.indexing_status == "completed",
+                        Document.enabled == True,
+                        Document.archived == False,
+                        Document.updated_at > clean_day,
+                    )
+                    .group_by(Document.dataset_id)
+                    .subquery()
+                )
+
+                # Subquery for counting old documents
+                document_subquery_old = (
+                    db.session.query(Document.dataset_id, func.count(Document.id).label("document_count"))
+                    .where(
+                        Document.indexing_status == "completed",
+                        Document.enabled == True,
+                        Document.archived == False,
+                        Document.updated_at < clean_day,
+                    )
+                    .group_by(Document.dataset_id)
+                    .subquery()
+                )
+
+                # Main query with join and filter
+                stmt = (
+                    select(Dataset)
+                    .outerjoin(document_subquery_new, Dataset.id == document_subquery_new.c.dataset_id)
+                    .outerjoin(document_subquery_old, Dataset.id == document_subquery_old.c.dataset_id)
+                    .where(
+                        Dataset.created_at < clean_day,
+                        func.coalesce(document_subquery_new.c.document_count, 0) == 0,
+                        func.coalesce(document_subquery_old.c.document_count, 0) > 0,
+                    )
+                    .order_by(Dataset.created_at.desc())
+                )
+
+                datasets = db.paginate(stmt, page=page, per_page=50, error_out=False)
+
+            except SQLAlchemyError:
+                raise
+
+            if datasets is None or datasets.items is None or len(datasets.items) == 0:
+                break
+
+            for dataset in datasets:
+                dataset_query = db.session.scalars(
+                    select(DatasetQuery).where(
+                        DatasetQuery.created_at > clean_day, DatasetQuery.dataset_id == dataset.id
+                    )
+                ).all()
+
+                if not dataset_query or len(dataset_query) == 0:
+                    try:
+                        should_clean = True
+
+                        # Check plan filter if specified
+                        if plan_filter:
+                            features_cache_key = f"features:{dataset.tenant_id}"
+                            plan_cache = redis_client.get(features_cache_key)
+                            if plan_cache is None:
+                                features = FeatureService.get_features(dataset.tenant_id)
+                                redis_client.setex(features_cache_key, 600, features.billing.subscription.plan)
+                                plan = features.billing.subscription.plan
+                            else:
+                                plan = plan_cache.decode()
+                            should_clean = plan == plan_filter
+
+                        if should_clean:
+                            # Add auto disable log if required
+                            if add_logs:
+                                documents = db.session.scalars(
+                                    select(Document).where(
+                                        Document.dataset_id == dataset.id,
+                                        Document.enabled == True,
+                                        Document.archived == False,
+                                    )
+                                ).all()
+                                for document in documents:
+                                    dataset_auto_disable_log = DatasetAutoDisableLog(
+                                        tenant_id=dataset.tenant_id,
+                                        dataset_id=dataset.id,
+                                        document_id=document.id,
+                                    )
+                                    db.session.add(dataset_auto_disable_log)
+
+                            # Remove index
+                            index_processor = IndexProcessorFactory(dataset.doc_form).init_index_processor()
+                            index_processor.clean(dataset, None)
+
+                            # Update document
+                            db.session.query(Document).filter_by(dataset_id=dataset.id).update(
+                                {Document.enabled: False}
+                            )
+                            db.session.commit()
+                            click.echo(click.style(f"Cleaned unused dataset {dataset.id} from db success!", fg="green"))
+                    except Exception as e:
+                        click.echo(click.style(f"clean dataset index error: {e.__class__.__name__} {str(e)}", fg="red"))
+
+            page += 1
+
+    end_at = time.perf_counter()
+    click.echo(click.style(f"Cleaned unused dataset from db success latency: {end_at - start_at}", fg="green"))

dify_1.9.0/api/schedule/clean_workflow_runlogs_precise.py

@@ -0,0 +1,155 @@
+import datetime
+import logging
+import time
+
+import click
+
+import app
+from configs import dify_config
+from extensions.ext_database import db
+from models.model import (
+    AppAnnotationHitHistory,
+    Conversation,
+    Message,
+    MessageAgentThought,
+    MessageAnnotation,
+    MessageChain,
+    MessageFeedback,
+    MessageFile,
+)
+from models.workflow import ConversationVariable, WorkflowAppLog, WorkflowNodeExecutionModel, WorkflowRun
+
+logger = logging.getLogger(__name__)
+
+
+MAX_RETRIES = 3
+BATCH_SIZE = dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE
+
+
+@app.celery.task(queue="dataset")
+def clean_workflow_runlogs_precise():
+    """Clean expired workflow run logs with retry mechanism and complete message cascade"""
+
+    click.echo(click.style("Start clean workflow run logs (precise mode with complete cascade).", fg="green"))
+    start_at = time.perf_counter()
+
+    retention_days = dify_config.WORKFLOW_LOG_RETENTION_DAYS
+    cutoff_date = datetime.datetime.now() - datetime.timedelta(days=retention_days)
+
+    try:
+        total_workflow_runs = db.session.query(WorkflowRun).where(WorkflowRun.created_at < cutoff_date).count()
+        if total_workflow_runs == 0:
+            logger.info("No expired workflow run logs found")
+            return
+        logger.info("Found %s expired workflow run logs to clean", total_workflow_runs)
+
+        total_deleted = 0
+        failed_batches = 0
+        batch_count = 0
+
+        while True:
+            workflow_runs = (
+                db.session.query(WorkflowRun.id).where(WorkflowRun.created_at < cutoff_date).limit(BATCH_SIZE).all()
+            )
+
+            if not workflow_runs:
+                break
+
+            workflow_run_ids = [run.id for run in workflow_runs]
+            batch_count += 1
+
+            success = _delete_batch_with_retry(workflow_run_ids, failed_batches)
+
+            if success:
+                total_deleted += len(workflow_run_ids)
+                failed_batches = 0
+            else:
+                failed_batches += 1
+                if failed_batches >= MAX_RETRIES:
+                    logger.error("Failed to delete batch after %s retries, aborting cleanup for today", MAX_RETRIES)
+                    break
+                else:
+                    # Calculate incremental delay times: 5, 10, 15 minutes
+                    retry_delay_minutes = failed_batches * 5
+                    logger.warning("Batch deletion failed, retrying in %s minutes...", retry_delay_minutes)
+                    time.sleep(retry_delay_minutes * 60)
+                    continue
+
+        logger.info("Cleanup completed: %s expired workflow run logs deleted", total_deleted)
+
+    except Exception:
+        db.session.rollback()
+        logger.exception("Unexpected error in workflow log cleanup")
+        raise
+
+    end_at = time.perf_counter()
+    execution_time = end_at - start_at
+    click.echo(click.style(f"Cleaned workflow run logs from db success latency: {execution_time:.2f}s", fg="green"))
+
+
+def _delete_batch_with_retry(workflow_run_ids: list[str], attempt_count: int) -> bool:
+    """Delete a single batch with a retry mechanism and complete cascading deletion"""
+    try:
+        with db.session.begin_nested():
+            message_data = (
+                db.session.query(Message.id, Message.conversation_id)
+                .where(Message.workflow_run_id.in_(workflow_run_ids))
+                .all()
+            )
+            message_id_list = [msg.id for msg in message_data]
+            conversation_id_list = list({msg.conversation_id for msg in message_data if msg.conversation_id})
+            if message_id_list:
+                db.session.query(AppAnnotationHitHistory).where(
+                    AppAnnotationHitHistory.message_id.in_(message_id_list)
+                ).delete(synchronize_session=False)
+
+                db.session.query(MessageAgentThought).where(MessageAgentThought.message_id.in_(message_id_list)).delete(
+                    synchronize_session=False
+                )
+
+                db.session.query(MessageChain).where(MessageChain.message_id.in_(message_id_list)).delete(
+                    synchronize_session=False
+                )
+
+                db.session.query(MessageFile).where(MessageFile.message_id.in_(message_id_list)).delete(
+                    synchronize_session=False
+                )
+
+                db.session.query(MessageAnnotation).where(MessageAnnotation.message_id.in_(message_id_list)).delete(
+                    synchronize_session=False
+                )
+
+                db.session.query(MessageFeedback).where(MessageFeedback.message_id.in_(message_id_list)).delete(
+                    synchronize_session=False
+                )
+
+                db.session.query(Message).where(Message.workflow_run_id.in_(workflow_run_ids)).delete(
+                    synchronize_session=False
+                )
+
+            db.session.query(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(workflow_run_ids)).delete(
+                synchronize_session=False
+            )
+
+            db.session.query(WorkflowNodeExecutionModel).where(
+                WorkflowNodeExecutionModel.workflow_run_id.in_(workflow_run_ids)
+            ).delete(synchronize_session=False)
+
+            if conversation_id_list:
+                db.session.query(ConversationVariable).where(
+                    ConversationVariable.conversation_id.in_(conversation_id_list)
+                ).delete(synchronize_session=False)
+
+                db.session.query(Conversation).where(Conversation.id.in_(conversation_id_list)).delete(
+                    synchronize_session=False
+                )
+
+            db.session.query(WorkflowRun).where(WorkflowRun.id.in_(workflow_run_ids)).delete(synchronize_session=False)
+
+        db.session.commit()
+        return True
+
+    except Exception:
+        db.session.rollback()
+        logger.exception("Batch deletion failed (attempt %s)", attempt_count + 1)
+        return False

dify_1.9.0/api/schedule/create_tidb_serverless_task.py

@@ -0,0 +1,61 @@
+import time
+
+import click
+
+import app
+from configs import dify_config
+from core.rag.datasource.vdb.tidb_on_qdrant.tidb_service import TidbService
+from extensions.ext_database import db
+from models.dataset import TidbAuthBinding
+
+
+@app.celery.task(queue="dataset")
+def create_tidb_serverless_task():
+    click.echo(click.style("Start create tidb serverless task.", fg="green"))
+    if not dify_config.CREATE_TIDB_SERVICE_JOB_ENABLED:
+        return
+    tidb_serverless_number = dify_config.TIDB_SERVERLESS_NUMBER
+    start_at = time.perf_counter()
+    while True:
+        try:
+            # check the number of idle tidb serverless
+            idle_tidb_serverless_number = (
+                db.session.query(TidbAuthBinding).where(TidbAuthBinding.active == False).count()
+            )
+            if idle_tidb_serverless_number >= tidb_serverless_number:
+                break
+            # create tidb serverless
+            iterations_per_thread = 20
+            create_clusters(iterations_per_thread)
+
+        except Exception as e:
+            click.echo(click.style(f"Error: {e}", fg="red"))
+            break
+
+    end_at = time.perf_counter()
+    click.echo(click.style(f"Create tidb serverless task success latency: {end_at - start_at}", fg="green"))
+
+
+def create_clusters(batch_size):
+    try:
+        # TODO: maybe we can set the default value for the following parameters in the config file
+        new_clusters = TidbService.batch_create_tidb_serverless_cluster(
+            batch_size=batch_size,
+            project_id=dify_config.TIDB_PROJECT_ID or "",
+            api_url=dify_config.TIDB_API_URL or "",
+            iam_url=dify_config.TIDB_IAM_API_URL or "",
+            public_key=dify_config.TIDB_PUBLIC_KEY or "",
+            private_key=dify_config.TIDB_PRIVATE_KEY or "",
+            region=dify_config.TIDB_REGION or "",
+        )
+        for new_cluster in new_clusters:
+            tidb_auth_binding = TidbAuthBinding(
+                cluster_id=new_cluster["cluster_id"],
+                cluster_name=new_cluster["cluster_name"],
+                account=new_cluster["account"],
+                password=new_cluster["password"],
+            )
+            db.session.add(tidb_auth_binding)
+        db.session.commit()
+    except Exception as e:
+        click.echo(click.style(f"Error: {e}", fg="red"))

dify_1.9.0/api/schedule/mail_clean_document_notify_task.py

@@ -0,0 +1,97 @@
+import logging
+import time
+from collections import defaultdict
+
+import click
+from sqlalchemy import select
+
+import app
+from configs import dify_config
+from extensions.ext_database import db
+from extensions.ext_mail import mail
+from libs.email_i18n import EmailType, get_email_i18n_service
+from models.account import Account, Tenant, TenantAccountJoin
+from models.dataset import Dataset, DatasetAutoDisableLog
+from services.feature_service import FeatureService
+
+logger = logging.getLogger(__name__)
+
+
+@app.celery.task(queue="dataset")
+def mail_clean_document_notify_task():
+    """
+    Async Send document clean notify mail
+
+    Usage: mail_clean_document_notify_task.delay()
+    """
+    if not mail.is_inited():
+        return
+
+    logger.info(click.style("Start send document clean notify mail", fg="green"))
+    start_at = time.perf_counter()
+
+    # send document clean notify mail
+    try:
+        dataset_auto_disable_logs = db.session.scalars(
+            select(DatasetAutoDisableLog).where(DatasetAutoDisableLog.notified == False)
+        ).all()
+        # group by tenant_id
+        dataset_auto_disable_logs_map: dict[str, list[DatasetAutoDisableLog]] = defaultdict(list)
+        for dataset_auto_disable_log in dataset_auto_disable_logs:
+            if dataset_auto_disable_log.tenant_id not in dataset_auto_disable_logs_map:
+                dataset_auto_disable_logs_map[dataset_auto_disable_log.tenant_id] = []
+            dataset_auto_disable_logs_map[dataset_auto_disable_log.tenant_id].append(dataset_auto_disable_log)
+        url = f"{dify_config.CONSOLE_WEB_URL}/datasets"
+        for tenant_id, tenant_dataset_auto_disable_logs in dataset_auto_disable_logs_map.items():
+            features = FeatureService.get_features(tenant_id)
+            plan = features.billing.subscription.plan
+            if plan != "sandbox":
+                knowledge_details = []
+                # check tenant
+                tenant = db.session.query(Tenant).where(Tenant.id == tenant_id).first()
+                if not tenant:
+                    continue
+                # check current owner
+                current_owner_join = (
+                    db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, role="owner").first()
+                )
+                if not current_owner_join:
+                    continue
+                account = db.session.query(Account).where(Account.id == current_owner_join.account_id).first()
+                if not account:
+                    continue
+
+                dataset_auto_dataset_map = {}  # type: ignore
+                for dataset_auto_disable_log in tenant_dataset_auto_disable_logs:
+                    if dataset_auto_disable_log.dataset_id not in dataset_auto_dataset_map:
+                        dataset_auto_dataset_map[dataset_auto_disable_log.dataset_id] = []
+                    dataset_auto_dataset_map[dataset_auto_disable_log.dataset_id].append(
+                        dataset_auto_disable_log.document_id
+                    )
+
+                for dataset_id, document_ids in dataset_auto_dataset_map.items():
+                    dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first()
+                    if dataset:
+                        document_count = len(document_ids)
+                        knowledge_details.append(rf"Knowledge base {dataset.name}: {document_count} documents")
+                if knowledge_details:
+                    email_service = get_email_i18n_service()
+                    email_service.send_email(
+                        email_type=EmailType.DOCUMENT_CLEAN_NOTIFY,
+                        language_code="en-US",
+                        to=account.email,
+                        template_context={
+                            "userName": account.email,
+                            "knowledge_details": knowledge_details,
+                            "url": url,
+                        },
+                    )
+
+            # update notified to True
+            for dataset_auto_disable_log in tenant_dataset_auto_disable_logs:
+                dataset_auto_disable_log.notified = True
+            db.session.commit()
+        end_at = time.perf_counter()
+        logger.info(click.style(f"Send document clean notify mail succeeded: latency: {end_at - start_at}", fg="green"))
+    except Exception:
+        logger.exception("Send document clean notify mail failed")

dify_1.9.0/api/schedule/queue_monitor_task.py

@@ -0,0 +1,73 @@
+import logging
+from datetime import datetime
+
+import click
+from kombu.utils.url import parse_url  # type: ignore
+from redis import Redis
+
+import app
+from configs import dify_config
+from extensions.ext_database import db
+from libs.email_i18n import EmailType, get_email_i18n_service
+
+redis_config = parse_url(dify_config.CELERY_BROKER_URL)
+celery_redis = Redis(
+    host=redis_config.get("hostname") or "localhost",
+    port=redis_config.get("port") or 6379,
+    password=redis_config.get("password") or None,
+    db=int(redis_config.get("virtual_host")) if redis_config.get("virtual_host") else 1,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@app.celery.task(queue="monitor")
+def queue_monitor_task():
+    queue_name = "dataset"
+    threshold = dify_config.QUEUE_MONITOR_THRESHOLD
+
+    if threshold is None:
+        logger.warning(click.style("QUEUE_MONITOR_THRESHOLD is not configured, skipping monitoring", fg="yellow"))
+        return
+
+    try:
+        queue_length = celery_redis.llen(f"{queue_name}")
+        logger.info(click.style(f"Start monitor {queue_name}", fg="green"))
+
+        if queue_length is None:
+            logger.error(
+                click.style(f"Failed to get queue length for {queue_name} - Redis may be unavailable", fg="red")
+            )
+            return
+
+        logger.info(click.style(f"Queue length: {queue_length}", fg="green"))
+
+        if queue_length >= threshold:
+            warning_msg = f"Queue {queue_name} task count exceeded the limit.: {queue_length}/{threshold}"
+            logging.warning(click.style(warning_msg, fg="red"))
+            alert_emails = dify_config.QUEUE_MONITOR_ALERT_EMAILS
+            if alert_emails:
+                to_list = alert_emails.split(",")
+                email_service = get_email_i18n_service()
+                for to in to_list:
+                    try:
+                        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                        email_service.send_email(
+                            email_type=EmailType.QUEUE_MONITOR_ALERT,
+                            language_code="en-US",
+                            to=to,
+                            template_context={
+                                "queue_name": queue_name,
+                                "queue_length": queue_length,
+                                "threshold": threshold,
+                                "alert_time": current_time,
+                            },
+                        )
+                    except Exception:
+                        logger.exception(click.style("Exception occurred during sending email", fg="red"))
+
+    except Exception:
+        logger.exception(click.style("Exception occurred during queue monitoring", fg="red"))
+    finally:
+        if db.session.is_active:
+            db.session.close()

dify_1.9.0/api/schedule/update_tidb_serverless_status_task.py

@@ -0,0 +1,50 @@
+import time
+from collections.abc import Sequence
+
+import click
+from sqlalchemy import select
+
+import app
+from configs import dify_config
+from core.rag.datasource.vdb.tidb_on_qdrant.tidb_service import TidbService
+from extensions.ext_database import db
+from models.dataset import TidbAuthBinding
+
+
+@app.celery.task(queue="dataset")
+def update_tidb_serverless_status_task():
+    click.echo(click.style("Update tidb serverless status task.", fg="green"))
+    start_at = time.perf_counter()
+    try:
+        # check the number of idle tidb serverless
+        tidb_serverless_list = db.session.scalars(
+            select(TidbAuthBinding).where(TidbAuthBinding.active == False, TidbAuthBinding.status == "CREATING")
+        ).all()
+        if len(tidb_serverless_list) == 0:
+            return
+        # update tidb serverless status
+        update_clusters(tidb_serverless_list)
+
+    except Exception as e:
+        click.echo(click.style(f"Error: {e}", fg="red"))
+
+    end_at = time.perf_counter()
+    click.echo(click.style(f"Update tidb serverless status task success latency: {end_at - start_at}", fg="green"))
+
+
+def update_clusters(tidb_serverless_list: Sequence[TidbAuthBinding]):
+    try:
+        # batch 20
+        for i in range(0, len(tidb_serverless_list), 20):
+            items = tidb_serverless_list[i : i + 20]
+            # TODO: maybe we can set the default value for the following parameters in the config file
+            TidbService.batch_update_tidb_serverless_cluster_status(
+                tidb_serverless_list=items,
+                project_id=dify_config.TIDB_PROJECT_ID or "",
+                api_url=dify_config.TIDB_API_URL or "",
+                iam_url=dify_config.TIDB_IAM_API_URL or "",
+                public_key=dify_config.TIDB_PUBLIC_KEY or "",
+                private_key=dify_config.TIDB_PRIVATE_KEY or "",
+            )
+    except Exception as e:
+        click.echo(click.style(f"Error: {e}", fg="red"))