def lambda_handler()

in custom_resource/sagemaker_query_drift.py [0:0]


def lambda_handler(event, context):
    if "ProcessingJobName" in event:
        job_name = event["ProcessingJobName"]
    else:
        raise KeyError("ProcessingJobName key not found in event: {}.".format(json.dumps(event)))
    try:
        # Parse the result uri
        status, exit_message, result_bucket, result_path = get_processing_job(job_name)
        logger.info("Processing job: {} has status:{}.".format(job_name, status))
        drift = None
        if status == "Completed":
            try:
                # Attempt to load the violations
                violations = get_s3_results_json(
                    result_bucket, result_path, "constraint_violations.json"
                )
                status = "CompletedWithViolations"
                logger.info("Has violations")
                drift = list(get_baseline_drift(violations))
            except Exception as e:
                print(e)
                logger.info("No violations")
        return {
            "statusCode": 200,
            "results": {
                "ProcessingJobName": job_name,
                "ProcessingJobStatus": status,
                "ExitMessage": exit_message,
                "BaselineDrift": drift,
            },
        }
    except Exception as e:
        message = "Failed to read processing status!"
        print(e)
        return {"statusCode": 500, "error": message}