def execute_code()

in parquet_flask/cdms_lambda_func/audit_tool/execute_lambda.py [0:0]


def execute_code(event, context):
    state = None
    s3 = boto3.client('s3')

    if event:
        if isinstance(event, str):
            event = json.loads(event)

        if 'State' in event:
            buf = BytesIO()

            s3.download_fileobj(event['State']['Bucket'], event['State']['Key'], buf)
            buf.seek(0)
            state = json.load(buf)

            print('Loaded persisted state from S3', flush=True)

    if state is None:
        try:
            buf = BytesIO()

            s3.download_fileobj(os.getenv('OPENSEARCH_BUCKET'), 'AUDIT_STATE.json', buf)
            buf.seek(0)
            state = json.load(buf)

            print('Loaded persisted state from S3', flush=True)
        except:
            print('Could not load state, starting from scratch')
            print(traceback.print_exc())
            state = {}

    if 'lastListTime' not in state:
        state['lastListTime'] = datetime(1, 1, 1, tzinfo=timezone.utc)
    else:
        state['lastListTime'] = datetime.strptime(state['lastListTime'], "%Y-%m-%dT%H:%M:%S%z")

    print('INVOKING AUDIT TOOL', flush=True)
    print('', flush=True)

    audit(
        'mock-s3',
        os.getenv('SQS_URL'),
        os.getenv('SNS_ARN'),
        state=state,
        lambda_ctx=context
    )