From 52711906832fbac9013ec8fdbd210243647ceee2 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Tue, 14 Apr 2026 15:47:10 +0800 Subject: [PATCH] add new endpoint for every new cluster --- .../tidb_on_qdrant/tidb_on_qdrant_vector.py | 7 +++- .../vdb/tidb_on_qdrant/tidb_service.py | 34 +++++++++++++++++++ api/models/dataset.py | 1 + api/schedule/create_tidb_serverless_task.py | 1 + 4 files changed, 42 insertions(+), 1 deletion(-) diff --git a/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py index 605cc5a08f..c40d71bf11 100644 --- a/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py +++ b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py @@ -439,6 +439,7 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory): idle_tidb_auth_binding.active = True idle_tidb_auth_binding.tenant_id = dataset.tenant_id db.session.commit() + tidb_auth_binding = idle_tidb_auth_binding TIDB_ON_QDRANT_API_KEY = f"{idle_tidb_auth_binding.account}:{idle_tidb_auth_binding.password}" else: new_cluster = TidbService.create_tidb_serverless_cluster( @@ -454,16 +455,20 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory): cluster_name=new_cluster["cluster_name"], account=new_cluster["account"], password=new_cluster["password"], + qdrant_endpoint=new_cluster.get("qdrant_endpoint"), tenant_id=dataset.tenant_id, active=True, status=TidbAuthBindingStatus.ACTIVE, ) db.session.add(new_tidb_auth_binding) db.session.commit() + tidb_auth_binding = new_tidb_auth_binding TIDB_ON_QDRANT_API_KEY = f"{new_tidb_auth_binding.account}:{new_tidb_auth_binding.password}" else: TIDB_ON_QDRANT_API_KEY = f"{tidb_auth_binding.account}:{tidb_auth_binding.password}" + qdrant_url = (tidb_auth_binding.qdrant_endpoint if tidb_auth_binding else None) or dify_config.TIDB_ON_QDRANT_URL or "" + if dataset.index_struct_dict: class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] collection_name = class_prefix @@ -478,7 +483,7 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory): collection_name=collection_name, group_id=dataset.id, config=TidbOnQdrantConfig( - endpoint=dify_config.TIDB_ON_QDRANT_URL or "", + endpoint=qdrant_url, api_key=TIDB_ON_QDRANT_API_KEY, root_path=str(config.root_path), timeout=dify_config.TIDB_ON_QDRANT_CLIENT_TIMEOUT, diff --git a/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py index 37114be6e7..f617fb081e 100644 --- a/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py +++ b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py @@ -20,6 +20,33 @@ _tidb_http_client: httpx.Client = get_pooled_http_client( class TidbService: + @staticmethod + def fetch_qdrant_endpoint( + api_url: str, public_key: str, private_key: str, cluster_id: str + ) -> str | None: + """Fetch the qdrant endpoint for a cluster by calling the Get Cluster API. + + The Get Cluster response contains ``status.connection_strings.standard.host`` + (e.g. ``gateway01.xx.tidbcloud.com``). We prepend ``qdrant-`` and wrap it + as an ``https://`` URL. + """ + try: + cluster_response = TidbService.get_tidb_serverless_cluster( + api_url, public_key, private_key, cluster_id + ) + if not cluster_response: + return None + # v1beta: status.connection_strings.standard.host + status = cluster_response.get("status") or {} + connection_strings = status.get("connection_strings") or {} + standard = connection_strings.get("standard") or {} + host = standard.get("host") + if host: + return f"https://qdrant-{host}" + except Exception: + pass + return None + @staticmethod def create_tidb_serverless_cluster( project_id: str, api_url: str, iam_url: str, public_key: str, private_key: str, region: str @@ -70,11 +97,15 @@ class TidbService: cluster_response = TidbService.get_tidb_serverless_cluster(api_url, public_key, private_key, cluster_id) if cluster_response["state"] == "ACTIVE": user_prefix = cluster_response["userPrefix"] + qdrant_endpoint = TidbService.fetch_qdrant_endpoint( + api_url, public_key, private_key, cluster_id + ) return { "cluster_id": cluster_id, "cluster_name": display_name, "account": f"{user_prefix}.root", "password": password, + "qdrant_endpoint": qdrant_endpoint, } time.sleep(30) # wait 30 seconds before retrying retry_count += 1 @@ -253,6 +284,9 @@ class TidbService: "cluster_name": item["displayName"], "account": "root", "password": cached_password.decode("utf-8"), + "qdrant_endpoint": TidbService.fetch_qdrant_endpoint( + api_url, public_key, private_key, item["clusterId"] + ), } cluster_infos.append(cluster_info) return cluster_infos diff --git a/api/models/dataset.py b/api/models/dataset.py index 97604848af..14dfcea5de 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -1250,6 +1250,7 @@ class TidbAuthBinding(TypeBase): ) account: Mapped[str] = mapped_column(String(255), nullable=False) password: Mapped[str] = mapped_column(String(255), nullable=False) + qdrant_endpoint: Mapped[str | None] = mapped_column(String(512), nullable=True, default=None) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) diff --git a/api/schedule/create_tidb_serverless_task.py b/api/schedule/create_tidb_serverless_task.py index 6ceb3ef856..511de002d2 100644 --- a/api/schedule/create_tidb_serverless_task.py +++ b/api/schedule/create_tidb_serverless_task.py @@ -57,6 +57,7 @@ def create_clusters(batch_size): cluster_name=new_cluster["cluster_name"], account=new_cluster["account"], password=new_cluster["password"], + qdrant_endpoint=new_cluster.get("qdrant_endpoint"), active=False, status=TidbAuthBindingStatus.CREATING, )