def update_poll_status()

in infra/src/lambda_poller/lambda_function.py [0:0]

• 29 lines of code
• 3 McCabe index (conditional complexity)

def update_poll_status(job_id, continuation_token):
    """
Updates the polling status for the given job
    :param job_id:
    :param continuation_token:
    """
    type_processor_map = {
        "frauddetector": FrauddetectorPoller()
    }

    job_parameters = json.loads(continuation_token)
    training_type = job_parameters['training_type']

    # Validate job type
    assert training_type in list(type_processor_map.keys()), "The training type must be in {} ".format(
        list(type_processor_map.keys()))

    # Update polling status in code pipeline
    code_pipeline = boto3.client('codepipeline')

    try:
        # Get poller
        poller = type_processor_map[training_type]

        # Construct boto3 client
        role = job_parameters['assume_role']
        access_key, secret_key, session_token = get_credentials_for_role(role)
        client = boto3.client(
            training_type,
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            aws_session_token=session_token,
        )

        # Check status
        is_complete = poller.poll_training_status(job_parameters, client)
        if is_complete:
            code_pipeline.put_job_success_result(jobId=job_id)
        else:
            code_pipeline.put_job_success_result(jobId=job_id, continuationToken=continuation_token)

    except Exception as e:
        code_pipeline.put_job_failure_result(jobId=job_id, failureDetails={
            'type': 'JobFailed',
            'message': f'Training job:  failed. Reason: {e} for user params {continuation_token}'
        })