From e3172ad5b6bf1436a63ee3b4d254f5f99804f7d3 Mon Sep 17 00:00:00 2001 From: "Matthew C. Morgan" Date: Tue, 25 Mar 2025 20:55:54 -0400 Subject: [PATCH] added emr.py --- .../aws_resource_management/managers/emr.py | 176 ++++++++++++++++++ 1 file changed, 176 insertions(+) diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py index d38345d0..bb214a0b 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py @@ -26,6 +26,182 @@ def __init__(self, credentials: Dict[str, str], account_id: str, account_name: O super().__init__(credentials, account_id, account_name) self.resource_type = "emr_cluster" + def discover_clusters(self, region: str, cluster_states: List[str] = None) -> List[Dict[str, Any]]: + """ + Discover EMR clusters in the given region with the specified states. + + Args: + region: AWS region + cluster_states: List of EMR cluster states to filter by + + Returns: + List of EMR cluster dictionaries + """ + try: + # Default states if not provided - EMR API only accepts specific states + if not cluster_states: + cluster_states = ['STARTING', 'BOOTSTRAPPING', 'RUNNING', 'WAITING', 'TERMINATING'] + + logger.debug(f"Discovering EMR clusters in {region} with states: {', '.join(cluster_states)}") + + emr_client = self.create_client('emr', region) + clusters = [] + + # ListClusters API only returns a subset of data, with pagination + paginator = emr_client.get_paginator('list_clusters') + page_iterator = paginator.paginate(ClusterStates=cluster_states) + + for page in page_iterator: + if 'Clusters' in page: + for cluster_summary in page['Clusters']: + # Get detailed info for each cluster + try: + cluster_id = cluster_summary.get('Id') + if not cluster_id: + logger.warning("Skipping cluster with missing ID") + continue + + # Format basic info + cluster = { + 'id': cluster_id, + 'name': cluster_summary.get('Name', 'Unknown'), + 'status': cluster_summary.get('Status', {}).get('State', 'UNKNOWN'), + 'region': region, + 'account_id': self.account_id, + 'account_name': self.account_name + } + + # Get tags + try: + response = emr_client.describe_cluster(ClusterId=cluster_id) + if 'Cluster' in response and 'Tags' in response['Cluster']: + tags = {} + for tag in response['Cluster']['Tags']: + tags[tag.get('Key')] = tag.get('Value') + cluster['tags'] = tags + except Exception as e: + logger.warning(f"Could not get tags for cluster {cluster_id}: {str(e)}") + cluster['tags'] = {} + + clusters.append(cluster) + + except Exception as e: + logger.warning(f"Error processing EMR cluster: {str(e)}") + + logger.info(f"Discovered {len(clusters)} EMR clusters in {region}") + return clusters + + except ClientError as e: + error_code = e.response.get('Error', {}).get('Code') + if error_code == 'AccessDeniedException' or error_code == 'UnauthorizedOperation': + logger.warning(f"Access denied to EMR in region {region}: {str(e)}") + elif error_code == 'EndpointConnectionError' or error_code == 'UnknownEndpoint': + logger.debug(f"EMR not supported in region {region}") + else: + logger.error(f"Error discovering EMR clusters in {region}: {str(e)}") + return [] + except Exception as e: + logger.error(f"Error discovering EMR clusters in {region}: {str(e)}") + return [] + + def discover_all_clusters(self, regions: List[str]) -> List[Dict[str, Any]]: + """ + Discover all EMR clusters across multiple regions. + + Args: + regions: List of AWS regions + + Returns: + List of EMR cluster dictionaries + """ + all_clusters = [] + + # First try RUNNING and WAITING clusters (typical use case) + active_states = ['STARTING', 'BOOTSTRAPPING', 'RUNNING', 'WAITING'] + for region in regions: + clusters = self.discover_clusters(region, active_states) + all_clusters.extend(clusters) + + # Then try TERMINATED and TERMINATED_WITH_ERRORS clusters (historical) + # Get STOPPED clusters (requires separate call due to EMR API limitations) + for region in regions: + try: + emr_client = self.create_client('emr', region) + logger.debug(f"Looking for STOPPED clusters in {region}") + + paginator = emr_client.get_paginator('list_clusters') + page_iterator = paginator.paginate(ClusterStates=['TERMINATED']) + + for page in page_iterator: + if 'Clusters' in page: + for cluster_summary in page['Clusters']: + # Check if this was a STOPPED cluster (vs actually terminated) + cluster_id = cluster_summary.get('Id') + if cluster_id: + try: + response = emr_client.describe_cluster(ClusterId=cluster_id) + if 'Cluster' in response and 'Status' in response['Cluster']: + status_reason = response['Cluster']['Status'].get('StateChangeReason', {}) + # EMR uses the TERMINATED state for stopped clusters with a specific code + if status_reason.get('Code') == 'USER_REQUEST' and 'stop clusters' in status_reason.get('Message', '').lower(): + cluster = { + 'id': cluster_id, + 'name': cluster_summary.get('Name', 'Unknown'), + 'status': 'STOPPED', # Override with our interpretation + 'region': region, + 'account_id': self.account_id, + 'account_name': self.account_name + } + + # Get tags + if 'Tags' in response['Cluster']: + tags = {} + for tag in response['Cluster']['Tags']: + tags[tag.get('Key')] = tag.get('Value') + cluster['tags'] = tags + else: + cluster['tags'] = {} + + all_clusters.append(cluster) + + except Exception as e: + logger.warning(f"Error checking if cluster {cluster_id} is stopped: {str(e)}") + + except Exception as e: + logger.error(f"Error discovering STOPPED EMR clusters in {region}: {str(e)}") + + logger.info(f"Discovered a total of {len(all_clusters)} EMR clusters across all regions") + return all_clusters + + def validate_credentials(self, region: str = None) -> bool: + """ + Validate if the credentials work for EMR service. + + Args: + region: AWS region to test credentials in + + Returns: + True if credentials are valid, False otherwise + """ + if not region: + # Use a default region based on the detected partition + if any(cred for cred, val in self.credentials.items() if val and 'gov' in val.lower()): + region = 'us-gov-west-1' + else: + region = 'us-east-1' + + try: + logger.info(f"Validating EMR credentials in region {region}") + emr_client = self.create_client('emr', region) + + # Just list a small number of clusters to verify permissions + response = emr_client.list_clusters(MaxResults=1) + logger.info(f"EMR credentials valid in region {region}") + return True + except Exception as e: + logger.warning(f"EMR credentials validation failed in {region}: {str(e)}") + return False + def stop(self, clusters: List[Dict[str, Any]], dry_run: bool = False) -> Dict[str, int]: """ Stop EMR clusters gracefully.