from enum import Enum from typing import Dict, List import requests from environs import Env from config.tiworkload_config import TIWORKLOAD_CONFIG # Load environment variables env = Env() env.read_env() class Workload(Enum): TPCC = {"name": "tpcc", "storage": "tpcc-1k-with-clustered-index"} SYSBENCH = {"name": "oltp_read_write", "storage": "sysbench-32-1000w-auto"} TPCH = {"name": "tpch", "storage": "tpch-1k-with-clustered-index"} # not checked # Determine which host to use based on environment # Set TEST_ENV=local for local testing, otherwise use production TEST_ENV = env("TEST_ENV", "production") # Default to production HOST = ( TIWORKLOAD_CONFIG["host_localtest"] if TEST_ENV == "local" else TIWORKLOAD_CONFIG["host"] ) BASE_URL = TIWORKLOAD_CONFIG["base_url"] USER_ID = TIWORKLOAD_CONFIG["user_id"] def get_workload_br_storage( workload: Workload, cloud_provider: str, region: str ) -> List[Dict[str, str]]: """ Get BR storage information for a specific workload, cloud provider, and region. Args: workload: The workload type (TPCC, SYSBENCH, or TPCH) cloud_provider: Cloud provider (e.g., 'aws', 'gcp', 'azure') region: Region name without provider prefix (e.g., 'us-west-2', 'useast2', 'us-west1') Returns: List of matching BR storage configurations with keys: - storage_url: S3/GCS storage URL - s3_endpoint: S3 endpoint URL - db_name: Database name - br_options: BR restore options - provider: Cloud provider - region: Region name """ url = f"{HOST}{BASE_URL}/service/workloads/{workload.value.get('name')}/datasets/br" headers = {"Accept": "application/json", "X-User-Id": str(USER_ID)} try: response = requests.get(url, headers=headers) response.raise_for_status() all_datasets = response.json() # Normalize cloud provider name for matching # API uses lowercase provider names normalized_provider = cloud_provider.lower() # Filter by provider and region matched_datasets = [ dataset for dataset in all_datasets if dataset.get("provider", "").lower() == normalized_provider and dataset.get("region", "") == region and workload.value.get("storage") in dataset.get("storage_url", "") ] return matched_datasets except requests.exceptions.RequestException as e: print(f"Error fetching workload BR storage for {workload.value}: {e}") if hasattr(e, "response") and e.response is not None: print(f"Response content: {e.response.text}") return [] def get_all_workload_datasets(workload: Workload) -> List[Dict[str, str]]: """ Get all BR storage configurations for a specific workload. Args: workload: The workload type (TPCC, SYSBENCH, or TPCH) Returns: List of all BR storage configurations """ url = f"{HOST}{BASE_URL}/service/workloads/{workload.value.get('name')}/datasets/br" headers = {"Accept": "application/json", "X-User-Id": str(USER_ID)} try: response = requests.get(url, headers=headers) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: print(f"Error fetching all datasets for {workload.value}: {e}") if hasattr(e, "response") and e.response is not None: print(f"Response content: {e.response.text}") return [] def get_tiworkload_execution(execution_id: str) -> Dict: """ Get TiWorkload execution information. Args: execution_id: The execution ID Returns: Dict containing execution information with tasks """ url = f"{HOST}{BASE_URL}/service/workload/exec/{execution_id}" headers = {"Accept": "application/json", "X-User-Id": str(USER_ID)} try: response = requests.get(url, headers=headers) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: print(f"Error fetching execution {execution_id}: {e}") if hasattr(e, "response") and e.response is not None: print(f"Response content: {e.response.text}") return {} def fetch_task_log(exec_id: str, task_id: str) -> str: """ Fetch stdout log from TiWorkload API. Args: exec_id: Execution ID (kept for backward compatibility, not used) task_id: Task ID timeout_seconds: Deprecated, kept for backward compatibility Returns: Complete log content as string """ url = f"{HOST}{BASE_URL}/service/workloads/tasks/{task_id}/logs/file" headers = {"X-User-Id": str(USER_ID)} try: response = requests.get(url, headers=headers) response.raise_for_status() return response.text except requests.exceptions.RequestException as e: print(f"Error fetching log for {exec_id}/{task_id}: {e}") return "" def get_task_metrics(exec_id: str, task_id: str) -> Dict: """ Get task metrics from TiWorkload API. Args: exec_id: Execution ID task_id: Task ID Returns: Dict containing task metrics with structure: { "TaskID": str, "Stats": [ { "Index": str (e.g., "cpu", "memory", "network/srtt"), "Metrics": [ { "Method": str (e.g., "max", "avg", "p80", "p95", "p99", "stddev"), "Value": float, "Unit": str } ], "PodName": str } ], "CollectTime": str } Note: Network metrics (network/srtt) are available for newer test runs. If metrics are empty or missing, the dict may be empty or incomplete. """ url = f"{HOST}{BASE_URL}/service/workloads/execs/{exec_id}/tasks/{task_id}/metrics" headers = {"Accept": "application/json", "X-User-Id": str(USER_ID)} try: response = requests.get(url, headers=headers) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: print(f"Error fetching metrics for {exec_id}/{task_id}: {e}") if hasattr(e, "response") and e.response is not None: print(f"Response content: {e.response.text}") return {} if __name__ == "__main__": # Example usage print(f"Using host: {HOST}") print("\n=== Testing TPCC workload for AWS us-west-2 ===") tpcc_datasets = get_workload_br_storage( workload=Workload.TPCC, cloud_provider="aws", region="us-west-2" ) print(f"Found {len(tpcc_datasets)} matching datasets:") for dataset in tpcc_datasets: print(f" - DB: {dataset['db_name']}, Storage: {dataset['storage_url']}") print("\n=== Testing TPCC workload for GCP us-west1 ===") gcp_datasets = get_workload_br_storage( workload=Workload.TPCC, cloud_provider="gcp", region="us-west1" ) print(f"Found {len(gcp_datasets)} matching datasets:") for dataset in gcp_datasets: print(f" - DB: {dataset['db_name']}, Storage: {dataset['storage_url']}") print("\n=== Testing all TPCC datasets ===") all_datasets = get_all_workload_datasets(Workload.TPCC) print(f"Total datasets: {len(all_datasets)}") for dataset in all_datasets: print( f" - Provider: {dataset['provider']}, Region: {dataset['region']}, DB: {dataset['db_name']}" )