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 5217322c3a..4e200ad974 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 @@ -233,8 +233,20 @@ class TidbService: 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" + if not cluster_info.qdrant_endpoint: + cluster_info.qdrant_endpoint = TidbService.extract_qdrant_endpoint( + item + ) or TidbService.fetch_qdrant_endpoint( + api_url, public_key, private_key, item["clusterId"] + ) + if cluster_info.qdrant_endpoint: + cluster_info.status = TidbAuthBindingStatus.ACTIVE + else: + logger.warning( + "Cluster %s is ACTIVE but qdrant endpoint is not ready; will retry later", + item["clusterId"], + ) db.session.add(cluster_info) db.session.commit() else: diff --git a/api/tests/unit_tests/core/rag/datasource/vdb/tidb_on_qdrant/test_tidb_service.py b/api/tests/unit_tests/core/rag/datasource/vdb/tidb_on_qdrant/test_tidb_service.py index 77fc4d31dc..abcaf1eff7 100644 --- a/api/tests/unit_tests/core/rag/datasource/vdb/tidb_on_qdrant/test_tidb_service.py +++ b/api/tests/unit_tests/core/rag/datasource/vdb/tidb_on_qdrant/test_tidb_service.py @@ -1,8 +1,10 @@ +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from core.rag.datasource.vdb.tidb_on_qdrant.tidb_service import TidbService +from models.enums import TidbAuthBindingStatus class TestExtractQdrantEndpoint: @@ -217,3 +219,90 @@ class TestBatchCreateEdgeCases: private_key="priv", region="us-east-1", ) + + +class TestBatchUpdateTidbServerlessClusterStatus: + """Verify that status updates only expose clusters after qdrant endpoint is ready.""" + + @patch("core.rag.datasource.vdb.tidb_on_qdrant.tidb_service.db") + @patch("core.rag.datasource.vdb.tidb_on_qdrant.tidb_service.httpx") + def test_sets_active_when_batch_response_contains_endpoint(self, mock_http, mock_db): + binding = SimpleNamespace( + cluster_id="c-1", + status=TidbAuthBindingStatus.CREATING, + account="root", + qdrant_endpoint=None, + ) + mock_http.get.return_value = MagicMock( + status_code=200, + json=lambda: { + "clusters": [ + { + "clusterId": "c-1", + "state": "ACTIVE", + "userPrefix": "pfx", + "endpoints": {"public": {"host": "gw.tidbcloud.com"}}, + } + ] + }, + ) + + TidbService.batch_update_tidb_serverless_cluster_status([binding], "proj", "url", "iam", "pub", "priv") + + assert binding.account == "pfx.root" + assert binding.qdrant_endpoint == "https://qdrant-gw.tidbcloud.com" + assert binding.status == TidbAuthBindingStatus.ACTIVE + mock_db.session.add.assert_called_once_with(binding) + mock_db.session.commit.assert_called_once() + + @patch.object(TidbService, "fetch_qdrant_endpoint", return_value="https://qdrant-gw.tidbcloud.com") + @patch("core.rag.datasource.vdb.tidb_on_qdrant.tidb_service.db") + @patch("core.rag.datasource.vdb.tidb_on_qdrant.tidb_service.httpx") + def test_fetches_endpoint_when_batch_response_omits_it(self, mock_http, mock_db, mock_fetch_endpoint): + binding = SimpleNamespace( + cluster_id="c-1", + status=TidbAuthBindingStatus.CREATING, + account="root", + qdrant_endpoint=None, + ) + mock_http.get.return_value = MagicMock( + status_code=200, + json=lambda: { + "clusters": [{"clusterId": "c-1", "state": "ACTIVE", "userPrefix": "pfx", "endpoints": {}}] + }, + ) + + TidbService.batch_update_tidb_serverless_cluster_status([binding], "proj", "url", "iam", "pub", "priv") + + assert binding.account == "pfx.root" + assert binding.qdrant_endpoint == "https://qdrant-gw.tidbcloud.com" + assert binding.status == TidbAuthBindingStatus.ACTIVE + mock_fetch_endpoint.assert_called_once_with("url", "pub", "priv", "c-1") + mock_db.session.add.assert_called_once_with(binding) + mock_db.session.commit.assert_called_once() + + @patch.object(TidbService, "fetch_qdrant_endpoint", return_value=None) + @patch("core.rag.datasource.vdb.tidb_on_qdrant.tidb_service.db") + @patch("core.rag.datasource.vdb.tidb_on_qdrant.tidb_service.httpx") + def test_keeps_creating_when_endpoint_is_not_ready(self, mock_http, mock_db, mock_fetch_endpoint): + binding = SimpleNamespace( + cluster_id="c-1", + status=TidbAuthBindingStatus.CREATING, + account="root", + qdrant_endpoint=None, + ) + mock_http.get.return_value = MagicMock( + status_code=200, + json=lambda: { + "clusters": [{"clusterId": "c-1", "state": "ACTIVE", "userPrefix": "pfx", "endpoints": {}}] + }, + ) + + TidbService.batch_update_tidb_serverless_cluster_status([binding], "proj", "url", "iam", "pub", "priv") + + assert binding.account == "pfx.root" + assert binding.qdrant_endpoint is None + assert binding.status == TidbAuthBindingStatus.CREATING + mock_fetch_endpoint.assert_called_once_with("url", "pub", "priv", "c-1") + mock_db.session.add.assert_called_once_with(binding) + mock_db.session.commit.assert_called_once()