""" Task Status tab UI component. """ import json from collections import defaultdict import pandas as pd import plotly.graph_objects as go import streamlit as st from src.database import db from src.services.cluster_utils import extract_cluster_ids_from_execution from src.services.report.metrics_fetcher import fetch_all_diagrams_with_cache, get_cluster_meta_sync from src.services.report.metrics_parser import compute_spu_identifier from src.services.workload_templates import canonicalize_workload_template_or_default def render_plotly_chart(df: pd.DataFrame, queries: list, title: str): """ Render an interactive Plotly chart with support for: - Interactive legends (click to hide/show) - Dual y-axes for different units - Proper unit labels Args: df: DataFrame with time series data (index is timestamp) queries: List of query configs with 'name' and 'unit' title: Chart title """ if df.empty: st.info(f"No data available for {title}") return # Build a mapping from query name to unit # Handle both full names and base names (without :by_xxx or :total) query_to_unit = {} for q in queries: query_name = q["name"] base_name = query_name.split(":")[0] query_to_unit[base_name] = q["unit"] query_to_unit[query_name] = q["unit"] # Group columns by unit units = defaultdict(list) for col in df.columns: # Try to match column with query matched_unit = None # First, try exact match if col in query_to_unit: matched_unit = query_to_unit[col] # If no exact match, try partial match if matched_unit is None: for query_name, unit in query_to_unit.items(): if query_name in col or col in query_name: matched_unit = unit break # If still no match, check label_replace values in expr if matched_unit is None: for q in queries: expr = q["expr"] # Check if this column name appears in label_replace if "label_replace" in expr and f'"{col}"' in expr: matched_unit = q["unit"] break # Default to first unit if still no match if matched_unit is None: matched_unit = queries[0]["unit"] if queries else "" units[matched_unit].append(col) # Create figure fig = go.Figure() # Add traces for each unit unit_list = list(units.keys()) for unit_idx, (unit, columns) in enumerate(units.items()): for col in columns: # Configure which y-axis to use and line style if unit_idx == 0: # First unit uses default y-axis (left) with solid lines yaxis_param = None # Uses y by default line_dash = "solid" else: # Second+ units use y2, y3, etc. (right) with dashed lines yaxis_param = f"y{unit_idx + 1}" # Alternate between dash and dot for multiple right axes line_dash = "dash" if unit_idx == 1 else "dot" # Add unit to legend name (skip if unit is "1") if unit == "1": legend_name = col else: legend_name = f"{col} ({unit})" fig.add_trace( go.Scatter( x=df.index, y=df[col], name=legend_name, mode="lines", yaxis=yaxis_param, line=dict(dash=line_dash, width=2), ) ) # Configure layout layout_config = { "title": title, "xaxis": {"title": "Time"}, "hovermode": "x unified", "legend": { "orientation": "h", "yanchor": "bottom", "y": -0.3, "xanchor": "left", "x": 0, }, } # Add y-axis configurations if len(unit_list) == 1: # Single unit - left axis only layout_config["yaxis"] = { "title": unit_list[0] if unit_list[0] != "1" else "", "side": "left", "showgrid": True, } elif len(unit_list) == 2: # Dual y-axes - left (solid lines) and right (dashed lines) layout_config["yaxis"] = { "title": unit_list[0] if unit_list[0] != "1" else "", "side": "left", "showgrid": True, "zeroline": False, } layout_config["yaxis2"] = { "title": unit_list[1] if unit_list[1] != "1" else "", "side": "right", "overlaying": "y", "showgrid": False, "zeroline": False, } else: # More than 2 units - stack on right side layout_config["yaxis"] = { "title": unit_list[0] if unit_list[0] != "1" else "", "side": "left", "showgrid": True, } for i in range(1, len(unit_list)): layout_config[f"yaxis{i + 1}"] = { "title": (unit_list[i] if unit_list[i] != "1" else ""), "side": "right", "overlaying": "y", "showgrid": False, "anchor": "free" if i > 1 else "x", "position": 1.0 if i == 1 else 1.0 + 0.05 * (i - 1), } fig.update_layout(**layout_config) # Display the chart st.plotly_chart(fig, width="stretch") def _cancel_task(task_id: int) -> bool: latest_task = db.get_task_by_id(task_id) if not latest_task: st.error(f"Task ID {task_id} not found") return False status = latest_task.get("status") workflow_execution_id = latest_task.get("workflow_execution_id") if status in {"REQUESTED", "SUBMITTING"}: updated = db.update_task_status_if_current( task_id, "CANCELLING", expected_current_statuses=["REQUESTED", "SUBMITTING"], ) if updated: st.success(f"Task {task_id} cancel requested ({status} → CANCELLING).") return True st.warning("Task status changed. Please refresh and try again.") return False if status in {"PENDING", "RUNNING"}: updated = db.update_task_status_if_current( task_id, "CANCELLING", expected_current_statuses=["PENDING", "RUNNING"], ) if updated: if workflow_execution_id: st.success( f"Task {task_id} cancel requested; scheduler will cancel the TCMS workflow." ) else: st.success( f"Task {task_id} cancel requested; scheduler will finalize cancellation." ) return True st.warning("Task status changed while cancelling. Please refresh to confirm.") return False if status == "CANCELLING": st.info(f"Task {task_id} is already cancelling.") return False st.info(f"Task {task_id} is {status}, cannot cancel.") return False @st.fragment def render_tab2(): """Dashboard tab - wrapped in fragment for isolation.""" st.header("Task List") if st.button("Refresh Data"): st.rerun(scope="fragment") tasks = db.get_all_tasks() if tasks: df = pd.DataFrame(tasks) # Use workload_template as workload_type df["workload_type"] = df["workload_template"].apply( canonicalize_workload_template_or_default ) # Add SPU identifier column df["spu"] = df.apply( lambda row: compute_spu_identifier( tidb_spu=row["tidb_spu"], tidb_count=row["tidb_count"], tikv_spu=row["tikv_spu"], tikv_count=row["tikv_count"], cloud_provider=row["cloud_provider"], region=row["region"], cluster_version=row["cluster_version"], workload_config=row.get("workload_config"), ), axis=1, ) # Filter options col_f1, col_f2, col_f3 = st.columns(3) with col_f1: filter_user = st.multiselect( "Filter by User", options=df["user_email"].unique() ) with col_f2: filter_status = st.multiselect( "Filter by Status", options=df["status"].unique() ) with col_f3: filter_workload = st.multiselect( "Filter by Workload Type", options=df["workload_type"].unique() ) if filter_user: df = df[df["user_email"].isin(filter_user)] if filter_status: df = df[df["status"].isin(filter_status)] if filter_workload: df = df[df["workload_type"].isin(filter_workload)] # Display main table (without index and result_type) st.dataframe( df[ [ "id", "user_email", "region", "spu", "workload_type", "status", "created_at", ] ], width="stretch", hide_index=True, ) st.markdown("---") # Detail view st.subheader("Task Details") selected_task_id = st.text_input( "Enter Task ID to view details", placeholder="e.g., 15", key="task_id_input", ) # Convert to integer if valid if selected_task_id: try: selected_task_id = int(selected_task_id) except ValueError: st.error("Please enter a valid Task ID (integer)") selected_task_id = None if selected_task_id: # Check if task exists matching_tasks = df[df["id"] == selected_task_id] if matching_tasks.empty: st.error(f"Task ID {selected_task_id} not found") else: task = matching_tasks.iloc[0] st.markdown(f"**Task ID:** {task['id']}") st.markdown(f"**Status:** {task['status']}") st.markdown(f"**TCMS Workflow ID:** {task['workflow_execution_id']}") st.markdown( f"**TCMS Workflow Name:** {task.get('workflow_execution_name')}" ) st.subheader("Actions") can_cancel = task["status"] in { "REQUESTED", "SUBMITTING", "PENDING", "RUNNING", } confirm_cancel = st.checkbox( "Confirm cancel (this may stop the TCMS workflow execution)", value=False, disabled=not can_cancel, key=f"confirm_cancel_{int(task['id'])}", ) cancel_clicked = st.button( "Cancel Task", type="primary", disabled=(not can_cancel) or (not confirm_cancel), key=f"cancel_task_{int(task['id'])}", ) if cancel_clicked: should_rerun = _cancel_task(int(task["id"])) if should_rerun: st.rerun(scope="fragment") col_d1, col_d2 = st.columns([1, 2]) with col_d1: st.markdown("#### Configuration") st.json( { "cloud": f"{task['cloud_provider']} - {task['region']}", "tidb": f"{task['tidb_count']}x {task['tidb_spu']}", "tikv": f"{task['tikv_count']}x {task['tikv_spu']}", "tiflash": f"{task['tiflash_count']}x {task['tiflash_spu']}", "tiproxy": f"{task['tiproxy_count']}x {task['tiproxy_spu']}", } ) with col_d2: st.markdown("#### Results") # Extract and display Clinic URL if available if task["raw_execution_response"]: try: clusters = {} raw_resp = task["raw_execution_response"] if isinstance(raw_resp, str): raw_resp = json.loads(raw_resp) # Navigate to cluster_id graph = raw_resp.get("graph", {}) completed_tasks = graph.get("completedTasks", []) cluster_id = None for task_item in completed_tasks: node_name = task_item.get("nodeName") if ( "tidbcloud_cluster_" in node_name ): # e.g., tidbcloud_cluster_create, tidbcloud_cluster_mix state_output = task_item.get("stateOutput", {}) testbed_info = state_output.get("testbedInfo", {}) resource_items = testbed_info.get( "resourceItems", [] ) for item in resource_items: if item.get("type") == "TIDBCLOUD_CLUSTER": details = item.get("details", {}) status = details.get("status", {}) cluster_id = status.get("clusterId") clusters[ node_name.removeprefix( "tidbcloud_cluster_" ) ] = cluster_id break # Validate cluster_id is a multi-digit integer string for key, cluster_id in clusters.items(): if ( cluster_id and isinstance(cluster_id, str) and cluster_id.isdigit() and len(cluster_id) > 1 ): cluster_meta = get_cluster_meta_sync(cluster_id) clinic_org_id = cluster_meta.clinic_org_id base_url = "https://staging-clinic.pingcap.com" clinic_url = ( f"{base_url}/portal/#/orgs/{clinic_org_id}/" f"clusters/{cluster_id}" ) if key == "create": st.markdown( f"**Clinic URL:** [{clinic_url}]({clinic_url})" ) else: st.markdown( f"**Clinic URL for {key}:** [{clinic_url}]({clinic_url})" ) except Exception: pass # Silently fail if cluster_id extraction fails if task["result_type"] == "SUCCESS": st.success("Test Succeeded") if task["result_data"]: result_data = task["result_data"] if isinstance(result_data, str): result_data = json.loads(result_data) # Filter nodes to show only specific fields (whitelist approach) if "nodes" in result_data and result_data["nodes"]: st.markdown("##### Execution Nodes") filtered_nodes = [] for node in result_data["nodes"]: # Only keep: node_name, summary, and metadata filtered_node = {} if "node_name" in node: filtered_node["node_name"] = node["node_name"] if "summary" in node: filtered_node["summary"] = node["summary"] # Process metadata: only keep exec_id and workload_tasks if "metadata" in node and isinstance( node["metadata"], dict ): metadata = node["metadata"] filtered_metadata = {} if "exec_id" in metadata: filtered_metadata["exec_id"] = metadata[ "exec_id" ] # Process workload_tasks to only keep 'id' if "workload_tasks" in metadata and isinstance( metadata["workload_tasks"], list ): filtered_metadata["workload_tasks"] = [ {"id": task.get("id")} for task in metadata["workload_tasks"] if isinstance(task, dict) ] if filtered_metadata: filtered_node["metadata"] = ( filtered_metadata ) filtered_nodes.append(filtered_node) st.json(filtered_nodes) # Display other metadata (workload_type, status, etc.) metadata = { "status": result_data.get("status"), "workload_type": result_data.get("workload_type"), "failed_node_name": result_data.get("failed_node_name"), "error": result_data.get("error"), } # Only show non-null metadata metadata = { k: v for k, v in metadata.items() if v is not None } if metadata: st.json(metadata) if task["raw_execution_response"]: with st.expander("View Raw Execution Response"): st.json(task["raw_execution_response"]) elif task["result_type"] == "FAILURE": st.error("Test Failed") if task["result_data"]: result_data = task["result_data"] if isinstance(result_data, str): result_data = json.loads(result_data) # Display error information first error_info = { "failed_node_name": result_data.get("failed_node_name"), "error": result_data.get("error"), } error_info = { k: v for k, v in error_info.items() if v is not None } if error_info: st.markdown("##### Error Information") st.json(error_info) # Filter nodes to show only specific fields (whitelist approach) if "nodes" in result_data and result_data["nodes"]: st.markdown("##### Execution Nodes") filtered_nodes = [] for node in result_data["nodes"]: # Only keep: node_name, summary, and metadata filtered_node = {} if "node_name" in node: filtered_node["node_name"] = node["node_name"] if "summary" in node: filtered_node["summary"] = node["summary"] # Process metadata: only keep exec_id and workload_tasks if "metadata" in node and isinstance( node["metadata"], dict ): metadata = node["metadata"] filtered_metadata = {} if "exec_id" in metadata: filtered_metadata["exec_id"] = metadata[ "exec_id" ] # Process workload_tasks to only keep 'id' if "workload_tasks" in metadata and isinstance( metadata["workload_tasks"], list ): filtered_metadata["workload_tasks"] = [ {"id": task.get("id")} for task in metadata["workload_tasks"] if isinstance(task, dict) ] if filtered_metadata: filtered_node["metadata"] = ( filtered_metadata ) filtered_nodes.append(filtered_node) st.json(filtered_nodes) st.markdown("---") # Metrics Charts Section st.subheader("Cluster Metrics") # Extract cluster IDs from raw execution response if task["raw_execution_response"]: try: # Use common utility function to extract cluster IDs clusters = extract_cluster_ids_from_execution( task["raw_execution_response"] ) if clusters: # Display metrics for each cluster for cluster_key, cluster_id in clusters.items(): if cluster_key == "create": st.markdown( f"#### Metrics for Cluster " f"{cluster_id}" ) else: st.markdown( f"#### Metrics for Cluster " f"{cluster_id} ({cluster_key})" ) # Fetch metrics with st.spinner( f"Loading metrics for cluster " f"{cluster_id}..." ): try: diagrams = fetch_all_diagrams_with_cache( task["id"], cluster_id ) # Display each diagram for ( diagram_name, diagram_data, ) in diagrams.items(): df = diagram_data["dataframe"] queries = diagram_data["queries"] if not df.empty: diagram_title = diagram_name.replace( "_", " " ).title() render_plotly_chart( df, queries, diagram_title ) else: st.info( f"No data available for " f"{diagram_name}" ) except Exception as e: st.error( f"Error loading metrics for " f"cluster {cluster_id}: {str(e)}" ) else: st.info("No cluster IDs found in task execution") except Exception as e: st.warning(f"Could not extract cluster IDs: {str(e)}") else: st.info("No tasks found in the database.")