add create tidb serverless job control (langgenius#10467)

Co-authored-by: crazywoola <[email protected]>
AceDataCloud · Nov 9, 2024 · 445a36c · 445a36c
1 parent df75448
commit 445a36c
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 0 deletions.
diff --git 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
diff --git 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
@@ -12,6 +12,8 @@
 @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
@@ -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
@@ -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
diff --git 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
Original file line number	Diff line number	Diff line change
Expand Up		@@ -396,3 +396,5 @@ POSITION_PROVIDER_EXCLUDES=

		# Reset password token expiry minutes
		RESET_PASSWORD_TOKEN_EXPIRY_MINUTES=5

		CREATE_TIDB_SERVICE_JOB_ENABLED=false