From 445dcfe4d09da16fecccf58a1b9985557cb4027f Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Fri, 8 Nov 2024 18:48:12 +0800
Subject: [PATCH] add create tidb serverless job control (#10467)

Co-authored-by: crazywoola <427733928@qq.com>
---
 api/.env.example                            | 2 ++
 api/configs/feature/__init__.py             | 5 +++++
 api/schedule/create_tidb_serverless_task.py | 2 ++
 docker-legacy/docker-compose.yaml           | 1 +
 docker/.env.example                         | 3 +++
 docker/docker-compose.yaml                  | 1 +
 6 files changed, 14 insertions(+)

diff --git a/api/.env.example b/api/.env.example
index 5dfc398df2..dbc402d34b 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -396,3 +396,5 @@ POSITION_PROVIDER_EXCLUDES=
 
 # Reset password token expiry minutes
 RESET_PASSWORD_TOKEN_EXPIRY_MINUTES=5
+
+CREATE_TIDB_SERVICE_JOB_ENABLED=false
\ No newline at end of file
diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py
index f011b638e3..23ccba60a2 100644
--- a/api/configs/feature/__init__.py
+++ b/api/configs/feature/__init__.py
@@ -611,6 +611,11 @@ class DataSetConfig(BaseSettings):
         default=500,
     )
 
+    CREATE_TIDB_SERVICE_JOB_ENABLED: bool = Field(
+        description="Enable or disable create tidb service job",
+        default=False,
+    )
+
 
 class WorkspaceConfig(BaseSettings):
     """
diff --git a/api/schedule/create_tidb_serverless_task.py b/api/schedule/create_tidb_serverless_task.py
index 42d6c04beb..a20b500308 100644
--- a/api/schedule/create_tidb_serverless_task.py
+++ b/api/schedule/create_tidb_serverless_task.py
@@ -12,6 +12,8 @@ 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:
diff --git a/docker-legacy/docker-compose.yaml b/docker-legacy/docker-compose.yaml
index 88650194ec..90110f49a2 100644
--- a/docker-legacy/docker-compose.yaml
+++ b/docker-legacy/docker-compose.yaml
@@ -384,6 +384,7 @@ services:
       NOTION_INTERNAL_SECRET: you-internal-secret
       # Indexing configuration
       INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH: 1000
+      CREATE_TIDB_SERVICE_JOB_ENABLED: false
     depends_on:
       - db
       - redis
diff --git a/docker/.env.example b/docker/.env.example
index 2dfb842a64..12443ee7f2 100644
--- a/docker/.env.example
+++ b/docker/.env.example
@@ -907,3 +907,6 @@ POSITION_PROVIDER_EXCLUDES=
 
 # CSP https://developer.mozilla.org/en-US/docs/Web/HTTP/CSP
 CSP_WHITELIST=
+
+# Enable or disable create tidb service job
+CREATE_TIDB_SERVICE_JOB_ENABLED=false
\ No newline at end of file
diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
index 60b0b9f257..6dcc3750c9 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -271,6 +271,7 @@ x-shared-env: &shared-api-worker-env
   OCEANBASE_VECTOR_DATABASE: ${OCEANBASE_VECTOR_DATABASE:-test}
   OCEANBASE_CLUSTER_NAME: ${OCEANBASE_CLUSTER_NAME:-difyai}
   OCEANBASE_MEMORY_LIMIT: ${OCEANBASE_MEMORY_LIMIT:-6G}
+  CREATE_TIDB_SERVICE_JOB_ENABLED: ${CREATE_TIDB_SERVICE_JOB_ENABLED:-false}
 
 services:
   # API service