fix: wait for qdrant endpoint before activating tidb cluster

This commit is contained in:
Yansong Zhang 2026-05-07 09:47:36 +08:00
parent e100e96279
commit 42c3163f90
2 changed files with 102 additions and 1 deletions

View File

@ -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:

View File

@ -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()