def main()

in cookbooks/aws-parallelcluster-computefleet/files/compute_fleet_status/compute_fleet_status.py [0:0]


def main():
    try:
        parser = argparse.ArgumentParser(description="Get or update compute fleet status of scheduler plugin.")
        parser.add_argument(
            "--table-name",
            type=str,
            required=True,
            help="DynamoDB table name",
        )
        parser.add_argument(
            "--region",
            type=str,
            required=True,
            help="Region of cluster",
        )
        parser.add_argument(
            "--status",
            type=str,
            required=False,
            help="Specify the compute fleet status to set, can be PROTECTED",
            choices={"PROTECTED"},
        )
        parser.add_argument(
            "--action",
            type=str,
            required=True,
            help="Get or update compute-fleet-status",
            choices={"update", "get"},
        )

        args = parser.parse_args()
        if args.action == "update" and not args.status:
            parser.error("ERROR: --status is required when 'action' is specified to 'update'.")
        elif args.action == "get" and args.status:
            parser.error("ERROR: --status can not be specified when 'action' is 'get'.")

        if args.action == "update":
            update_status_with_last_updated_time(args.table_name, args.region, args.status)
        else:
            get_status_with_last_updated_time(args.table_name, args.region)
    except Exception as e:
        print(f"ERROR: Failed to {args.action} compute fleet status, exception: {e}", file=sys.stderr)
        sys.exit(1)