""" Cluster manager for TiDB cluster operations. """ import json import subprocess import tempfile import time from pathlib import Path from typing import Optional from datetime import datetime, timezone from .utils import Colors, Logger from .config import ClusterManagerConfig class ClusterManager: """Manages TiDB cluster operations.""" def __init__(self, config: ClusterManagerConfig): """ Initialize ClusterManager. Args: config: ClusterManagerConfig instance """ self.tiop_home = config.tiop_home self.script_dir = config.script_dir self.libsonnet_dir = config.script_dir / 'tisonnet' self.binary_manager = config.binary_manager self.run_kubectl = config.run_kubectl self.run_jsonnet = config.run_jsonnet self.namespace = config.namespace def _kubectl_with_namespace(self, args, **kwargs): """Run kubectl with namespace automatically added.""" # Insert namespace flags after the subcommand # e.g., ['get', 'pods'] -> ['get', 'pods', '-n', 'my-namespace'] new_args = list(args) if self.namespace is not None: new_args.extend(['-n', self.namespace]) return self.run_kubectl(new_args, **kwargs) def _kubectl_cmd_hint(self, cmd: str) -> str: """Generate kubectl command hint with namespace if needed.""" if self.namespace is not None: return f"kubectl -n {self.namespace} {cmd}" return f"kubectl {cmd}" def _generate_cluster_yaml(self, name: str, config_file: str) -> str: """Generate cluster YAML from config file.""" config_path = Path(config_file) if not config_path.exists(): Logger.die(f"Config file not found: {config_file}") Logger.info(f"Generating resources from config: {config_file}") # Create a wrapper jsonnet that imports user config and wraps with maker config_dir = config_path.parent config_filename = config_path.name wrapper_jsonnet = f""" local maker = import 'maker.libsonnet'; local clusterName = std.extVar('clusterName'); local userConfig = import '{config_filename}'; // Add tiop management labels to all resources // Do NOT override namespace - let kubectl handle it via -n flag // This respects user's namespace settings in their config local addTiopLabels(resources) = std.map( function(res) res {{ metadata+: {{ labels+: {{ 'app.kubernetes.io/managed-by': 'tiop', 'tiop/cluster': clusterName, }}, }}, }}, resources ); // User should provide an array of configurations // tiop.py will call maker.makeByClusterConfigurations automatically local resources = maker.makeByClusterConfigurations(clusterName, userConfig); addTiopLabels(resources) """ # Write wrapper to temp file in the same directory as config with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonnet', dir=config_dir, delete=False) as f: f.write(wrapper_jsonnet) wrapper_path = f.name try: # Run jsonnet with the cluster name # Namespace is handled by kubectl -n flag, not in the generated resources today = datetime.now(timezone.utc).strftime('%Y%m%d') result = self.run_jsonnet([ wrapper_path, '--ext-str', f'clusterName={name}', '--ext-str', f'utcToday={today}', '-J', str(self.libsonnet_dir), '-y' ], capture_output=True, check=True) return result.stdout except subprocess.CalledProcessError as e: # Show the jsonnet error directly if hasattr(e, 'stderr') and e.stderr: print(e.stderr) Logger.die(f"Failed to generate resources from {config_file}") finally: # Clean up temp file Path(wrapper_path).unlink(missing_ok=True) Path(wrapper_path).unlink(missing_ok=True) def create(self, name: str, config_file: Optional[str] = None): """Create a new TiDB cluster.""" if not name: Logger.die("Cluster name is required") # Determine config file if not config_file: config_file = str(self.script_dir.parent / 'simple' / 'simple-cluster.jsonnet') Logger.info(f"No config file specified, using default: {config_file}") Logger.info(f"Creating TiDB cluster '{name}' with config: {config_file}") # Generate resources using jsonnet Logger.info(f"Generating resources...") yaml_content = self._generate_cluster_yaml(name, config_file) # Apply resources directly using stdin Logger.info("Creating resources in Kubernetes...") try: self._kubectl_with_namespace( ['create', '-f', '-'], input=yaml_content.encode(), check=True ) Logger.success(f"TiDB cluster '{name}' resources created") Logger.info(f"Resources are labeled with app.kubernetes.io/managed-by=tiop") # Wait for cluster to be ready with periodic status updates Logger.info(f"Waiting for TiDB cluster '{name}' to be ready...") Logger.info("This may take several minutes as components are being deployed...") print() # Add blank line for better readability max_wait_time = 600 # 10 minutes check_interval = 10 # Check every 10 seconds elapsed_time = 0 while elapsed_time < max_wait_time: # Check cluster status result = self._kubectl_with_namespace([ 'get', 'cluster', name, '-o', 'jsonpath={.status.conditions[?(@.type=="Available")].status}' ], capture_output=True, check=False) if result.returncode == 0 and result.stdout.strip() == 'True': Logger.success(f"TiDB cluster '{name}' is ready! (took {elapsed_time}s)") break # Get pod status summary pod_result = self._kubectl_with_namespace([ 'get', 'pods', '-l', f'app.kubernetes.io/instance={name}', '-o', 'json' ], capture_output=True, check=False) if pod_result.returncode == 0: try: pods_data = json.loads(pod_result.stdout) pods = pods_data.get('items', []) running = sum(1 for p in pods if p.get('status', {}).get('phase') == 'Running') pending = sum(1 for p in pods if p.get('status', {}).get('phase') == 'Pending') total = len(pods) # Count ready containers ready_containers = 0 total_containers = 0 for pod in pods: container_statuses = pod.get('status', {}).get('containerStatuses', []) total_containers += len(container_statuses) ready_containers += sum(1 for cs in container_statuses if cs.get('ready', False)) Logger.info(f"[{elapsed_time}s] Pods: {running}/{total} Running, {pending} Pending | Containers: {ready_containers}/{total_containers} Ready") except json.JSONDecodeError: Logger.info(f"[{elapsed_time}s] Waiting...") else: Logger.info(f"[{elapsed_time}s] Waiting for pods to be created...") time.sleep(check_interval) elapsed_time += check_interval if elapsed_time >= max_wait_time: Logger.warn(f"Timeout waiting for cluster '{name}' to be ready after {max_wait_time}s") Logger.info(f"Monitor the cluster with: {self._kubectl_cmd_hint(f'get cluster {name} -w')}") Logger.info(f"Check pods with: {self._kubectl_cmd_hint(f'get pods -l app.kubernetes.io/instance={name}')}") else: # Show final status print() # Add blank line Logger.info("Cluster status:") self._kubectl_with_namespace(['get', 'cluster', name]) print() # Add blank line Logger.info("Cluster pods:") self._kubectl_with_namespace(['get', 'pods', '-l', f'app.kubernetes.io/instance={name}']) except subprocess.CalledProcessError: Logger.die(f"Failed to create TiDB cluster '{name}'") Logger.info(f"Check pods with: {self._kubectl_cmd_hint(f'get pods -l app.kubernetes.io/instance={name}')}") def apply(self, name: str, config_file: Optional[str] = None): """Apply or update a TiDB cluster.""" if not name: Logger.die("Cluster name is required") # Config file is now required since we don't save references if not config_file: Logger.die("Config file is required. Use --config to specify the configuration file.") Logger.info(f"Applying TiDB cluster '{name}' with config: {config_file}") # Generate resources using jsonnet Logger.info(f"Generating resources...") yaml_content = self._generate_cluster_yaml(name, config_file) # Apply resources directly using stdin Logger.info("Applying resources to Kubernetes...") try: self._kubectl_with_namespace( ['apply', '-f', '-'], input=yaml_content.encode(), capture_output=False, check=True ) Logger.success(f"TiDB cluster '{name}' applied successfully") Logger.info(f"Resources are labeled with app.kubernetes.io/managed-by=tiop") Logger.info(f"Monitor the cluster with: {self._kubectl_cmd_hint(f'get cluster {name} -w')}") Logger.info(f"Check pods with: {self._kubectl_cmd_hint(f'get pods -l app.kubernetes.io/instance={name}')}") except subprocess.CalledProcessError: Logger.die(f"Failed to apply TiDB cluster '{name}'") def delete(self, name: str): """Delete a TiDB cluster.""" if not name: Logger.die("Cluster name is required") # Check if cluster exists by querying with label selector Logger.info(f"Checking if cluster '{name}' exists...") result = self._kubectl_with_namespace([ 'get', 'cluster', name, '-o', 'jsonpath={.metadata.labels.tiop\.pingcap\.com/cluster}' ], capture_output=True, check=False) if result.returncode != 0: Logger.die(f"Cluster '{name}' not found") # Verify it's managed by tiop managed_by_result = self._kubectl_with_namespace([ 'get', 'cluster', name, '-o', 'jsonpath={.metadata.labels.app\.kubernetes\.io/managed-by}' ], capture_output=True, check=False) if managed_by_result.returncode == 0 and managed_by_result.stdout.strip() != 'tiop': Logger.warn(f"Cluster '{name}' is not managed by tiop (managed-by: {managed_by_result.stdout.strip()})") Logger.warn("This command only deletes clusters created by tiop") Logger.die("Operation aborted") response = input(f"{Colors.YELLOW}This will delete the TiDB cluster '{name}'. Are you sure? (y/N){Colors.NC} ") if response.lower() not in ['y', 'yes']: Logger.info("Operation cancelled") return Logger.info(f"Deleting TiDB cluster '{name}' and all its resources...") try: tidb_operator_types = [ "cluster", "pdgroups", "tikvgroups", "tiflashgroups", "tsogroups", "schedulergroups", "tiproxygroups", "tidbgroups", "ticdcgroups", "tibrgcs", "tibrs", 'configmaps', 'ingress', 'services', 'deployments', 'pvc', 'hpa' ] self._kubectl_with_namespace([ 'delete', ','.join(tidb_operator_types), '-l', f'tiop/cluster={name},app.kubernetes.io/managed-by=tiop' ], capture_output=False, check=False) # Don't fail if resource type doesn't exist Logger.success(f"TiDB cluster '{name}' deleted") except subprocess.CalledProcessError: Logger.error(f"Failed to delete TiDB cluster '{name}'") def status(self, name: str): """Show TiDB cluster status.""" if not name: Logger.die("Cluster name is required") Logger.info(f"Checking status of TiDB cluster '{name}'...") # Check if cluster exists result = self._kubectl_with_namespace(['get', 'cluster', name], capture_output=True, check=False) if result.returncode != 0: Logger.error(f"TiDB cluster '{name}' not found") return print() Logger.info("Cluster resource:") self._kubectl_with_namespace(['get', 'cluster', name, '-o', 'wide']) print() Logger.info("All resources:") self._kubectl_with_namespace(['get', 'all', '-l', f'app.kubernetes.io/instance={name}']) def list(self): """List all TiDB clusters.""" Logger.info("Listing TiDB clusters managed by tiop...") # Get all clusters with tiop management label result = self._kubectl_with_namespace([ 'get', 'cluster', '-l', 'app.kubernetes.io/managed-by=tiop', '-o', 'wide' ], capture_output=True, check=False) if result.returncode == 0: if result.stdout.strip(): print(result.stdout) else: Logger.info("No clusters managed by tiop found") else: Logger.error("Failed to list clusters") print() Logger.info("To see all clusters (including non-tiop managed):") Logger.info(f" {self._kubectl_cmd_hint('get cluster')}") def inspect(self, name: Optional[str] = None, config_file: Optional[str] = None): """Inspect cluster configuration without applying.""" if not name and not config_file: Logger.die("Either cluster name or config file is required") # If name is provided without config, try to get from Kubernetes if name and not config_file: Logger.info(f"Fetching current configuration for cluster '{name}' from Kubernetes...") result = self._kubectl_with_namespace([ 'get', 'cluster', name, '-o', 'yaml' ], capture_output=True, check=False) if result.returncode == 0: Logger.info(f"Current Cluster resource for '{name}':") Logger.info("=" * 60) print(result.stdout) # Also show related resources print() Logger.info("=" * 60) Logger.info(f"Related resources (managed by tiop):") Logger.info("=" * 60) self._kubectl_with_namespace([ 'get', 'all', '-l', f'tiop/cluster={name}', '-l', 'app.kubernetes.io/managed-by=tiop' ]) return else: Logger.die(f"Cluster '{name}' not found in Kubernetes. Use --config to generate from a config file.") # If config file is provided, generate and show the YAML if not config_file: config_file = str(self.script_dir.parent / 'simple' / 'simple-cluster.jsonnet') Logger.info(f"No config file specified, using default: {config_file}") # Use name from parameter or extract from config cluster_name = name if name else 'preview' Logger.info(f"Inspecting configuration from: {config_file}") Logger.info("=" * 60) yaml_content = self._generate_cluster_yaml(cluster_name, config_file) print(yaml_content) def template(self, template_name: Optional[str] = None): """Print available templates or show a specific template.""" templates_dir = self.script_dir / 'templates' if not templates_dir.exists(): Logger.die(f"Templates directory not found: {templates_dir}") # List available templates available_templates = sorted([f.stem for f in templates_dir.glob('*.jsonnet')]) if not available_templates: Logger.die(f"No templates found in {templates_dir}") # If no template specified, list all available templates if not template_name: Logger.info("Available templates:") for tmpl in available_templates: Logger.info(f" - {tmpl}") Logger.info("") Logger.info("Use 'tiop.py cluster template ' to view a specific template") return # Show specific template template_file = templates_dir / f"{template_name}.jsonnet" if not template_file.exists(): Logger.error(f"Template '{template_name}' not found") Logger.info("Available templates:") for tmpl in available_templates: Logger.info(f" - {tmpl}") return Logger.info(f"Template: {template_name}") Logger.info(f"Path: {template_file}") Logger.info("=" * 60) print(template_file.read_text())