Underground-Digital/Workflow-Engine

1import time2 3import click4 5import app6from configs import dify_config7from core.rag.datasource.vdb.tidb_on_qdrant.tidb_service import TidbService8from models.dataset import TidbAuthBinding9 10 11@app.celery.task(queue="dataset")12def update_tidb_serverless_status_task():13    click.echo(click.style("Update tidb serverless status task.", fg="green"))14    start_at = time.perf_counter()15    while True:16        try:17            # check the number of idle tidb serverless18            tidb_serverless_list = TidbAuthBinding.query.filter(19                TidbAuthBinding.active == False, TidbAuthBinding.status == "CREATING"20            ).all()21            if len(tidb_serverless_list) == 0:22                break23            # update tidb serverless status24            iterations_per_thread = 2025            update_clusters(tidb_serverless_list)26 27        except Exception as e:28            click.echo(click.style(f"Error: {e}", fg="red"))29            break30 31    end_at = time.perf_counter()32    click.echo(33        click.style("Update tidb serverless status task success latency: {}".format(end_at - start_at), fg="green")34    )35 36 37def update_clusters(tidb_serverless_list: list[TidbAuthBinding]):38    try:39        # batch 2040        for i in range(0, len(tidb_serverless_list), 20):41            items = tidb_serverless_list[i : i + 20]42            TidbService.batch_update_tidb_serverless_cluster_status(43                items,44                dify_config.TIDB_PROJECT_ID,45                dify_config.TIDB_API_URL,46                dify_config.TIDB_IAM_API_URL,47                dify_config.TIDB_PUBLIC_KEY,48                dify_config.TIDB_PRIVATE_KEY,49            )50    except Exception as e:51        click.echo(click.style(f"Error: {e}", fg="red"))52