def test_get_single_waiter_config(self):
single_waiter = {
'description':
'Waiter description',
'operation':
'HeadBucket',
'delay':
5,
'maxAttempts':
20,
'acceptors': [
{
'state': 'success',
'matcher': 'status',
'expected': 200
},
{
'state': 'retry',
'matcher': 'status',
'expected': 404
},
],
}
waiters = {
'version': 2,
'waiters': {
'BucketExists': single_waiter,
}
}
model = WaiterModel(waiters)
config = model.get_waiter('BucketExists')
self.assertEqual(config.operation, 'HeadBucket')
def test_get_waiter_does_not_exist(self):
waiters = {
'version': 2,
'waiters': {}
}
model = WaiterModel(waiters)
with self.assertRaises(ValueError):
model.get_waiter('UnknownWaiter')
def test_waiter_names(self):
waiters = {
'version': 2,
'waiters': {
'BarWaiter': {},
'FooWaiter': {},
}
}
self.assertEqual(
WaiterModel(waiters).waiter_names, ['BarWaiter', 'FooWaiter'])
def test_get_single_waiter_config(self):
single_waiter = {
'description': 'Waiter description',
'operation': 'HeadBucket',
'delay': 5,
'maxAttempts': 20,
'acceptors': [
{'state': 'success', 'matcher': 'status', 'expected': 200},
{'state': 'retry', 'matcher': 'status', 'expected': 404},
],
}
waiters = {
'version': 2,
'waiters': {
'BucketExists': single_waiter,
}
}
model = WaiterModel(waiters)
config = model.get_waiter('BucketExists')
self.assertEqual(config.operation, 'HeadBucket')
def test_add_waiters_no_waiter_names(self):
self.session.get_waiter_model.return_value = WaiterModel({
'version': 2,
# No waiters are specified.
'waiters': {}
})
command_table = {}
add_waiters(command_table, self.session, self.command_object)
# Make sure that no wait command was added since the service object
# has no waiters.
self.assertEqual(command_table, {})
def setUp(self):
self.model = WaiterModel({
'version': 2,
'waiters': {
'Foo': {
'operation': 'foo', 'maxAttempts': 1, 'delay': 1,
'acceptors': [],
}
}
})
self.service_object = mock.Mock()
self.cmd = WaitCommand(self.model, self.service_object)
def get_waiter_model(self, service, api_version=None):
"""
Get the waiter model for the service
"""
service = os.path.join('aws', service)
model_version = self.loader.determine_latest(service, api_version)
# Some wierd formatting required to get the name of the model
# correct. Right now this is returned: YYYY-MM-DD.api
# We need: YYYY-MM-DD
model_version = ''.join(model_version.split('.')[:-1])
waiter_model = model_version + '.waiters'
return WaiterModel(self.loader.load_data(waiter_model))
def setUp(self):
self.waiter_config = {
'version': 2,
'waiters': {
'WaiterName': {
'operation': 'Foo',
'delay': 1,
'maxAttempts': 1,
'acceptors': [],
},
},
}
self.waiter_model = WaiterModel(self.waiter_config)
def setUp(self):
self.service_model = mock.Mock()
self.session = mock.Mock()
self.command_object = mock.Mock()
self.command_object.service_model = self.service_model
# Set up the mock session.
self.session.get_waiter_model.return_value = WaiterModel({
'version': 2,
'waiters': {
'FooExists': {},
}
})
def tgw_attachment_waiter(self, desired_state: AttachmentState,
attachment_id: str) -> None:
"""[summary]
Args:
desired_state (str): desired state of the tgw attachment
attachment_id (str): attachment-id
"""
delay = 10
max_attempts = 15
waiter_name = "TGWAttachmentInPendingAcceptance"
waiter_config = {
"version": 2,
"waiters": {
"TGWAttachmentInPendingAcceptance": {
"operation":
"DescribeTransitGatewayPeeringAttachments",
"delay":
delay,
"maxAttempts":
max_attempts,
"acceptors": [
{
"matcher": "path",
"expected": desired_state.value,
"argument":
"TransitGatewayPeeringAttachments[0].State",
"state": "success",
},
{
"matcher": "path",
"expected": AttachmentState.FAILED.value,
"argument":
"TransitGatewayPeeringAttachments[0].State",
"state": "failure",
},
],
}
},
}
waiter_model = WaiterModel(waiter_config)
custom_waiter = create_waiter_with_client(waiter_name, waiter_model,
self.ec2_client)
try:
custom_waiter.wait(TransitGatewayAttachmentIds=[attachment_id])
except WaiterError as err:
self.logger.error(str(err))
raise
def _load_prefect_waiter(boto_client: "boto3.client", client_str: str,
waiter_name: str):
"""
Load a custom waiter from the ./waiters directory.
"""
try:
# Instantiate waiter from accompanying client json file
with pkg_resources.open_text(waiters,
f"{client_str}.json") as handle:
waiter_model = WaiterModel(json.load(handle))
return create_waiter_with_client(waiter_name, waiter_model,
boto_client)
except Exception as err:
raise ValueError(
f"Unable to load waiter '{waiter_name}' for AWS client '{client_str}'."
) from err
def setUp(self):
self.waiter_config = {
'version': 2,
'waiters': {
'WaiterName': {
'operation': 'Foo',
'delay': 1,
'maxAttempts': 1,
'acceptors': [],
},
},
}
self.waiter_model = WaiterModel(self.waiter_config)
self.service_json_model = {
'metadata': {
'serviceFullName': 'Amazon MyService'
},
'operations': {
'Foo': {
'name': 'Foo',
'input': {
'shape': 'FooInputOutput'
},
'output': {
'shape': 'FooInputOutput'
}
}
},
'shapes': {
'FooInputOutput': {
'type': 'structure',
'members': {
'bar': {
'shape': 'String',
'documentation': 'Documents bar'
}
}
},
'String': {
'type': 'string'
}
}
}
self.service_model = ServiceModel(self.service_json_model, 'myservice')
self.client = mock.Mock()
self.client.meta.service_model = self.service_model
def gen_available_waiter(WAITER_ID, operation, delay, maxAttempts, path):
return WaiterModel({
'version': 2,
'waiters': {
WAITER_ID: {
'operation':
operation,
'delay':
delay,
'maxAttempts':
maxAttempts,
'acceptors': [{
'state': 'success',
'matcher': 'path',
'argument': f"{path} == 'available'",
'expected': True
}]
}
}
})
def handle_waiters(client, client_name, class_name, service_name, service_path,
sidebar_lines):
waiter_config = client._get_waiter_config()
waiter_model = WaiterModel(
waiter_config) if 'waiters' in waiter_config else None
if not waiter_model:
return
waiters_path = f'{service_path}/waiters'
sidebar_lines.append(f' - [Waiters]({waiters_path})')
docs_waiters_path = f'docs/{waiters_path}.md'
waiter_names = waiter_model.waiter_names
example_waiter_name = waiter_names[0]
waiter_list_items = create_waiter_index(docs_waiters_path, client_name,
service_name, example_waiter_name)
for name in waiter_names:
handle_waiter(class_name, client_name, name, service_path,
waiter_list_items, waiter_model, waiters_path)
write_lines(docs_waiters_path, waiter_list_items)
def delete_fargate_profile(
profile_name: str,
cluster_name: str,
) -> None:
_logger.debug(f"Deleting EKS Fargate Profile: {profile_name}")
if describe_fargate_profile(profile_name=profile_name,
cluster_name=cluster_name) is None:
_logger.debug(f"EKS Fargate Profile not found: {profile_name}")
return
eks_client = boto3_client("eks")
eks_client.delete_fargate_profile(
fargateProfileName=profile_name,
clusterName=cluster_name,
)
waiter_model = WaiterModel(WAITER_CONFIG)
waiter = create_waiter_with_client("FargateProfileDeleted", waiter_model,
eks_client)
waiter.wait(fargateProfileName=profile_name, clusterName=cluster_name)
_logger.debug(f"Deleted EKS Fargate Profile: {profile_name}")
def get_class_output(client_name):
method_signatures = []
shapes_in_classes = []
client = boto3.client(client_name)
class_name = type(client).__name__
service_model = client._service_model
waiter_config = client._get_waiter_config()
waiter_model = WaiterModel(
waiter_config) if 'waiters' in waiter_config else None
try:
paginator_model = botocore.session.get_session().get_paginator_model(
client_name)
except botocore.exceptions.UnknownServiceError:
paginator_model = None # meaning it probably doesn't have paginators
for name in service_model.operation_names:
method_signatures.append(
get_method_signature(service_model, name, shapes_in_classes,
class_name))
return get_class_signature(client_name, class_name,
service_model.documentation, method_signatures,
shapes_in_classes, waiter_model,
paginator_model)
def _create_delete_traffic_session_waiter(ec2_client):
delete_session_model = WaiterModel({
"version": 2,
"waiters": {
"TrafficMirrorDeleted": {
"delay":
15,
"operation":
"DescribeTrafficMirrorSessions",
"maxAttempts":
40,
"acceptors": [{
"matcher": "error",
"expected": "InvalidTrafficMirrorSessionId.NotFound",
"state": "success",
}],
}
},
})
delete_session_waiter = create_waiter_with_client(
"TrafficMirrorDeleted", delete_session_model, ec2_client)
return delete_session_waiter
def setUp(self):
self.service_object = mock.Mock()
# Create some waiters.
self.model = WaiterModel({
'version': 2,
'waiters': {
'InstanceRunning': {
'description': 'my waiter description',
'delay': 1,
'maxAttempts': 10,
'operation': 'MyOperation',
},
'BucketExists': {
'description': 'my waiter description',
'operation': 'MyOperation',
'delay': 1,
'maxAttempts': 10,
}
}
})
self.waiter_builder = WaiterStateCommandBuilder(
self.model, self.service_object)
def get_waiter_model(self, service, api_version=None):
"""Get the waiter model for the service."""
return WaiterModel(
self.loader.load_service_model(service,
type_name='waiters-2',
api_version=api_version))
def get_waiter_model(self, service, api_version=None):
"""Get the waiter model for the service."""
with mock.patch('botocore.loaders.Loader.list_available_services',
return_value=[service]):
return WaiterModel(self.loader.load_service_model(
service, type_name='waiters-2', api_version=api_version))
def create_environment(self, auth_credentials: Dict, user: User,
environment: Environment):
org_client = self._get_client("organizations")
# Create an account. Requires organizations:CreateAccount permission
account_request = org_client.create_account(
Email=user.email,
AccountName=uuid4().hex,
IamUserAccessToBilling="ALLOW")
# Configuration for our CreateAccount Waiter.
# A waiter is a boto3 helper which can be configured to poll a given status
# endpoint until it succeeds or fails. boto3 has many built in waiters, but none
# for the organizations service so we're building our own here.
waiter_config = {
"version": 2,
"waiters": {
"AccountCreated": {
"operation":
"DescribeCreateAccountStatus",
"delay":
20,
"maxAttempts":
self.MAX_CREATE_ACCOUNT_ATTEMPTS,
"acceptors": [
{
"matcher": "path",
"expected": "SUCCEEDED",
"argument": "CreateAccountStatus.State",
"state": "success",
},
{
"matcher": "path",
"expected": "IN_PROGRESS",
"argument": "CreateAccountStatus.State",
"state": "retry",
},
{
"matcher": "path",
"expected": "FAILED",
"argument": "CreateAccountStatus.State",
"state": "failure",
},
],
}
},
}
waiter_model = WaiterModel(waiter_config)
account_waiter = create_waiter_with_client("AccountCreated",
waiter_model, org_client)
try:
# Poll until the CreateAccount request either succeeds or fails.
account_waiter.wait(
CreateAccountRequestId=account_request["CreateAccountStatus"]
["Id"])
except WaiterError:
# TODO: Possible failure reasons:
# 'ACCOUNT_LIMIT_EXCEEDED'|'EMAIL_ALREADY_EXISTS'|'INVALID_ADDRESS'|'INVALID_EMAIL'|'CONCURRENT_ACCOUNT_MODIFICATION'|'INTERNAL_FAILURE'
raise EnvironmentCreationException(environment.id,
"Failed to create account.")
# We need to re-fetch this since the Waiter throws away the success response for some reason.
created_account_status = org_client.describe_create_account_status(
CreateAccountRequestId=account_request["CreateAccountStatus"]
["Id"])
account_id = created_account_status["CreateAccountStatus"]["AccountId"]
return account_id
def lambda_handler(event, context):
# Set base parameters for each simulation job. Order from Lambda event, then environment variables.
if 'simulationJobParams' in event:
if 'vpcConfig' in event['simulationJobParams']:
sim_job_params['vpcConfig'] = sim_job_params['vpcConfig']
if 'iamRole' in event['simulationJobParams']:
sim_job_params['iamRole'] = sim_job_params['iamRole']
if 'outputLocation' in event['simulationJobParams']:
sim_job_params['outputLocation'] = sim_job_params['outputLocation']
if 'simulationApplicationArn' in event:
app_arn = event['simulationApplicationArn']
if 'serverIP' in event:
private_ip = event['serverIP']
else:
# Launch Server.
server_app_params = create_application_config(event['server'], True, 'localhost')
server_job_params = deepcopy(sim_job_params)
server_job_params['simulationApplications'].append(server_app_params)
server_job_response = robomaker.create_simulation_job(
iamRole=server_job_params["iamRole"],
maxJobDurationInSeconds=server_job_params["maxJobDurationInSeconds"],
simulationApplications=[server_app_params],
vpcConfig=server_job_params["vpcConfig"],
loggingConfig=server_job_params["loggingConfig"],
outputLocation=server_job_params["outputLocation"]
)
# Wait for server to be available.
waiter_name = 'SimJobCreated'
waiter_model = WaiterModel(waiter_config)
custom_waiter = create_waiter_with_client(waiter_name, waiter_model, robomaker)
custom_waiter.wait(job=server_job_response['arn'])
desc_result = robomaker.describe_simulation_job( job = server_job_response['arn'] )
private_ip = desc_result['networkInterface']['privateIpAddress']
# Launch multiple robot batches.
batch_job_requests = []
client_app_params = {}
client_job_params = {}
for robot in event['robots']:
client_app_params[robot['name']] = create_application_config(robot, False, private_ip)
client_job_params[robot['name']] = deepcopy(sim_job_params)
client_job_params[robot['name']]['simulationApplications'].append(client_app_params[robot['name']])
batch_job_requests.append(client_job_params[robot['name']])
response = robomaker.start_simulation_job_batch(
batchPolicy={
'timeoutInSeconds': DEFAULT_MAX_DURATION,
'maxConcurrency': len(event['robots'])
},
createSimulationJobRequests=batch_job_requests,
tags = {
'launcher': 'multi_robot_fleet'
})
return {
'statusCode': 200,
'body': response['arn']
}
def get_waiter_model(config_file):
waiter_dir = Path(__file__).parent.parent.joinpath("waiters")
with open(waiter_dir.joinpath(config_file)) as f:
config = json.load(f)
return WaiterModel(config)
def delete_fargate_profile(
profile_name: str,
cluster_name: str,
) -> None:
_logger.debug(f"Deleting EKS Fargate Profile: {profile_name}")
if describe_fargate_profile(profile_name=profile_name,
cluster_name=cluster_name) is None:
_logger.debug(f"EKS Fargate Profile not found: {profile_name}")
return
eks_client = boto3_client("eks")
try:
timeout = 0
_logger.debug(f"Checking {cluster_name} cluster current state")
res = eks_client.describe_cluster(name=cluster_name)
eks_status = res.get("cluster").get("status")
_logger.debug(f"{cluster_name} current state: {eks_status}")
while eks_status != "ACTIVE" and timeout < 15:
_logger.debug(
f"EKS is an {eks_status} state. Retrying in 1min - attempt {timeout}"
)
time.sleep(60)
timeout += 1
res = eks_client.describe_cluster(name=cluster_name)
eks_status = res.get("cluster").get("status")
timeout = 0
_logger.debug(f"Checking fargate profile {profile_name} current state")
fargate_profile_res = eks_client.describe_fargate_profile(
clusterName=cluster_name, fargateProfileName=profile_name)
fargate_profile_status = fargate_profile_res.get("fargateProfile").get(
"status")
_logger.debug(
f"{profile_name} current state: {fargate_profile_status}")
while fargate_profile_status != "ACTIVE" and timeout < 15:
_logger.debug(
f"Fargate profile is in {fargate_profile_status} state. Retrying in 1min - attempt {timeout}"
)
time.sleep(60)
timeout += 1
fargate_profile_res = eks_client.describe_fargate_profile(
clusterName=cluster_name, fargateProfileName=profile_name)
fargate_profile_status = fargate_profile_res.get(
"fargateProfile").get("status")
_logger.debug(f"Deleting fargate profile {profile_name}")
eks_client.delete_fargate_profile(
fargateProfileName=profile_name,
clusterName=cluster_name,
)
except eks_client.exceptions.ResourceNotFoundException as err:
_logger.debug(err)
waiter_model = WaiterModel(WAITER_CONFIG)
waiter = create_waiter_with_client("FargateProfileDeleted", waiter_model,
eks_client)
waiter.wait(fargateProfileName=profile_name, clusterName=cluster_name)
_logger.debug(f"Deleted EKS Fargate Profile: {profile_name}")
def __init__(self, client):
waiter_json_filename = os.path.join(utils.__path__[0], 'cfn-waiters-2.json')
with open(waiter_json_filename, 'r') as waiter_json_file:
self.waiter_json_model = json.load(waiter_json_file)
self.waiter_model = WaiterModel(self.waiter_json_model)
self.waiter = create_waiter_with_client('StackAvailable', self.waiter_model, client.meta.client)
def test_unsupported_waiter_version(self):
waiters = {'version': 1, 'waiters': {}}
with self.assertRaises(WaiterConfigError):
WaiterModel(waiters)
def test_waiter_version(self):
self.assertEqual(WaiterModel({'version': 2, 'waiters': {}}).version, 2)
def test_get_waiter_does_not_exist(self):
waiters = {'version': 2, 'waiters': {}}
model = WaiterModel(waiters)
with self.assertRaises(ValueError):
model.get_waiter('UnknownWaiter')
def run(
self,
client: str = None,
waiter_name: str = None,
waiter_definition: dict = None,
waiter_kwargs: dict = None,
credentials: str = None,
):
"""
Task for waiting on a long-running AWS job. Uses the underlying boto3 waiter functionality.
Args:
- client (str): The AWS client on which to wait (e.g., 'batch', 'ec2', etc)
- waiter_name (str, optional): The name of the waiter to instantiate. Can be a boto-supported
waiter or one of prefect's custom waiters. Currently, prefect offers three additional
waiters for AWS Batch: `"JobExists"` waits for a job to be instantiated, `"JobRunning"`
waits for a job to start running, and `"JobComplete"` waits for a job to finish. You can
find the definitions for all prefect-defined waiters [here](https://github.com/PrefectHQ/prefect/tree/master/src/prefect/tasks/aws/waiters). # noqa
You may also use a custom waiter name, if you supply an accompanying waiter definition
dict.
- waiter_definition (dict, optional): A valid custom waiter model, as a dict. Note that if
you supply a custom definition, it is assumed that the provided 'waiter_name' is
contained within the waiter definition dict.
- waiter_kwargs (dict, optional): Arguments to pass to the `waiter.wait(...)` method. Will
depend upon the specific waiter being called.
- credentials (dict, optional): your AWS credentials passed from an upstream
Secret task; this Secret must be a JSON string
with two keys: `ACCESS_KEY` and `SECRET_ACCESS_KEY` which will be
passed directly to `boto3`. If not provided here or in context, `boto3`
will fall back on standard AWS rules for authentication.
"""
if not client:
raise ValueError("An AWS client string must be provided.")
if not waiter_name:
raise ValueError("A waiter name must be provided.")
if not waiter_kwargs:
waiter_kwargs = {}
boto_client = get_boto_client(client,
credentials=credentials,
**self.boto_kwargs)
if waiter_definition:
# Use user-provided waiter definition
waiter_model = WaiterModel(waiter_definition)
waiter = create_waiter_with_client(waiter_name, waiter_model,
boto_client)
else:
# Use either boto-provided or prefect-provided waiter
if waiter_name in boto_client.waiter_names:
waiter = boto_client.get_waiter(waiter_name)
else:
waiter = self._load_prefect_waiter(boto_client, client,
waiter_name)
try:
waiter.wait(**waiter_kwargs)
except WaiterError as e:
raise FAIL(
f"AWS {client} waiter '{waiter_name}' failed with: {str(e)}"
) from e
