"""Seach for SFN execution(s).""" import argparse import json from datetime import datetime from enum import Enum from pprint import pprint import boto3 STATUSES = Enum( 'Status', ['RUNNING', 'SUCCEEDED', 'FAILED', 'TIMED_OUT', 'ABORTED'] ) MATCH_TYPES = Enum( 'MatchType', [ 'CONTAINS', 'STARTSWITH', 'ENDSWITH', 'NOT_CONTAINS', 'NOT_STARTSWITH', 'NOT_ENDSWITH' ] ) MATCH_LOCS = Enum( 'MatchLoc', ['NAME', 'INPUT', 'OUTPUT', 'CAUSE', 'ERROR'] ) MATCH_LOC_STATUSES = { MATCH_LOCS.OUTPUT.name: STATUSES.SUCCEEDED.name, MATCH_LOCS.CAUSE.name: STATUSES.FAILED.name, MATCH_LOCS.ERROR.name: STATUSES.FAILED.name } def main(): """Entrypoint.""" args = parse_args() # read metadata about current aws user sts_client = boto3.client('sts') aws_region = sts_client.meta.region_name aws_account_id = sts_client.get_caller_identity().get('Account') # format inputs into variables for calls sfn_arn = f'arn:aws:states:us-east-1:{aws_account_id}:stateMachine:{args.sfn_name}' # noqa:E501 status_filter = args.status match_list = args.match.split('|') if args.match else [] match_type = args.match_type match_loc = args.match_loc start_date_str = args.start_date start_date_fmt = args.start_date_fmt start_date = datetime.strptime( start_date_str, start_date_fmt ) if start_date_str else None end_date_str = args.end_date end_date_fmt = args.end_date_fmt end_date = datetime.strptime( end_date_str, end_date_fmt ) if end_date_str else None max_results = args.max_results # additional validation for MATCH_LOC and STATUS_FILTER pair if status_filter and match_loc in MATCH_LOC_STATUSES: if MATCH_LOC_STATUSES[match_loc] != status_filter: print(f'Invalid match_loc to status_filter pair : {match_loc} => {status_filter}') # noqa:E501 exit(1) # fetch executions in batches sfn_client = boto3.client('stepfunctions') results = { 'arn': sfn_arn, 'executions': [] } next_token = None num_matched = 0 while True: params = { 'stateMachineArn': sfn_arn, 'maxResults': 1000 } if status_filter: params['statusFilter'] = status_filter if next_token: params['nextToken'] = next_token sfn_response = sfn_client.list_executions(**params) # exit on unexpected response if sfn_response['ResponseMetadata']['HTTPStatusCode'] != 200: pprint(sfn_response) exit(1) # process list of execution metadata executions = sfn_response['executions'] for exc in executions: exc_name = exc['name'] exc_arn = exc['executionArn'] exc_start_date = exc['startDate'] exc_stop_date = exc.get('stopDate') status = exc['status'] # skip until earlier than start date if start_date and start_date.replace(tzinfo=exc_start_date.tzinfo) <= exc_start_date: # noqa:E501 continue # end processing if past end date if end_date and end_date.replace(tzinfo=exc_start_date.tzinfo) >= exc_start_date: # noqa:E501 _end(results) # pick what exc data to examine match_str = None if match_loc == MATCH_LOCS.NAME.name: match_str = exc_name else: sfn_exc_response = sfn_client.describe_execution(executionArn=exc_arn) # noqa:E501 if sfn_exc_response['ResponseMetadata']['HTTPStatusCode'] != 200: # noqa:E501 pprint(sfn_exc_response) exit(1) match_str = sfn_exc_response[match_loc.lower()] # check if exc data matches if _matches(match_str, match_list, match_type): results['executions'].append({ 'arn': exc_arn, 'name': exc_name, 'url': _exc_url(aws_region, aws_account_id, exc_arn), 'start': str(exc_start_date), 'stop': str(exc_stop_date), 'status': status }) num_matched += 1 # end processing by short circuit if max_results and num_matched >= max_results: _end(results) # no more results, end processing next_token = sfn_response.get('nextToken') if not next_token: _end(results) def _end(results): print(json.dumps(results, indent=4)) exit(0) def _exc_url(aws_region, aws_account_id, exc_arn): return f'https://{aws_region}.console.aws.amazon.com/states/home?region={aws_region}#/v2/executions/details/{exc_arn}' # noqa:E501 def _matches(match_str, match_list, match_type): if not match_list or not match_type: return True for match_target in match_list: if match_type == MATCH_TYPES.CONTAINS.name and match_target in match_str: # noqa:E501 return True elif match_type == MATCH_TYPES.STARTSWITH.name and match_str.startswith(match_target): # noqa:E501 return True elif match_type == MATCH_TYPES.ENDSWITH.name and match_str.endswith(match_target): # noqa:E501 return True elif match_type == MATCH_TYPES.NOT_CONTAINS.name and match_target not in match_str: # noqa:E501 return True elif match_type == MATCH_TYPES.NOT_STARTSWITH.name and not match_str.startswith(match_target): # noqa:E501 return True elif match_type == MATCH_TYPES.NOT_ENDSWITH.name and not match_str.endswith(match_target): # noqa:E501 return True return False def parse_args(): """Read CLI args.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) # required positional arguments parser.add_argument('sfn_name', help='Name of SFN to search executions') # optional flag arguments parser.add_argument( '--status', choices=[x.name for x in STATUSES], help='Filter in executions of status type' ) parser.add_argument( '--start-date', type=str, help='Ignore until this date + time' ) parser.add_argument( '--start-date-fmt', type=str, default='%Y-%m-%d', help='Format for --start-date using datetime.strptime' ) parser.add_argument( '--end-date', type=str, help='Search until this date + time' ) parser.add_argument( '--end-date-fmt', type=str, default='%Y-%m-%d', help='Format for --end-date using datetime.strptime' ) parser.add_argument( '--match', type=str, help='Match value, use "|" as "OR" statement' ) parser.add_argument( '--match-type', choices=[x.name for x in MATCH_TYPES], default=[x.name for x in MATCH_TYPES][0], help='Matching strategy' ) parser.add_argument( '--match-loc', choices=[x.name for x in MATCH_LOCS], default=[x.name for x in MATCH_LOCS][0], help='Data to examine for match' ) parser.add_argument( '--max-results', type=int, help='Exit after finding number of matches' ) # load and validate arguments return parser.parse_args() if __name__ == '__main__': main()