add new endpoint for every new cluster

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,

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

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
     )

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,
             )