import logging import os import time from dataclasses import dataclass import requests from tenacity import ( before_sleep_log, retry, retry_if_exception, stop_after_attempt, wait_fixed, ) LOGGER = logging.getLogger(__name__) SYNC_PENDING_STATUSES = {"pending", "syncing"} SYNC_SUCCESS_STATUS = "synced" SYNC_FAILED_STATUS = "failed" CLUSTER_TYPES = ("SHARED_POOL", "LOGICAL_CLUSTER") class ConfigCenterError(Exception): pass class RetryableRequestError(ConfigCenterError): pass @dataclass(frozen=True) class RegisterArgs: env: str logical_region_name: str server_address: str is_byoc: bool request_timeout_seconds: int poll_interval_seconds: int poll_timeout_seconds: int max_sync_retries: int def _is_retryable_request_error(exc: BaseException) -> bool: return isinstance(exc, (requests.Timeout, requests.ConnectionError, RetryableRequestError)) class ConfigCenterClient: def __init__(self, env: str, request_timeout_seconds: int): auth_key = os.environ.get("GLOBAL_AUTH_KEY", "") if auth_key == "": raise ConfigCenterError("GLOBAL_AUTH_KEY is required") self.base_url = f"https://inner-central.{env}.tidbcloud.com/inner/config-center" self.request_timeout_seconds = request_timeout_seconds self.session = requests.Session() self.session.headers.update( { "Authorization": f"Basic {auth_key}", "Content-Type": "application/json", } ) @retry( retry=retry_if_exception(_is_retryable_request_error), wait=wait_fixed(3), stop=stop_after_attempt(5), before_sleep=before_sleep_log(LOGGER, logging.WARNING), reraise=True, ) def _request( self, method: str, path: str, payload: dict | None = None, expect_json: bool = True, ) -> dict: url = f"{self.base_url}{path}" response = self.session.request(method, url, json=payload, timeout=self.request_timeout_seconds) if response.status_code == 429 or response.status_code >= 500: raise RetryableRequestError( f"retryable response for {method} {url}: status={response.status_code}, body={response.text}" ) if response.status_code >= 400: raise ConfigCenterError( f"request failed for {method} {url}: status={response.status_code}, body={response.text}" ) if response.text.strip() == "": return {} if not expect_json: return {} try: return response.json() except ValueError as exc: raise ConfigCenterError( f"invalid json response for {method} {url}: {response.text}" ) from exc def register_regional_server( self, cluster_type: str, logical_region_name: str, server_address: str, is_byoc: bool, ) -> dict: payload = { "logicalRegion": logical_region_name, "serverAddress": server_address, "isByoc": is_byoc, } return self._request( "POST", f"/cluster-types/{cluster_type}/regional-servers", payload=payload, ) def get_regional_server(self, cluster_type: str, logical_region_name: str) -> dict: return self._request( "GET", f"/cluster-types/{cluster_type}/regional-servers/{logical_region_name}", ) def retry_sync(self, cluster_type: str, logical_region_name: str) -> dict: return self._request( "POST", f"/cluster-types/{cluster_type}/regional-servers/{logical_region_name}/sync", ) def delete_regional_server(self, cluster_type: str, logical_region_name: str) -> dict: return self._request( "DELETE", f"/cluster-types/{cluster_type}/regional-servers/{logical_region_name}", expect_json=False, ) def extract_server_address(response: dict) -> str: server_address = response.get("serverAddress") or response.get("server_address") if not server_address: raise ConfigCenterError(f"server address missing in response: {response}") return str(server_address) def extract_sync_status(response: dict) -> str: status = response.get("syncStatus") or response.get("sync_status") if not status: raise ConfigCenterError(f"sync status missing in response: {response}") return str(status).lower() def wait_until_synced(client: ConfigCenterClient, cluster_type: str, args: RegisterArgs) -> None: deadline = time.monotonic() + args.poll_timeout_seconds sync_retry_count = 0 while time.monotonic() < deadline: response = client.get_regional_server(cluster_type, args.logical_region_name) status = extract_sync_status(response) LOGGER.info( "cluster_type=%s logical_region=%s sync_status=%s sync_retry_count=%s", cluster_type, args.logical_region_name, status, sync_retry_count, ) if status == SYNC_SUCCESS_STATUS: return if status in SYNC_PENDING_STATUSES: time.sleep(args.poll_interval_seconds) continue if status == SYNC_FAILED_STATUS: if sync_retry_count >= args.max_sync_retries: raise ConfigCenterError( f"cluster_type={cluster_type} exceeded max sync retries: {args.max_sync_retries}" ) sync_retry_count += 1 LOGGER.warning( "cluster_type=%s logical_region=%s sync failed, retrying %s/%s", cluster_type, args.logical_region_name, sync_retry_count, args.max_sync_retries, ) client.retry_sync(cluster_type, args.logical_region_name) time.sleep(args.poll_interval_seconds) continue raise ConfigCenterError( f"cluster_type={cluster_type} got unknown sync status: {status}" ) raise ConfigCenterError( f"cluster_type={cluster_type} was not synced within {args.poll_timeout_seconds} seconds" ) def register_all(args: RegisterArgs) -> None: client = ConfigCenterClient(args.env, args.request_timeout_seconds) for cluster_type in CLUSTER_TYPES: LOGGER.info( "registering cluster_type=%s logical_region=%s server_address=%s", cluster_type, args.logical_region_name, args.server_address, ) response = client.register_regional_server( cluster_type, args.logical_region_name, args.server_address, args.is_byoc, ) returned_server_address = extract_server_address(response) if returned_server_address != args.server_address: raise ConfigCenterError( f"cluster_type={cluster_type} returned unexpected server address: {returned_server_address}" ) wait_until_synced(client, cluster_type, args) LOGGER.info( "cluster_type=%s logical_region=%s synced", cluster_type, args.logical_region_name, ) def delete_all(env: str, logical_region_name: str, request_timeout_seconds: int) -> None: client = ConfigCenterClient(env, request_timeout_seconds) for cluster_type in reversed(CLUSTER_TYPES): LOGGER.info( "deleting cluster_type=%s logical_region=%s", cluster_type, logical_region_name, ) client.delete_regional_server(cluster_type, logical_region_name)