import base64 import gzip import json import pickle from contextlib import contextmanager from typing import Dict, List, Literal, Optional, Sequence import boto3 from botocore.exceptions import ClientError, NoCredentialsError, TokenRetrievalError from environs import Env from sqlalchemy import create_engine, or_ from sqlalchemy.orm import sessionmaker from sqlalchemy.sql import func from src.models.models import Base, ClusterMetrics, TestResult, TestTask # Load environment variables env = Env() env.read_env() DATABASE_URL = env("SPU_STORAGE_DSN") # TiDB has a practical single-row size limit (~6MiB). Storing full TCMS execution # payloads (often containing large logs) into JSON columns can exceed it. # These limits are guardrails; tune via env if you need to keep more data. TEST_RESULT_ROW_MAX_BYTES = env.int("SPU_TEST_RESULT_ROW_MAX_BYTES", 6_000_000) TEST_RESULT_JSON_FIELD_MAX_BYTES = env.int( "SPU_TEST_RESULT_JSON_FIELD_MAX_BYTES", 5_000_000 ) TEST_RESULT_ERROR_MAX_BYTES = env.int("SPU_TEST_RESULT_ERROR_MAX_BYTES", 64_000) TEST_RESULT_SIZE_DEBUG = env.bool("SPU_TEST_RESULT_SIZE_DEBUG", False) def _json_size_bytes(value) -> int: if value is None: return 0 try: payload = json.dumps( value, ensure_ascii=False, separators=(",", ":"), default=str, ).encode("utf-8") return len(payload) except Exception: return len(str(value).encode("utf-8", errors="ignore")) def _truncate_text_to_bytes(text: Optional[str], max_bytes: int) -> Optional[str]: if text is None: return None encoded = text.encode("utf-8", errors="ignore") if len(encoded) <= max_bytes: return text suffix = "\n...[TRUNCATED]..." suffix_bytes = suffix.encode("utf-8") # If the suffix alone would exceed the limit, just hard-clamp the original text. if len(suffix_bytes) >= max_bytes: clipped = encoded[:max_bytes].decode("utf-8", errors="ignore") return clipped allowed_content_bytes = max_bytes - len(suffix_bytes) clipped = encoded[:allowed_content_bytes].decode("utf-8", errors="ignore") return f"{clipped}{suffix}" def _deep_truncate_json( value, *, max_depth: int = 8, max_keys: int = 200, max_list_items: int = 200, max_string_bytes: int = 16_000, ): if value is None: return None if max_depth <= 0: return "...[TRUNCATED:DEPTH]..." if isinstance(value, str): return _truncate_text_to_bytes(value, max_string_bytes) if isinstance(value, (int, float, bool)): return value if isinstance(value, dict): truncated = {} for idx, (key, item) in enumerate(value.items()): if idx >= max_keys: truncated["...[TRUNCATED:KEYS]..."] = len(value) - max_keys break truncated[key] = _deep_truncate_json( item, max_depth=max_depth - 1, max_keys=max_keys, max_list_items=max_list_items, max_string_bytes=max_string_bytes, ) return truncated if isinstance(value, (list, tuple)): sliced = value[:max_list_items] truncated_list = [ _deep_truncate_json( item, max_depth=max_depth - 1, max_keys=max_keys, max_list_items=max_list_items, max_string_bytes=max_string_bytes, ) for item in sliced ] if len(value) > max_list_items: truncated_list.append( {"...[TRUNCATED:LIST_ITEMS]...": len(value) - max_list_items} ) return truncated_list # Fallback for non-JSON-serializable objects return _truncate_text_to_bytes(str(value), max_string_bytes) def _extract_tcms_cluster_ids(raw_execution_response) -> Dict[str, str]: if not isinstance(raw_execution_response, dict): return {} graph = raw_execution_response.get("graph", {}) if not isinstance(graph, dict): return {} completed_tasks = graph.get("completedTasks", []) if not isinstance(completed_tasks, list): return {} clusters: Dict[str, str] = {} for task_item in completed_tasks: if not isinstance(task_item, dict): continue node_name = task_item.get("nodeName") or "" if "tidbcloud_cluster_" not in node_name: continue state_output = task_item.get("stateOutput", {}) if not isinstance(state_output, dict): continue testbed_info = state_output.get("testbedInfo", {}) if not isinstance(testbed_info, dict): continue resource_items = testbed_info.get("resourceItems", []) if not isinstance(resource_items, list): continue for item in resource_items: if not isinstance(item, dict): continue if item.get("type") != "TIDBCLOUD_CLUSTER": continue details = item.get("details", {}) if not isinstance(details, dict): continue status = details.get("status", {}) if not isinstance(status, dict): continue cluster_id = status.get("clusterId") if ( cluster_id and isinstance(cluster_id, str) and cluster_id.isdigit() and len(cluster_id) > 1 ): key = node_name.removeprefix("tidbcloud_cluster_") clusters[key] = cluster_id break return clusters def _summarize_tcms_execution_response(raw_execution_response): if not isinstance(raw_execution_response, dict): return raw_execution_response summary = {} for key in ( "id", "name", "status", "createdAt", "updatedAt", "workflowID", "runnerID", "creatorEmail", "output", ): if key in raw_execution_response: summary[key] = raw_execution_response.get(key) clusters = _extract_tcms_cluster_ids(raw_execution_response) if clusters: summary["clusters"] = clusters if "output" in summary and isinstance(summary["output"], str): summary["output"] = _truncate_text_to_bytes(summary["output"], 16_000) return summary def _gzip_base64_encode_text(text: str) -> str: raw = text.encode("utf-8", errors="ignore") compressed = gzip.compress(raw, compresslevel=6) return base64.b64encode(compressed).decode("ascii") def _compress_raw_logs_in_execution_result_dict(result_data): if not isinstance(result_data, dict): return result_data, False nodes = result_data.get("nodes") if not isinstance(nodes, list): return result_data, False changed = False new_nodes = [] for node in nodes: if not isinstance(node, dict): new_nodes.append(node) continue raw_log = node.get("raw_log") # Preserve the public schema: raw_log should remain a string/None. # Store the encoded representation in a separate internal field. if isinstance(node.get("_raw_log_encoded"), dict): new_nodes.append(node) continue # Normalize legacy dict schema if encountered (older versions may have # rewritten raw_log into a dict). if isinstance(raw_log, dict): copied = dict(node) if raw_log.get("encoding") == "gzip_base64" and isinstance( raw_log.get("data"), str ): copied["_raw_log_encoded"] = raw_log copied["raw_log"] = "[raw_log compressed]" elif raw_log.get("truncated") is True: copied.pop("_raw_log_encoded", None) copied["raw_log"] = "[raw_log truncated]" else: copied.pop("_raw_log_encoded", None) copied["raw_log"] = None new_nodes.append(copied) changed = True continue if not isinstance(raw_log, str) or not raw_log: new_nodes.append(node) continue raw_bytes = len(raw_log.encode("utf-8", errors="ignore")) encoded = _gzip_base64_encode_text(raw_log) encoded_bytes = len(encoded.encode("ascii", errors="ignore")) # Base64 adds ~33% overhead; only keep encoded form when it actually helps. if encoded_bytes >= raw_bytes: new_nodes.append(node) continue copied = dict(node) copied["_raw_log_encoded"] = {"encoding": "gzip_base64", "data": encoded} copied["raw_log"] = f"[raw_log compressed; original size: {raw_bytes} bytes]" new_nodes.append(copied) changed = True if not changed: return result_data, False updated = dict(result_data) updated["nodes"] = new_nodes return updated, True def _truncate_raw_logs_in_execution_result_dict(result_data): if not isinstance(result_data, dict): return result_data, False nodes = result_data.get("nodes") if not isinstance(nodes, list): return result_data, False changed = False new_nodes = [] for node in nodes: if not isinstance(node, dict): new_nodes.append(node) continue raw_log = node.get("raw_log") if raw_log is None: new_nodes.append(node) continue # Already truncated. if raw_log == "[raw_log truncated]": new_nodes.append(node) continue copied = dict(node) copied.pop("_raw_log_encoded", None) copied["raw_log"] = "[raw_log truncated]" new_nodes.append(copied) changed = True if not changed: return result_data, False updated = dict(result_data) updated["nodes"] = new_nodes return updated, True def _fit_test_result_payload( *, result_data, error_msg: Optional[str], raw_execution_response, ): if isinstance(error_msg, str): error_msg = _truncate_text_to_bytes(error_msg, TEST_RESULT_ERROR_MAX_BYTES) # Prefer keeping parsed result_data; shrink raw_execution_response first. def _total_size(): return ( _json_size_bytes(result_data) + _json_size_bytes(raw_execution_response) + ( len(error_msg.encode("utf-8", errors="ignore")) if isinstance(error_msg, str) else 0 ) ) if TEST_RESULT_SIZE_DEBUG: error_bytes = ( len(error_msg.encode("utf-8", errors="ignore")) if isinstance(error_msg, str) else 0 ) print( "[save_result] payload sizes (bytes): " f"result_data={_json_size_bytes(result_data)}, " f"raw_execution_response={_json_size_bytes(raw_execution_response)}, " f"error_msg={error_bytes}, " f"total={_total_size()}, " f"limit={TEST_RESULT_ROW_MAX_BYTES}" ) if _total_size() <= TEST_RESULT_ROW_MAX_BYTES: return result_data, error_msg, raw_execution_response # 1) Shrink raw execution response first. if raw_execution_response is not None: raw_bytes = _json_size_bytes(raw_execution_response) # Only summarize raw when it is meaningfully large; otherwise keep the full # raw payload and shrink other fields (e.g. result_data raw logs) instead. if raw_bytes > min(TEST_RESULT_JSON_FIELD_MAX_BYTES, 512_000): raw_execution_response = _summarize_tcms_execution_response( raw_execution_response ) if TEST_RESULT_SIZE_DEBUG: error_bytes = ( len(error_msg.encode("utf-8", errors="ignore")) if isinstance(error_msg, str) else 0 ) print( "[save_result] after raw shrink (bytes): " f"result_data={_json_size_bytes(result_data)}, " f"raw_execution_response={_json_size_bytes(raw_execution_response)}, " f"error_msg={error_bytes}, " f"total={_total_size()}" ) if _total_size() <= TEST_RESULT_ROW_MAX_BYTES: return result_data, error_msg, raw_execution_response # 2) Then shrink error message. if isinstance(error_msg, str): error_msg = _truncate_text_to_bytes(error_msg, 8_000) if TEST_RESULT_SIZE_DEBUG: print( "[save_result] after error shrink (bytes): " f"result_data={_json_size_bytes(result_data)}, " f"raw_execution_response={_json_size_bytes(raw_execution_response)}, " f"error_msg={len(error_msg.encode('utf-8', errors='ignore'))}, " f"total={_total_size()}" ) if _total_size() <= TEST_RESULT_ROW_MAX_BYTES: return result_data, error_msg, raw_execution_response # 3) Finally shrink result_data. if result_data is not None: # First try compressing raw logs (lossless). compressed, changed = _compress_raw_logs_in_execution_result_dict(result_data) if changed: result_data = compressed if _total_size() <= TEST_RESULT_ROW_MAX_BYTES: return result_data, error_msg, raw_execution_response # If still too large, drop raw logs but keep the overall result_data shape. truncated_logs, changed = _truncate_raw_logs_in_execution_result_dict(result_data) if changed: result_data = truncated_logs if _total_size() <= TEST_RESULT_ROW_MAX_BYTES: return result_data, error_msg, raw_execution_response result_data = _deep_truncate_json( result_data, max_string_bytes=8_000, max_list_items=100 ) if TEST_RESULT_SIZE_DEBUG: error_bytes = ( len(error_msg.encode("utf-8", errors="ignore")) if isinstance(error_msg, str) else 0 ) print( "[save_result] after result_data shrink (bytes): " f"result_data={_json_size_bytes(result_data)}, " f"raw_execution_response={_json_size_bytes(raw_execution_response)}, " f"error_msg={error_bytes}, " f"total={_total_size()}" ) if _total_size() <= TEST_RESULT_ROW_MAX_BYTES: return result_data, error_msg, raw_execution_response # Last-resort minimal payloads. if raw_execution_response is not None: raw_execution_response = _summarize_tcms_execution_response( raw_execution_response ) if result_data is not None: if isinstance(result_data, dict): # Preserve the overall shape used by consumers. result_data = { "status": result_data.get("status"), "workload_type": result_data.get("workload_type"), "nodes": [], "failed_node_name": result_data.get("failed_node_name"), "error": result_data.get("error"), } else: result_data = None if isinstance(error_msg, str): error_msg = _truncate_text_to_bytes(error_msg, 2_000) return result_data, error_msg, raw_execution_response # Create engine engine = create_engine(DATABASE_URL, pool_pre_ping=True) # Create session factory SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) @contextmanager def get_db(): """Provide a transactional scope around a series of operations.""" db = SessionLocal() try: yield db db.commit() except Exception: db.rollback() raise finally: db.close() def init_db(): """Initialize database tables. Note: Use Alembic for production.""" try: Base.metadata.create_all(bind=engine) print("Database tables created successfully") except Exception as e: print(f"Error initializing database: {e}") def create_task(task_data: Dict) -> Optional[int]: """Creates a new test task.""" try: with get_db() as db: task = TestTask(**task_data) db.add(task) db.flush() return task.id except Exception as e: print(f"Error creating task: {e}") return None def get_pending_tasks() -> List[Dict]: """Retrieves non-deleted tasks with status REQUESTED, SUBMITTING, PENDING, RUNNING or CANCELLING.""" try: with get_db() as db: tasks = ( db.query(TestTask) .filter( TestTask.status.in_( ["REQUESTED", "SUBMITTING", "PENDING", "RUNNING", "CANCELLING"] ), TestTask.deleted == False, ) .all() ) return [task_to_dict(task) for task in tasks] except Exception as e: print(f"Error fetching pending tasks: {e}") return [] def get_ended_tasks() -> List[Dict]: """Retrieves non-deleted tasks with status SUCCESS, FAILURE or ERROR that haven't been parsed yet.""" try: with get_db() as db: # Get tasks that are ended but don't have result_data parsed yet tasks = ( db.query(TestTask) .filter( TestTask.status.in_(["SUCCESS", "FAILURE", "ERROR"]), TestTask.deleted == False, ) .all() ) # Filter tasks that haven't been parsed yet unparsed_tasks = [] for task in tasks: results = ( db.query(TestResult) .filter_by(task_id=task.id) .filter(TestResult.deleted == False) .order_by(TestResult.id.desc()) .all() ) # If no results or parsed is False, it needs parsing # Check the most recent result (first in list) if not results or not results[0].parsed: unparsed_tasks.append(task_to_dict(task)) return unparsed_tasks except Exception as e: print(f"Error fetching ended tasks: {e}") return [] def update_task_status(task_id: int, status: str, workflow_id: Optional[str] = None): """Updates the status and workflow_id of a task. Returns: True if a row was updated, otherwise False. """ return update_task_status_if_current( task_id=task_id, status=status, workflow_id=workflow_id, expected_current_statuses=None, ) def update_task_status_if_current( task_id: int, status: str, workflow_id: Optional[str] = None, expected_current_statuses: Optional[Sequence[str]] = None, workflow_name: Optional[str] = None, ) -> bool: """Updates task status (and optional workflow_id) only if current status matches. This is used as an optimistic concurrency control mechanism to avoid overwriting newer status changes (e.g., user-cancel vs. scheduler updates). Args: task_id: Task ID status: New task status workflow_id: Optional workflow execution ID to set expected_current_statuses: If provided, only update when current status is in it Returns: True if a row was updated, otherwise False. """ try: with get_db() as db: query = db.query(TestTask).filter( TestTask.id == task_id, TestTask.deleted == False, ) if expected_current_statuses is not None: query = query.filter( TestTask.status.in_(list(expected_current_statuses)) ) values = {"status": status, "updated_at": func.now()} if workflow_id is not None: values["workflow_execution_id"] = workflow_id if workflow_name is not None: values["workflow_execution_name"] = workflow_name updated_rows = query.update(values, synchronize_session=False) return updated_rows > 0 except Exception as e: print(f"Error updating task status: {e}") return False def get_task_by_id(task_id: int) -> Optional[Dict]: """Retrieves a single non-deleted task by ID.""" try: with get_db() as db: task = ( db.query(TestTask) .filter( TestTask.id == task_id, TestTask.deleted == False, ) .first() ) return task_to_dict(task) if task else None except Exception as e: print(f"Error fetching task by id: {e}") return None def attach_workflow_execution_id_if_missing( task_id: int, workflow_id: str, expected_current_statuses: Optional[Sequence[str]] = None, ) -> bool: """Attach workflow_execution_id if it's not set yet. Returns: True if attached, otherwise False. """ try: with get_db() as db: query = db.query(TestTask).filter( TestTask.id == task_id, TestTask.deleted == False, or_( TestTask.workflow_execution_id.is_(None), TestTask.workflow_execution_id == "", ), ) if expected_current_statuses is not None: query = query.filter( TestTask.status.in_(list(expected_current_statuses)) ) updated_rows = query.update( { "workflow_execution_id": workflow_id, "updated_at": func.now(), }, synchronize_session=False, ) return updated_rows > 0 except Exception as e: print(f"Error attaching workflow_execution_id: {e}") return False def acquire_submitting_lease_if_stale(task_id: int, cutoff_dt) -> bool: """Acquire a lease for SUBMITTING tasks if the last update is older than cutoff_dt. This is used to avoid concurrent re-submission by multiple schedulers. Args: task_id: Task ID cutoff_dt: Only acquire if updated_at < cutoff_dt Returns: True if the lease is acquired (row updated), otherwise False. """ try: with get_db() as db: updated_rows = ( db.query(TestTask) .filter( TestTask.id == task_id, TestTask.deleted == False, TestTask.status == "SUBMITTING", or_( TestTask.workflow_execution_id.is_(None), TestTask.workflow_execution_id == "", ), TestTask.updated_at < cutoff_dt, ) .update({"updated_at": func.now()}, synchronize_session=False) ) return updated_rows > 0 except Exception as e: print(f"Error acquiring submitting lease: {e}") return False def save_result( task_id: int, result_type: str, result_data: Optional[Dict] = None, error_msg: Optional[str] = None, raw_execution_response: Optional[Dict] = None, parsed: bool = False, ): """Saves or updates the result of a task (upsert operation).""" try: result_data, error_msg, raw_execution_response = _fit_test_result_payload( result_data=result_data, error_msg=error_msg, raw_execution_response=raw_execution_response, ) with get_db() as db: # Check if result already exists for this task # Use the most recent one (highest id) if multiple exist existing_result = ( db.query(TestResult) .filter_by(task_id=task_id) .filter(TestResult.deleted == False) .order_by(TestResult.id.desc()) .first() ) if existing_result: # Update existing result existing_result.result_type = result_type if result_data is not None: existing_result.result_data = result_data if error_msg is not None: existing_result.error_msg = error_msg if raw_execution_response is not None: existing_result.raw_execution_response = raw_execution_response existing_result.parsed = parsed else: # Create new result result = TestResult( task_id=task_id, result_type=result_type, result_data=result_data, error_msg=error_msg, raw_execution_response=raw_execution_response, parsed=parsed, ) db.add(result) except Exception as e: print(f"Error saving result: {e}") def get_all_tasks() -> List[Dict]: """Retrieves all non-deleted tasks joined with their results.""" try: with get_db() as db: # Only get non-deleted tasks tasks = db.query(TestTask).filter(TestTask.deleted == False).all() result_list = [] for task in tasks: task_dict = task_to_dict(task) # Get associated results (only non-deleted results) # Use the most recent one (highest id) if multiple exist results = ( db.query(TestResult) .filter_by(task_id=task.id) .filter(TestResult.deleted == False) .order_by(TestResult.id.desc()) .all() ) if results: result = results[0] # Take most recent result task_dict.update( { "result_type": result.result_type, "result_data": result.result_data, "error_msg": result.error_msg, "raw_execution_response": result.raw_execution_response, } ) else: task_dict.update( { "result_type": None, "result_data": None, "error_msg": None, "raw_execution_response": None, } ) result_list.append(task_dict) return result_list except Exception as e: print(f"Error fetching all tasks: {e}") return [] def task_to_dict(task: TestTask) -> Dict: """Convert TestTask object to dictionary.""" return { "id": task.id, "user_email": task.user_email, "cloud_provider": task.cloud_provider, "region": task.region, "cluster_version": task.cluster_version, "org_id": task.org_id, "project_id": task.project_id, "tidb_spu": task.tidb_spu, "tidb_count": task.tidb_count, "tidb_config": task.tidb_config, "tikv_spu": task.tikv_spu, "tikv_count": task.tikv_count, "tikv_config": task.tikv_config, "tikv_storage_type": task.tikv_storage_type, "tikv_storage_size_gi": task.tikv_storage_size_gi, "tiflash_spu": task.tiflash_spu, "tiflash_count": task.tiflash_count, "tiflash_config": task.tiflash_config, "tiflash_storage_type": task.tiflash_storage_type, "tiflash_storage_size_gi": task.tiflash_storage_size_gi, "pd_spu": task.pd_spu, "pd_count": task.pd_count, "pd_config": task.pd_config, "tiproxy_spu": task.tiproxy_spu, "tiproxy_count": task.tiproxy_count, "tiproxy_config": task.tiproxy_config, "workload_template": task.workload_template, "workload_config": task.workload_config, "workflow_execution_id": task.workflow_execution_id, "workflow_execution_name": task.workflow_execution_name, "status": task.status, "deleted": task.deleted, "created_at": task.created_at, "updated_at": task.updated_at, } def compress_diagrams( diagrams: Dict, format_type: Literal["pickle", "json"] = "json" ) -> bytes: """ Compress diagrams dictionary (with DataFrames) to bytes. Args: diagrams: Dict mapping diagram names to their data dict with: - 'dataframe': pd.DataFrame - 'queries': list of query configs format_type: Serialization format - 'pickle' or 'json'(default) - 'pickle': Python-only, smallest size, supports any Python object - 'json': Cross-language compatible, slightly larger (~20-30% more) Returns: Compressed bytes (format + gzip) """ if format_type == "pickle": serialized = pickle.dumps(diagrams) elif format_type == "json": # Convert DataFrames to JSON-serializable format json_data = { name: { "dataframe": data["dataframe"].to_dict(orient="records"), "queries": data["queries"], } for name, data in diagrams.items() } serialized = json.dumps(json_data, default=str).encode("utf-8") else: raise ValueError(f"Unsupported format_type: {format_type}") return gzip.compress(serialized) def decompress_diagrams( data: bytes, format_type: Literal["pickle", "json"] = "json" ) -> Dict: """ Decompress bytes back to diagrams dictionary. Args: data: Compressed bytes from compress_diagrams() format_type: Serialization format - 'pickle' or 'json'(default) Must match the format used in compress_diagrams() Returns: Dict mapping diagram names to their data dict """ decompressed = gzip.decompress(data) if format_type == "pickle": return pickle.loads(decompressed) elif format_type == "json": # Import pandas here to avoid circular dependency import pandas as pd json_data = json.loads(decompressed.decode("utf-8")) # Convert back to DataFrames diagrams = { name: { "dataframe": pd.DataFrame(data["dataframe"]), "queries": data["queries"], } for name, data in json_data.items() } return diagrams else: raise ValueError(f"Unsupported format_type: {format_type}") def get_s3_client(): """ Create S3 client that works in both local (SSO) and EC2 (instance profile) environments. Behavior: - If AWS_PROFILE is set in .env: Uses that profile (ignores env vars like AWS_ACCESS_KEY_ID) - If AWS_PROFILE not set: Uses default credential chain (env vars or EC2 instance profile) - If METRICS_S3_ROLE_ARN is set: Assumes that role using above credentials Returns: boto3 S3 client Raises: TokenRetrievalError: If local SSO token expired (run: aws sso login) NoCredentialsError: If no credentials found """ region = env.str("METRICS_S3_REGION", "us-west-2") profile_name = env.str("AWS_PROFILE", None) # None if not set in .env try: # Create session with explicit profile if set in .env # This forces boto3 to use the profile, ignoring AWS_ACCESS_KEY_ID env vars if profile_name: session = boto3.Session(profile_name=profile_name, region_name=region) else: # Use default credential chain (env vars or EC2 instance profile) session = boto3.Session(region_name=region) return session.client("s3") except TokenRetrievalError as e: print(f"❌ AWS SSO token has expired. Please run:") print(f" aws sso login --profile {profile_name or 'default'}") raise except NoCredentialsError as e: print(f"❌ No AWS credentials found. Please run:") print(f" aws sso login --profile {profile_name or 'default'}") raise def upload_metrics_to_s3(task_id: int, cluster_id: str, data: bytes) -> Optional[str]: """ Upload compressed metrics data to S3. Args: task_id: Task ID cluster_id: Cluster ID data: Compressed metrics bytes (from compress_diagrams) Returns: S3 URI (e.g., s3://bucket/path/123/cluster-456.json.gz) or None on error """ bucket = env.str("METRICS_S3_BUCKET") # Support optional path prefix (e.g., "benchmark-metrics") path_prefix = env.str("METRICS_S3_PATH", "benchmark-metrics").strip("/") s3_key = f"{path_prefix}/{task_id}/{cluster_id}.json.gz" try: s3_client = get_s3_client() s3_client.put_object( Bucket=bucket, Key=s3_key, Body=data, ContentType="application/gzip", Metadata={ "task_id": str(task_id), "cluster_id": cluster_id, }, ) s3_uri = f"s3://{bucket}/{s3_key}" print(f"✓ Uploaded metrics to S3: {s3_uri}") return s3_uri except ClientError as e: print(f"Error uploading to S3: {e}") return None except Exception as e: print(f"Unexpected error uploading to S3: {e}") return None def download_metrics_from_s3(s3_uri: str) -> Optional[bytes]: """ Download compressed metrics data from S3. Args: s3_uri: S3 URI (e.g., s3://bucket/metrics/123/cluster-456.json.gz) Returns: Compressed metrics bytes or None on error """ # Parse s3://bucket/key if not s3_uri.startswith("s3://"): print(f"Invalid S3 URI format: {s3_uri}") return None parts = s3_uri.replace("s3://", "").split("/", 1) if len(parts) != 2: print(f"Invalid S3 URI format: {s3_uri}") return None bucket, key = parts[0], parts[1] try: s3_client = get_s3_client() response = s3_client.get_object(Bucket=bucket, Key=key) data = response["Body"].read() print(f"✓ Downloaded metrics from S3: {s3_uri}") return data except ClientError as e: if e.response["Error"]["Code"] == "NoSuchKey": print(f"S3 object not found: {s3_uri}") else: print(f"Error downloading from S3: {e}") return None except Exception as e: print(f"Unexpected error downloading from S3: {e}") return None def save_cluster_metrics( task_id: int, cluster_id: str, diagrams: Dict, format_type: Literal["pickle", "json"] = "json", use_s3: bool = True, ) -> Optional[int]: """ Save cluster metrics to database (upsert operation). If metrics already exist for (task_id, cluster_id), update them. Otherwise create new record. Args: task_id: Task ID cluster_id: Cluster ID diagrams: Diagrams dict from fetch_all_diagrams_for_cluster() format_type: Serialization format - 'pickle' or 'json'(default) use_s3: If True, store data in S3 and save URI to database (default) If False, store data directly in database (legacy mode) Returns: ClusterMetrics record ID, or None on error """ try: compressed_data = compress_diagrams(diagrams, format_type=format_type) with get_db() as db: # Check if record already exists existing = ( db.query(ClusterMetrics) .filter_by( task_id=task_id, cluster_id=cluster_id, deleted=False, ) .first() ) if use_s3: # Upload to S3 and save URI s3_uri = upload_metrics_to_s3(task_id, cluster_id, compressed_data) if not s3_uri: print(f"Failed to upload to S3, falling back to database storage") # Fall back to database storage if existing: existing.metrics_data = compressed_data existing.metrics_s3_uri = None return existing.id else: metrics = ClusterMetrics( task_id=task_id, cluster_id=cluster_id, metrics_data=compressed_data, ) db.add(metrics) db.flush() return metrics.id # Successfully uploaded to S3 if existing: existing.metrics_s3_uri = s3_uri return existing.id else: metrics = ClusterMetrics( task_id=task_id, cluster_id=cluster_id, metrics_s3_uri=s3_uri, ) db.add(metrics) db.flush() return metrics.id else: # Legacy mode: store directly in database if existing: existing.metrics_data = compressed_data existing.metrics_s3_uri = None return existing.id else: metrics = ClusterMetrics( task_id=task_id, cluster_id=cluster_id, metrics_data=compressed_data, ) db.add(metrics) db.flush() return metrics.id except Exception as e: print(f"Error saving cluster metrics: {e}") return None def get_cluster_metrics( task_id: int, cluster_id: str, format_type: Literal["pickle", "json"] = "json" ) -> Optional[Dict]: """ Retrieve cluster metrics from S3. Only reads from S3 storage. If metrics_s3_uri is not set or S3 download fails, returns None to trigger a re-fetch from Clinic API. Args: task_id: Task ID cluster_id: Cluster ID format_type: Serialization format - 'pickle' or 'json'(default) Must match the format used when saving Returns: Diagrams dict, or None if not found """ try: with get_db() as db: metrics = ( db.query(ClusterMetrics) .filter_by( task_id=task_id, cluster_id=cluster_id, deleted=False, ) .first() ) if not metrics: return None # Only use S3 - ignore metrics_data column completely if metrics.metrics_s3_uri: compressed_data = download_metrics_from_s3(metrics.metrics_s3_uri) if compressed_data: return decompress_diagrams(compressed_data, format_type=format_type) else: print( f"Failed to download from S3 for task_id={task_id}, cluster_id={cluster_id}" ) return None # No S3 URI - return None to trigger re-fetch print(f"No S3 URI found for task_id={task_id}, cluster_id={cluster_id}") return None except Exception as e: print(f"Error retrieving cluster metrics: {e}") return None if __name__ == "__main__": init_db()