def execute(self, context: 'Context') -> List[str]: emr_hook = EmrHook(aws_conn_id=self.aws_conn_id) emr = emr_hook.get_conn() job_flow_id = self.job_flow_id or emr_hook.get_cluster_id_by_name( str(self.job_flow_name), self.cluster_states) if not job_flow_id: raise AirflowException( f'No cluster found for name: {self.job_flow_name}') if self.do_xcom_push: context['ti'].xcom_push(key='job_flow_id', value=job_flow_id) self.log.info('Adding steps to %s', job_flow_id) # steps may arrive as a string representing a list # e.g. if we used XCom or a file then: steps="[{ step1 }, { step2 }]" steps = self.steps if isinstance(steps, str): steps = ast.literal_eval(steps) response = emr.add_job_flow_steps(JobFlowId=job_flow_id, Steps=steps) if not response['ResponseMetadata']['HTTPStatusCode'] == 200: raise AirflowException(f'Adding steps failed: {response}') else: self.log.info('Steps %s added to JobFlow', response['StepIds']) return response['StepIds']
def execute(self, context: 'Context') -> int: emr_hook = EmrHook(aws_conn_id=self.aws_conn_id) emr = emr_hook.get_conn() if self.do_xcom_push: context['ti'].xcom_push(key='cluster_id', value=self.cluster_id) EmrClusterLink.persist( context=context, operator=self, region_name=emr_hook.conn_region_name, aws_partition=emr_hook.conn_partition, job_flow_id=self.cluster_id, ) self.log.info('Modifying cluster %s', self.cluster_id) response = emr.modify_cluster( ClusterId=self.cluster_id, StepConcurrencyLevel=self.step_concurrency_level) if response['ResponseMetadata']['HTTPStatusCode'] != 200: raise AirflowException(f'Modify cluster failed: {response}') else: self.log.info('Steps concurrency level %d', response['StepConcurrencyLevel']) return response['StepConcurrencyLevel']
def execute(self, context: Dict[str, Any]) -> int: emr_hook = EmrHook(aws_conn_id=self.aws_conn_id) emr = emr_hook.get_conn() if self.do_xcom_push: context['ti'].xcom_push(key='cluster_id', value=self.cluster_id) self.log.info('Modifying cluster %s', self.cluster_id) response = emr.modify_cluster( ClusterId=self.cluster_id, StepConcurrencyLevel=self.step_concurrency_level ) if response['ResponseMetadata']['HTTPStatusCode'] != 200: raise AirflowException(f'Modify cluster failed: {response}') else: self.log.info('Steps concurrency level %d', response['StepConcurrencyLevel']) return response['StepConcurrencyLevel']
def execute(self, context: 'Context') -> None: emr_hook = EmrHook(aws_conn_id=self.aws_conn_id) emr = emr_hook.get_conn() EmrClusterLink.persist( context=context, operator=self, region_name=emr_hook.conn_region_name, aws_partition=emr_hook.conn_partition, job_flow_id=self.job_flow_id, ) self.log.info('Terminating JobFlow %s', self.job_flow_id) response = emr.terminate_job_flows(JobFlowIds=[self.job_flow_id]) if not response['ResponseMetadata']['HTTPStatusCode'] == 200: raise AirflowException(f'JobFlow termination failed: {response}') else: self.log.info('JobFlow with id %s terminated', self.job_flow_id)
def execute(self, context): emr_hook = EmrHook(aws_conn_id=self.aws_conn_id) emr = emr_hook.get_conn() job_flow_id = self.job_flow_id or emr_hook.get_cluster_id_by_name( self.job_flow_name, self.cluster_states) if not job_flow_id: raise AirflowException( f'No cluster found for name: {self.job_flow_name}') if self.do_xcom_push: context['ti'].xcom_push(key='job_flow_id', value=job_flow_id) self.log.info('Adding steps to %s', job_flow_id) response = emr.add_job_flow_steps(JobFlowId=job_flow_id, Steps=self.steps) if not response['ResponseMetadata']['HTTPStatusCode'] == 200: raise AirflowException('Adding steps failed: %s' % response) else: self.log.info('Steps %s added to JobFlow', response['StepIds']) return response['StepIds']
def test_get_conn_returns_a_boto3_connection(self): hook = EmrHook(aws_conn_id='aws_default', region_name='ap-southeast-2') self.assertIsNotNone(hook.get_conn().list_clusters())