added emr.py

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.