feat(aws-data-processing-mcp-server): Add Support for EMR-S Job and Application (#1881)

* feat(aws-data-processing-mcp-server): Add Support for EMR-S Job and Application

* chore(aws-data-processing-mcp-server): Fix pyright errors

* chore(aws-data-processing-mcp-server): Add more test coverage

* chore(aws-dataprocessing-mcp-server):Add more tests for aws_helper

Author: naikvaib

src/aws-dataprocessing-mcp-server/README.md

@@ -373,6 +373,13 @@ Controls whether the MCP server adds and verifies MCP-managed tags on resources.
 |-----------|-------------|----------------|--------------|
 | manage_aws_emr_ec2_steps | Manage Amazon EMR steps for processing data on EMR clusters | add-steps, cancel-steps, describe-step, list-steps | --allow-write flag for add/cancel operations, appropriate AWS permissions |
 
+### EMR Serverless Handler Tools
+
+| Tool Name | Description | Key Operations | Requirements |
+|-----------|-------------|----------------|--------------|
+| manage_aws_emr_serverless_applications | Manage Amazon EMR Serverless applications with comprehensive lifecycle control | create-application, get-application, update-application, delete-application, list-applications, start-application, stop-application | --allow-write flag for create/update/delete/start/stop operations, appropriate AWS permissions |
+| manage_aws_emr_serverless_job_runs | Manage Amazon EMR Serverless job runs for executing data processing workloads | start-job-run, get-job-run, cancel-job-run, list-job-runs, get-dashboard-for-job-run | --allow-write flag for start/cancel operations, application must exist, appropriate AWS permissions |
+
 ### Athena Query Handler Tools
 
 | Tool Name | Description | Key Operations | Requirements |

src/aws-dataprocessing-mcp-server/awslabs/aws_dataprocessing_mcp_server/handlers/emr/emr_serverless_application_handler.py

@@ -465,3 +465,123 @@ class WaitClusterResponse(CallToolResult):
     cluster_id: str = Field(..., description='ID of the cluster')
     state: str = Field(..., description='Current state of the cluster')
     operation: str = Field(default='wait', description='Operation performed')
+
+
+# Response models for EMR Serverless Operations
+
+
+class CreateApplicationResponse(CallToolResult):
+    """Response model for create EMR Serverless application operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application_id: str = Field(..., description='ID of the created application')
+    name: str = Field(..., description='Name of the created application')
+    arn: str = Field(..., description='ARN of the created application')
+    operation: str = Field(default='create-application', description='Operation performed')
+
+
+class GetApplicationResponse(CallToolResult):
+    """Response model for get EMR Serverless application operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application: Dict[str, Any] = Field(..., description='Application details')
+    operation: str = Field(default='get-application', description='Operation performed')
+
+
+class UpdateApplicationResponse(CallToolResult):
+    """Response model for update EMR Serverless application operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application: Dict[str, Any] = Field(..., description='Updated application details')
+    operation: str = Field(default='update-application', description='Operation performed')
+
+
+class DeleteApplicationResponse(CallToolResult):
+    """Response model for delete EMR Serverless application operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application_id: str = Field(..., description='ID of the deleted application')
+    operation: str = Field(default='delete-application', description='Operation performed')
+
+
+class ListApplicationsResponse(CallToolResult):
+    """Response model for list EMR Serverless applications operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    applications: List[Dict[str, Any]] = Field(..., description='List of applications')
+    count: int = Field(..., description='Number of applications found')
+    next_token: Optional[str] = Field(None, description='Token for pagination')
+    operation: str = Field(default='list-applications', description='Operation performed')
+
+
+class StartApplicationResponse(CallToolResult):
+    """Response model for start EMR Serverless application operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application_id: str = Field(..., description='ID of the started application')
+    operation: str = Field(default='start-application', description='Operation performed')
+
+
+class StopApplicationResponse(CallToolResult):
+    """Response model for stop EMR Serverless application operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application_id: str = Field(..., description='ID of the stopped application')
+    operation: str = Field(default='stop-application', description='Operation performed')
+
+
+class StartJobRunResponse(CallToolResult):
+    """Response model for start EMR Serverless job run operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application_id: str = Field(..., description='ID of the application')
+    job_run_id: str = Field(..., description='ID of the started job run')
+    arn: str = Field(..., description='ARN of the job run')
+    operation: str = Field(default='start-job-run', description='Operation performed')
+
+
+class GetJobRunResponse(CallToolResult):
+    """Response model for get EMR Serverless job run operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    job_run: Dict[str, Any] = Field(..., description='Job run details')
+    operation: str = Field(default='get-job-run', description='Operation performed')
+
+
+class CancelJobRunResponse(CallToolResult):
+    """Response model for cancel EMR Serverless job run operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    application_id: str = Field(..., description='ID of the application')
+    job_run_id: str = Field(..., description='ID of the cancelled job run')
+    operation: str = Field(default='cancel-job-run', description='Operation performed')
+
+
+class ListJobRunsResponse(CallToolResult):
+    """Response model for list EMR Serverless job runs operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    job_runs: List[Dict[str, Any]] = Field(..., description='List of job runs')
+    count: int = Field(..., description='Number of job runs found')
+    next_token: Optional[str] = Field(None, description='Token for pagination')
+    operation: str = Field(default='list-job-runs', description='Operation performed')
+
+
+class GetDashboardForJobRunResponse(CallToolResult):
+    """Response model for get dashboard for EMR Serverless job run operation."""
+
+    isError: bool = Field(default=False, description='Whether the operation resulted in an error')
+    content: List[Content] = Field(..., description='Content of the response')
+    url: str = Field(..., description='Dashboard URL for the job run')
+    operation: str = Field(default='get-dashboard-for-job-run', description='Operation performed')

src/aws-dataprocessing-mcp-server/awslabs/aws_dataprocessing_mcp_server/handlers/emr/emr_serverless_job_run_handler.py

@@ -45,6 +45,12 @@
 from awslabs.aws_dataprocessing_mcp_server.handlers.emr.emr_ec2_steps_handler import (
     EMREc2StepsHandler,
 )
+from awslabs.aws_dataprocessing_mcp_server.handlers.emr.emr_serverless_application_handler import (
+    EMRServerlessApplicationHandler,
+)
+from awslabs.aws_dataprocessing_mcp_server.handlers.emr.emr_serverless_job_run_handler import (
+    EMRServerlessJobRunHandler,
+)
 from awslabs.aws_dataprocessing_mcp_server.handlers.glue.crawler_handler import (
     CrawlerHandler,
 )
@@ -372,6 +378,19 @@ def main():
         allow_write=allow_write,
         allow_sensitive_data_access=allow_sensitive_data_access,
     )
+
+    EMRServerlessApplicationHandler(
+        mcp,
+        allow_write=allow_write,
+        allow_sensitive_data_access=allow_sensitive_data_access,
+    )
+
+    EMRServerlessJobRunHandler(
+        mcp,
+        allow_write=allow_write,
+        allow_sensitive_data_access=allow_sensitive_data_access,
+    )
+
     CommonResourceHandler(mcp, allow_write=allow_write)
 
     # Run server

src/aws-dataprocessing-mcp-server/awslabs/aws_dataprocessing_mcp_server/models/emr_models.py

@@ -20,6 +20,7 @@
     CUSTOM_TAGS_ENV_VAR,
     DEFAULT_RESOURCE_TAGS,
     EMR_CLUSTER_RESOURCE_TYPE,
+    EMR_SERVERLESS_APPLICATION_RESOURCE_TYPE,
     MCP_CREATION_TIME_TAG_KEY,
     MCP_MANAGED_TAG_KEY,
     MCP_MANAGED_TAG_VALUE,
@@ -425,3 +426,64 @@ def verify_athena_data_catalog_managed_by_mcp(
         except Exception as e:
             result['error_message'] = f'Error getting data catalog: {str(e)}'
             return result
+
+    @classmethod
+    def verify_emr_serverless_application_managed_by_mcp(
+        cls,
+        emr_serverless_client: Any,
+        application_id: str,
+        expected_resource_type: str = EMR_SERVERLESS_APPLICATION_RESOURCE_TYPE,
+    ) -> Dict[str, Any]:
+        """Verify if an EMR Serverless application is managed by the MCP server and has the expected resource type.
+
+        This method checks if the EMR Serverless application has the MCP managed tag and the correct resource type tag.
+
+        Args:
+            emr_serverless_client: EMR Serverless boto3 client
+            application_id: ID of the EMR Serverless application to verify
+            expected_resource_type: The expected resource type value (default: EMR_SERVERLESS_APPLICATION_RESOURCE_TYPE)
+
+        Returns:
+            Dictionary with verification result:
+                - is_valid: True if verification passed, False otherwise
+                - error_message: Error message if verification failed, None otherwise
+        """
+        # If custom tags are enabled, skip verification
+        if cls.is_custom_tags_enabled():
+            return {'is_valid': True, 'error_message': None}
+
+        result = {'is_valid': False, 'error_message': None}
+
+        try:
+            response = emr_serverless_client.get_application(applicationId=application_id)
+            tags_dict = response.get('application', {}).get('tags', {})
+
+            # Convert tags dictionary to list format for verification
+            tags_list = [{'Key': key, 'Value': value} for key, value in tags_dict.items()]
+
+            # Check if the resource is managed by MCP
+            if not cls.verify_resource_managed_by_mcp(tags_list):
+                result['error_message'] = (
+                    f'Application {application_id} is not managed by MCP (missing required tags)'
+                )
+                return result
+
+            # Check if the resource has the expected resource type
+            actual_type = tags_dict.get(MCP_RESOURCE_TYPE_TAG_KEY, 'unknown')
+            if (
+                actual_type != expected_resource_type
+                and actual_type != EMR_SERVERLESS_APPLICATION_RESOURCE_TYPE
+            ):
+                result['error_message'] = (
+                    f'Application {application_id} has incorrect type (expected {expected_resource_type}, got {actual_type})'
+                )
+                return result
+
+            # All checks passed
+            result['is_valid'] = True
+            return result
+
+        except ClientError as e:
+            # If we can't get the application information, return error
+            result['error_message'] = f'Error retrieving application {application_id}: {str(e)}'
+            return result

src/aws-dataprocessing-mcp-server/awslabs/aws_dataprocessing_mcp_server/server.py

@@ -31,3 +31,5 @@
 EMR_INSTANCE_FLEET_RESOURCE_TYPE = 'EMRInstanceFleet'
 EMR_INSTANCE_GROUP_RESOURCE_TYPE = 'EMRInstanceGroup'
 EMR_STEPS_RESOURCE_TYPE = 'EMRSteps'
+EMR_SERVERLESS_APPLICATION_RESOURCE_TYPE = 'EMRServerlessApplication'
+EMR_SERVERLESS_JOB_RUN_RESOURCE_TYPE = 'EMRServerlessJobRun'