[autofix.ci] apply automated fixes

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py

@@ -0,0 +1,292 @@
+import logging
+import time
+import uuid
+from collections.abc import Sequence
+
+import httpx
+from httpx import DigestAuth
+
+from configs import dify_config
+from core.helper.http_client_pooling import get_pooled_http_client
+from extensions.ext_database import db
+from extensions.ext_redis import redis_client
+from models.dataset import TidbAuthBinding
+from models.enums import TidbAuthBindingStatus
+
+logger = logging.getLogger(__name__)
+
+# Reuse a pooled HTTP client for all TiDB Cloud requests to minimize connection churn
+_tidb_http_client: httpx.Client = get_pooled_http_client(
+    "tidb:cloud",
+    lambda: httpx.Client(limits=httpx.Limits(max_keepalive_connections=50, max_connections=100)),
+)
+
+
+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:
+            logger.exception("Failed to fetch qdrant endpoint for cluster %s", cluster_id)
+        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
+    ):
+        """
+        Creates a new TiDB Serverless cluster.
+        :param project_id: The project ID of the TiDB Cloud project (required).
+        :param api_url: The URL of the TiDB Cloud API (required).
+        :param iam_url: The URL of the TiDB Cloud IAM API (required).
+        :param public_key: The public key for the API (required).
+        :param private_key: The private key for the API (required).
+        :param region: The region where the cluster will be created (required).
+
+        :return: The response from the API.
+        """
+
+        region_object = {
+            "name": region,
+        }
+
+        labels = {
+            "tidb.cloud/project": project_id,
+        }
+
+        spending_limit = {
+            "monthly": dify_config.TIDB_SPEND_LIMIT,
+        }
+        password = str(uuid.uuid4()).replace("-", "")[:16]
+        display_name = str(uuid.uuid4()).replace("-", "")[:16]
+        cluster_data = {
+            "displayName": display_name,
+            "region": region_object,
+            "labels": labels,
+            "spendingLimit": spending_limit,
+            "rootPassword": password,
+        }
+
+        response = _tidb_http_client.post(
+            f"{api_url}/clusters", json=cluster_data, auth=DigestAuth(public_key, private_key)
+        )
+
+        if response.status_code == 200:
+            response_data = response.json()
+            cluster_id = response_data["clusterId"]
+            retry_count = 0
+            max_retries = 30
+            while retry_count < max_retries:
+                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
+        else:
+            response.raise_for_status()
+
+    @staticmethod
+    def delete_tidb_serverless_cluster(api_url: str, public_key: str, private_key: str, cluster_id: str):
+        """
+        Deletes a specific TiDB Serverless cluster.
+
+        :param api_url: The URL of the TiDB Cloud API (required).
+        :param public_key: The public key for the API (required).
+        :param private_key: The private key for the API (required).
+        :param cluster_id: The ID of the cluster to be deleted (required).
+        :return: The response from the API.
+        """
+
+        response = _tidb_http_client.delete(
+            f"{api_url}/clusters/{cluster_id}", auth=DigestAuth(public_key, private_key)
+        )
+
+        if response.status_code == 200:
+            return response.json()
+        else:
+            response.raise_for_status()
+
+    @staticmethod
+    def get_tidb_serverless_cluster(api_url: str, public_key: str, private_key: str, cluster_id: str):
+        """
+        Deletes a specific TiDB Serverless cluster.
+
+        :param api_url: The URL of the TiDB Cloud API (required).
+        :param public_key: The public key for the API (required).
+        :param private_key: The private key for the API (required).
+        :param cluster_id: The ID of the cluster to be deleted (required).
+        :return: The response from the API.
+        """
+
+        response = _tidb_http_client.get(f"{api_url}/clusters/{cluster_id}", auth=DigestAuth(public_key, private_key))
+
+        if response.status_code == 200:
+            return response.json()
+        else:
+            response.raise_for_status()
+
+    @staticmethod
+    def change_tidb_serverless_root_password(
+        api_url: str, public_key: str, private_key: str, cluster_id: str, account: str, new_password: str
+    ):
+        """
+        Changes the root password of a specific TiDB Serverless cluster.
+
+        :param api_url: The URL of the TiDB Cloud API (required).
+        :param public_key: The public key for the API (required).
+        :param private_key: The private key for the API (required).
+        :param cluster_id: The ID of the cluster for which the password is to be changed (required).+
+        :param account: The account for which the password is to be changed (required).
+        :param new_password: The new password for the root user (required).
+        :return: The response from the API.
+        """
+
+        body = {"password": new_password, "builtinRole": "role_admin", "customRoles": []}
+
+        response = _tidb_http_client.patch(
+            f"{api_url}/clusters/{cluster_id}/sqlUsers/{account}",
+            json=body,
+            auth=DigestAuth(public_key, private_key),
+        )
+
+        if response.status_code == 200:
+            return response.json()
+        else:
+            response.raise_for_status()
+
+    @staticmethod
+    def batch_update_tidb_serverless_cluster_status(
+        tidb_serverless_list: Sequence[TidbAuthBinding],
+        project_id: str,
+        api_url: str,
+        iam_url: str,
+        public_key: str,
+        private_key: str,
+    ):
+        """
+        Update the status of a new TiDB Serverless cluster.
+        :param tidb_serverless_list: The TiDB serverless list (required).
+        :param project_id: The project ID of the TiDB Cloud project (required).
+        :param api_url: The URL of the TiDB Cloud API (required).
+        :param iam_url: The URL of the TiDB Cloud IAM API (required).
+        :param public_key: The public key for the API (required).
+        :param private_key: The private key for the API (required).
+
+        :return: The response from the API.
+        """
+        tidb_serverless_list_map = {item.cluster_id: item for item in tidb_serverless_list}
+        cluster_ids = [item.cluster_id for item in tidb_serverless_list]
+        params = {"clusterIds": cluster_ids, "view": "BASIC"}
+        response = _tidb_http_client.get(
+            f"{api_url}/clusters:batchGet", params=params, auth=DigestAuth(public_key, private_key)
+        )
+
+        if response.status_code == 200:
+            response_data = response.json()
+            for item in response_data["clusters"]:
+                state = item["state"]
+                userPrefix = item["userPrefix"]
+                if state == "ACTIVE" and len(userPrefix) > 0:
+                    cluster_info = tidb_serverless_list_map[item["clusterId"]]
+                    cluster_info.status = TidbAuthBindingStatus.ACTIVE
+                    cluster_info.account = f"{userPrefix}.root"
+                    db.session.add(cluster_info)
+            db.session.commit()
+        else:
+            response.raise_for_status()
+
+    @staticmethod
+    def batch_create_tidb_serverless_cluster(
+        batch_size: int, project_id: str, api_url: str, iam_url: str, public_key: str, private_key: str, region: str
+    ) -> list[dict]:
+        """
+        Creates a new TiDB Serverless cluster.
+        :param batch_size: The batch size (required).
+        :param project_id: The project ID of the TiDB Cloud project (required).
+        :param api_url: The URL of the TiDB Cloud API (required).
+        :param iam_url: The URL of the TiDB Cloud IAM API (required).
+        :param public_key: The public key for the API (required).
+        :param private_key: The private key for the API (required).
+        :param region: The region where the cluster will be created (required).
+
+        :return: The response from the API.
+        """
+        clusters = []
+        for _ in range(batch_size):
+            region_object = {
+                "name": region,
+            }
+
+            labels = {
+                "tidb.cloud/project": project_id,
+            }
+
+            spending_limit = {
+                "monthly": dify_config.TIDB_SPEND_LIMIT,
+            }
+            password = str(uuid.uuid4()).replace("-", "")[:16]
+            display_name = str(uuid.uuid4()).replace("-", "")
+            cluster_data = {
+                "cluster": {
+                    "displayName": display_name,
+                    "region": region_object,
+                    "labels": labels,
+                    "spendingLimit": spending_limit,
+                    "rootPassword": password,
+                }
+            }
+            cache_key = f"tidb_serverless_cluster_password:{display_name}"
+            redis_client.setex(cache_key, 3600, password)
+            clusters.append(cluster_data)
+
+        request_body = {"requests": clusters}
+        response = _tidb_http_client.post(
+            f"{api_url}/clusters:batchCreate", json=request_body, auth=DigestAuth(public_key, private_key)
+        )
+
+        if response.status_code == 200:
+            response_data = response.json()
+            cluster_infos = []
+            for item in response_data["clusters"]:
+                cache_key = f"tidb_serverless_cluster_password:{item['displayName']}"
+                cached_password = redis_client.get(cache_key)
+                if not cached_password:
+                    continue
+                cluster_info = {
+                    "cluster_id": item["clusterId"],
+                    "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
+        else:
+            response.raise_for_status()
+            return []

api/migrations/versions/2026_04_14_1500-8574b23a38fd_add_qdrant_endpoint_to_tidb_auth_bindings.py

@@ -0,0 +1,26 @@
+"""add qdrant_endpoint to tidb_auth_bindings
+
+Revision ID: 8574b23a38fd
+Revises: 6b5f9f8b1a2c
+Create Date: 2026-04-14 15:00:00.000000
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "8574b23a38fd"
+down_revision = "6b5f9f8b1a2c"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table("tidb_auth_bindings", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("qdrant_endpoint", sa.String(length=512), nullable=True))
+
+
+def downgrade():
+    with op.batch_alter_table("tidb_auth_bindings", schema=None) as batch_op:
+        batch_op.drop_column("qdrant_endpoint")

api/models/dataset.py

@@ -1305,6 +1305,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/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -440,6 +440,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(
@@ -455,16 +456,22 @@ 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
@@ -479,7 +486,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/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_service.py

@@ -1,3 +1,4 @@
+import logging
 import time
 import uuid
 from collections.abc import Sequence
@@ -12,6 +13,8 @@ from extensions.ext_redis import redis_client
 from models.dataset import TidbAuthBinding
 from models.enums import TidbAuthBindingStatus
 
+logger = logging.getLogger(__name__)
+
 # Reuse a pooled HTTP client for all TiDB Cloud requests to minimize connection churn
 _tidb_http_client: httpx.Client = get_pooled_http_client(
     "tidb:cloud",
@@ -20,6 +23,29 @@ _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:
+            logger.exception("Failed to fetch qdrant endpoint for cluster %s", cluster_id)
+        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 +96,13 @@ 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 +281,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/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_on_qdrant_vector.py

@@ -114,14 +114,12 @@ class TestTidbOnQdrantVectorDeleteByIds:
 
         assert exc_info.value.status_code == 500
 
-    def test_delete_by_ids_with_large_batch(self, vector_instance):
-        """Test deletion with a large batch of IDs."""
-        # Create 1000 IDs
+    def test_delete_by_ids_with_exactly_1000(self, vector_instance):
+        """Test deletion with exactly 1000 IDs triggers a single batch."""
         ids = [f"doc_{i}" for i in range(1000)]
 
         vector_instance.delete_by_ids(ids)
 
-        # Verify single delete call with all IDs
         vector_instance._client.delete.assert_called_once()
         call_args = vector_instance._client.delete.call_args
 
@@ -129,11 +127,28 @@ class TestTidbOnQdrantVectorDeleteByIds:
         filter_obj = filter_selector.filter
         field_condition = filter_obj.must[0]
 
-        # Verify all 1000 IDs are in the batch
         assert len(field_condition.match.any) == 1000
         assert "doc_0" in field_condition.match.any
         assert "doc_999" in field_condition.match.any
 
+    def test_delete_by_ids_splits_into_batches(self, vector_instance):
+        """Test deletion with >1000 IDs triggers multiple batched calls."""
+        ids = [f"doc_{i}" for i in range(2500)]
+
+        vector_instance.delete_by_ids(ids)
+
+        assert vector_instance._client.delete.call_count == 3
+
+        batches = []
+        for call in vector_instance._client.delete.call_args_list:
+            filter_selector = call[1]["points_selector"]
+            field_condition = filter_selector.filter.must[0]
+            batches.append(field_condition.match.any)
+
+        assert len(batches[0]) == 1000
+        assert len(batches[1]) == 1000
+        assert len(batches[2]) == 500
+
     def test_delete_by_ids_filter_structure(self, vector_instance):
         """Test that the filter structure is correctly constructed."""
         ids = ["doc1", "doc2"]

api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_service.py

@@ -0,0 +1,117 @@
+from unittest.mock import MagicMock, patch
+
+from dify_vdb_tidb_on_qdrant.tidb_service import TidbService
+
+
+class TestFetchQdrantEndpoint:
+    """Unit tests for TidbService.fetch_qdrant_endpoint."""
+
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    def test_returns_endpoint_when_host_present(self, mock_get_cluster):
+        mock_get_cluster.return_value = {
+            "status": {"connection_strings": {"standard": {"host": "gateway01.us-east-1.tidbcloud.com"}}}
+        }
+        result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
+        assert result == "https://qdrant-gateway01.us-east-1.tidbcloud.com"
+
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    def test_returns_none_when_cluster_response_is_none(self, mock_get_cluster):
+        mock_get_cluster.return_value = None
+        result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
+        assert result is None
+
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    def test_returns_none_when_host_missing(self, mock_get_cluster):
+        mock_get_cluster.return_value = {"status": {"connection_strings": {"standard": {}}}}
+        result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
+        assert result is None
+
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    def test_returns_none_when_status_missing(self, mock_get_cluster):
+        mock_get_cluster.return_value = {}
+        result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
+        assert result is None
+
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    def test_returns_none_when_connection_strings_missing(self, mock_get_cluster):
+        mock_get_cluster.return_value = {"status": {}}
+        result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
+        assert result is None
+
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    def test_returns_none_on_exception(self, mock_get_cluster):
+        mock_get_cluster.side_effect = RuntimeError("network error")
+        result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
+        assert result is None
+
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    def test_returns_none_when_standard_key_missing(self, mock_get_cluster):
+        mock_get_cluster.return_value = {"status": {"connection_strings": {}}}
+        result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
+        assert result is None
+
+
+class TestCreateTidbServerlessClusterQdrantEndpoint:
+    """Verify that create_tidb_serverless_cluster includes qdrant_endpoint in its result."""
+
+    @patch.object(TidbService, "fetch_qdrant_endpoint", return_value="https://qdrant-gw.tidbcloud.com")
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    @patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
+    @patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
+    def test_result_contains_qdrant_endpoint(self, mock_config, mock_http, mock_get_cluster, mock_fetch_ep):
+        mock_config.TIDB_SPEND_LIMIT = 10
+        mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
+        mock_get_cluster.return_value = {"state": "ACTIVE", "userPrefix": "pfx"}
+
+        result = TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
+
+        assert result is not None
+        assert result["qdrant_endpoint"] == "https://qdrant-gw.tidbcloud.com"
+        mock_fetch_ep.assert_called_once_with("url", "pub", "priv", "c-1")
+
+    @patch.object(TidbService, "fetch_qdrant_endpoint", return_value=None)
+    @patch.object(TidbService, "get_tidb_serverless_cluster")
+    @patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
+    @patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
+    def test_result_qdrant_endpoint_none_when_fetch_fails(
+        self, mock_config, mock_http, mock_get_cluster, mock_fetch_ep
+    ):
+        mock_config.TIDB_SPEND_LIMIT = 10
+        mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
+        mock_get_cluster.return_value = {"state": "ACTIVE", "userPrefix": "pfx"}
+
+        result = TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
+
+        assert result is not None
+        assert result["qdrant_endpoint"] is None
+
+
+class TestBatchCreateTidbServerlessClusterQdrantEndpoint:
+    """Verify that batch_create includes qdrant_endpoint per cluster."""
+
+    @patch.object(TidbService, "fetch_qdrant_endpoint", return_value="https://qdrant-gw.tidbcloud.com")
+    @patch("dify_vdb_tidb_on_qdrant.tidb_service.redis_client")
+    @patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
+    @patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
+    def test_batch_result_contains_qdrant_endpoint(self, mock_config, mock_http, mock_redis, mock_fetch_ep):
+        mock_config.TIDB_SPEND_LIMIT = 10
+        cluster_name = "abc123"
+        mock_http.post.return_value = MagicMock(
+            status_code=200,
+            json=lambda: {"clusters": [{"clusterId": "c-1", "displayName": cluster_name}]},
+        )
+        mock_redis.setex = MagicMock()
+        mock_redis.get.return_value = b"password123"
+
+        result = TidbService.batch_create_tidb_serverless_cluster(
+            batch_size=1,
+            project_id="proj",
+            api_url="url",
+            iam_url="iam",
+            public_key="pub",
+            private_key="priv",
+            region="us-east-1",
+        )
+
+        assert len(result) == 1
+        assert result[0]["qdrant_endpoint"] == "https://qdrant-gw.tidbcloud.com"

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