diff --git a/aws-mwaa-environment/aws-mwaa-environment.json b/aws-mwaa-environment/aws-mwaa-environment.json index 47afdd6..00d0039 100644 --- a/aws-mwaa-environment/aws-mwaa-environment.json +++ b/aws-mwaa-environment/aws-mwaa-environment.json @@ -3,24 +3,6 @@ "description": "Resource schema for AWS::MWAA::Environment", "sourceUrl": "https://github.com/aws-cloudformation/aws-cloudformation-resource-providers-mwaa.git", "definitions": { - "TagKey": { - "type": "string", - "description": "", - "minLength": 1, - "maxLength": 128, - "pattern": "^([\\p{L}\\p{Z}\\p{N}_.:/=+\\-@]*)$" - }, - "TagValue": { - "type": "string", - "description": "", - "minLength": 1, - "maxLength": 256, - "pattern": "^([\\p{L}\\p{Z}\\p{N}_.:/=+\\-@]*)$" - }, - "TagMap": { - "type": "object", - "description": "A map of tags for the environment." - }, "EnvironmentName": { "type": "string", "description": "Customer-defined identifier for the environment, unique per customer region.", @@ -69,21 +51,21 @@ "description": "", "minLength": 1, "maxLength": 1224, - "pattern": "^arn:aws(-[a-z]+)?:airflow:[a-z0-9\\-]+:\\d{12}:environment/\\w+" + "pattern": "^arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b)(-[a-z]+)?:airflow:[a-z0-9\\-]+:\\d{12}:environment/\\w+" }, "EnvironmentArn": { "type": "string", "description": "ARN for the MWAA environment.", "minLength": 1, "maxLength": 1224, - "pattern": "^arn:aws(-[a-z]+)?:airflow:[a-z0-9\\-]+:\\d{12}:environment/\\w+" + "pattern": "^arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b)(-[a-z]+)?:airflow:[a-z0-9\\-]+:\\d{12}:environment/\\w+" }, "S3BucketArn": { "type": "string", "description": "ARN for the AWS S3 bucket to use as the source of DAGs and plugins for the environment.", "minLength": 1, "maxLength": 1224, - "pattern": "^arn:aws(-[a-z]+)?:s3:::[a-z0-9.\\-]+$" + "pattern": "^arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b)(-[a-z]+)?:s3:::[a-z0-9.\\-]+$" }, "CreatedAt": { "type": "string", @@ -104,19 +86,19 @@ "type": "string", "description": "IAM role to be used by tasks.", "maxLength": 1224, - "pattern": "^arn:aws(-[a-z]+)?:iam::\\d{12}:role/?[a-zA-Z_0-9+=,.@\\-_/]+$" + "pattern": "^arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b)(-[a-z]+)?:iam::\\d{12}:role/?[a-zA-Z_0-9+=,.@\\-_/]+$" }, "ServiceRoleArn": { "type": "string", "description": "IAM role to be used by MWAA to perform AWS API calls on behalf of the customer.", "maxLength": 1224, - "pattern": "^arn:aws(-[a-z]+)?:iam::\\d{12}:role/?[a-zA-Z_0-9+=,.@\\-_/]+$" + "pattern": "^arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b)(-[a-z]+)?:iam::\\d{12}:role/?[a-zA-Z_0-9+=,.@\\-_/]+$" }, "KmsKey": { "type": "string", "description": "The identifier of the AWS Key Management Service (AWS KMS) customer master key (CMK) to use for MWAA data encryption.\n\n You can specify the CMK using any of the following:\n\n Key ID. For example, key/1234abcd-12ab-34cd-56ef-1234567890ab.\n\n Key alias. For example, alias/ExampleAlias.\n\n Key ARN. For example, arn:aws:kms:us-east-1:012345678910:key/abcd1234-a123-456a-a12b-a123b4cd56ef.\n\n Alias ARN. For example, arn:aws:kms:us-east-1:012345678910:alias/ExampleAlias.\n\n AWS authenticates the CMK asynchronously. Therefore, if you specify an ID, alias, or ARN that is not valid, the action can appear to complete, but eventually fails.", "maxLength": 1224, - "pattern": "^(((arn:aws(-[a-z]+)?:kms:[a-z]{2}-[a-z]+-\\d:\\d+:)?key\\/)?[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}|(arn:aws:kms:[a-z]{2}-[a-z]+-\\d:\\d+:)?alias/.+)$" + "pattern": "^(((arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b)(-[a-z]+)?:kms:[a-z]{2}-[a-z]+-\\d:\\d+:)?key\\/)?[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}|(arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b):kms:[a-z]{2}-[a-z]+-\\d:\\d+:)?alias/.+)$" }, "AirflowVersion": { "type": "string", @@ -159,7 +141,7 @@ "type": "string", "description": "", "maxLength": 1224, - "pattern": "^arn:aws(-[a-z]+)?:logs:[a-z0-9\\-]+:\\d{12}:log-group:\\w+" + "pattern": "^arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b)(-[a-z]+)?:logs:[a-z0-9\\-]+:\\d{12}:log-group:\\w+" }, "LoggingEnabled": { "type": "boolean", @@ -192,6 +174,11 @@ "description": "Minimum worker compute units.", "minimum": 1 }, + "Schedulers": { + "type": "integer", + "description": "Scheduler compute units.", + "minimum": 1 + }, "NetworkConfiguration": { "type": "object", "description": "Configures the network resources of the environment.", @@ -199,6 +186,7 @@ "properties": { "SubnetIds": { "type": "array", + "insertionOrder": true, "description": "A list of subnets to use for the environment. These must be private subnets, in the same VPC, in two different availability zones.", "minItems": 2, "maxItems": 2, @@ -208,6 +196,7 @@ }, "SecurityGroupIds": { "type": "array", + "insertionOrder": true, "description": "A list of security groups to use for the environment.", "minItems": 1, "maxItems": 5, @@ -368,15 +357,18 @@ "PluginsS3ObjectVersion": { "$ref": "#/definitions/S3ObjectVersion" }, - "MinWorkers": { - "$ref": "#/definitions/MinWorkers" - }, "RequirementsS3Path": { "$ref": "#/definitions/RelativePath" }, "RequirementsS3ObjectVersion": { "$ref": "#/definitions/S3ObjectVersion" }, + "StartupScriptS3Path": { + "$ref": "#/definitions/RelativePath" + }, + "StartupScriptS3ObjectVersion": { + "$ref": "#/definitions/S3ObjectVersion" + }, "AirflowConfigurationOptions": { "type": "object", "description": "Key/value pairs representing Airflow configuration variables.\n Keys are prefixed by their section:\n\n [core]\n dags_folder={AIRFLOW_HOME}/dags\n\n Would be represented as\n\n \"core.dags_folder\": \"{AIRFLOW_HOME}/dags\"" @@ -387,6 +379,12 @@ "MaxWorkers": { "$ref": "#/definitions/MaxWorkers" }, + "MinWorkers": { + "$ref": "#/definitions/MinWorkers" + }, + "Schedulers": { + "$ref": "#/definitions/Schedulers" + }, "NetworkConfiguration": { "$ref": "#/definitions/NetworkConfiguration" }, @@ -397,7 +395,8 @@ "$ref": "#/definitions/WeeklyMaintenanceWindowStart" }, "Tags": { - "$ref": "#/definitions/TagMap" + "type": "object", + "description": "A map of tags for the environment." }, "WebserverAccessMode": { "$ref": "#/definitions/WebserverAccessMode" @@ -421,6 +420,7 @@ "/properties/LoggingConfiguration/WorkerLogs/CloudWatchLogGroupArn", "/properties/LoggingConfiguration/TaskLogs/CloudWatchLogGroupArn" ], + "taggable": true, "primaryIdentifier": [ "/properties/Name" ], @@ -438,7 +438,9 @@ }, "update": { "permissions": [ - "airflow:UpdateEnvironment" + "airflow:UpdateEnvironment", + "airflow:TagResource", + "airflow:UntagResource" ], "timeoutInMinutes": 480 }, @@ -453,4 +455,4 @@ ] } } -} \ No newline at end of file +} diff --git a/aws-mwaa-environment/resource-role.yaml b/aws-mwaa-environment/resource-role.yaml index 64dce2e..ae8693d 100644 --- a/aws-mwaa-environment/resource-role.yaml +++ b/aws-mwaa-environment/resource-role.yaml @@ -15,6 +15,13 @@ Resources: Principal: Service: resources.cloudformation.amazonaws.com Action: sts:AssumeRole + Condition: + StringEquals: + aws:SourceAccount: + Ref: AWS::AccountId + StringLike: + aws:SourceArn: + Fn::Sub: arn:${AWS::Partition}:cloudformation:${AWS::Region}:${AWS::AccountId}:type/resource/AWS-MWAA-Environment/* Path: "/" Policies: - PolicyName: ResourceTypePolicy @@ -27,6 +34,8 @@ Resources: - "airflow:DeleteEnvironment" - "airflow:GetEnvironment" - "airflow:ListEnvironments" + - "airflow:TagResource" + - "airflow:UntagResource" - "airflow:UpdateEnvironment" Resource: "*" Outputs: diff --git a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/BaseHandlerStd.java b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/BaseHandlerStd.java index 02b165d..e0db79d 100644 --- a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/BaseHandlerStd.java +++ b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/BaseHandlerStd.java @@ -104,7 +104,6 @@ protected Optional getEnvironmentStatus( } } - protected Environment getEnvironment(final ProxyClient mwaaClientProxy, final GetEnvironmentRequest awsRequest) { final GetEnvironmentResponse response = doReadEnvironment(awsRequest, mwaaClientProxy); @@ -182,4 +181,8 @@ protected GetEnvironmentResponse doReadEnvironment( throw new CfnNotFoundException(ResourceModel.TYPE_NAME, request.name(), e); } } + + protected Logger getLogger() { + return this.logger; + } } diff --git a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/CreateEnvironmentRetryListener.java b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/CreateEnvironmentRetryListener.java new file mode 100644 index 0000000..d8c018c --- /dev/null +++ b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/CreateEnvironmentRetryListener.java @@ -0,0 +1,56 @@ +// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved. + +package software.amazon.mwaa.environment; + +import com.github.rholder.retry.Attempt; +import com.github.rholder.retry.RetryListener; +import software.amazon.cloudformation.proxy.Logger; + +/** + * Listener for CreateEnvironment retries. + */ +public class CreateEnvironmentRetryListener implements RetryListener { + private long attemptNumber = 0; + private long delaySinceFirstAttempt = 0; + + private final Logger logger; + private final int maxRetries; + private final String environmentName; + + /** + * + * @param logger logger for retry logging. + * @param maxRetries maximum retry attempts before failure. + * @param environmentName name of environment to be created. + */ + public CreateEnvironmentRetryListener(Logger logger, int maxRetries, String environmentName) { + this.logger = logger; + this.maxRetries = maxRetries; + this.environmentName = environmentName; + } + + @Override + public void onRetry(Attempt attempt) { + attemptNumber = attempt.getAttemptNumber(); + delaySinceFirstAttempt = attempt.getDelaySinceFirstAttempt(); + + if (attempt.hasResult()) { + log("CreateEnvironment [%s]: retry attempt %d/%d successful. Total delay since first attempt: %dms", + environmentName, attemptNumber, maxRetries, delaySinceFirstAttempt); + } else { + log("CreateEnvironment [%s]: retry attempt %d/%d failed with error message: %s. " + + "Total delay since first attempt: %dms", + environmentName, + attemptNumber, + maxRetries, + attempt.getExceptionCause().getMessage(), + delaySinceFirstAttempt); + } + } + + private void log(final String format, final Object... args) { + if (logger != null) { + logger.log(String.format(format, args)); + } + } +} diff --git a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/CreateHandler.java b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/CreateHandler.java index 48f8ec5..0126463 100644 --- a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/CreateHandler.java +++ b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/CreateHandler.java @@ -2,12 +2,21 @@ package software.amazon.mwaa.environment; +import com.github.rholder.retry.Attempt; +import com.github.rholder.retry.RetryException; +import com.github.rholder.retry.RetryListener; +import com.github.rholder.retry.Retryer; +import com.github.rholder.retry.RetryerBuilder; +import com.github.rholder.retry.StopStrategies; +import com.github.rholder.retry.WaitStrategies; import java.time.Duration; import java.util.Optional; +import java.util.concurrent.ExecutionException; import software.amazon.awssdk.services.mwaa.MwaaClient; import software.amazon.awssdk.services.mwaa.model.CreateEnvironmentRequest; import software.amazon.awssdk.services.mwaa.model.CreateEnvironmentResponse; import software.amazon.awssdk.services.mwaa.model.EnvironmentStatus; +import software.amazon.awssdk.services.mwaa.model.InternalServerException; import software.amazon.awssdk.services.mwaa.model.ValidationException; import software.amazon.cloudformation.exceptions.CfnInvalidRequestException; import software.amazon.cloudformation.proxy.HandlerErrorCode; @@ -24,6 +33,7 @@ */ public class CreateHandler extends BaseHandlerStd { private static final Duration CALLBACK_DELAY = Duration.ofMinutes(1); + public static final int MAX_RETRIES = 14; protected ProgressEvent handleRequest( final Proxies proxies, @@ -91,14 +101,36 @@ private CreateEnvironmentResponse doCreateEnvironment( try { log("Creating %s [%s]", ResourceModel.TYPE_NAME, name); - final CreateEnvironmentResponse response = mwaaClientProxy.injectCredentialsAndInvokeV2( + + final CreateEnvironmentRetryListener listener = new CreateEnvironmentRetryListener( + getLogger(), MAX_RETRIES, name); + final Retryer retryer = getCreateEnvironmentRetryer(listener); + + final CreateEnvironmentResponse response = retryer.call(() -> mwaaClientProxy.injectCredentialsAndInvokeV2( awsRequest, - mwaaClientProxy.client()::createEnvironment); + mwaaClientProxy.client()::createEnvironment)); log("Create submitted %s [%s]", ResourceModel.TYPE_NAME, name); callbackContext.setStabilizing(true); return response; - } catch (final ValidationException e) { - throw new CfnInvalidRequestException(e.getMessage(), e); + } catch (final RetryException e) { + final Attempt lastAttempt = e.getLastFailedAttempt(); + Throwable rootCause = lastAttempt.getExceptionCause(); + log("CreateEnvironment [%s]: Reached maximum number of retires. Total delay since first attempt: %dms", + name, + lastAttempt.getDelaySinceFirstAttempt()); + throw new CfnInvalidRequestException(rootCause.getMessage(), e); + } catch (ExecutionException e) { + throw new CfnInvalidRequestException(e.getCause().getMessage(), e); } } + + private Retryer getCreateEnvironmentRetryer(RetryListener listener) { + return RetryerBuilder.newBuilder() + .retryIfExceptionOfType(ValidationException.class) + .retryIfExceptionOfType(InternalServerException.class) + .withRetryListener(listener) + .withWaitStrategy(WaitStrategies.exponentialWait()) + .withStopStrategy(StopStrategies.stopAfterAttempt(MAX_RETRIES)) + .build(); + } } diff --git a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/UpdateHandler.java b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/UpdateHandler.java index d3b814a..1da7fee 100644 --- a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/UpdateHandler.java +++ b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/environment/UpdateHandler.java @@ -65,15 +65,13 @@ protected ProgressEvent handleRequest( HandlerErrorCode.NotStabilized, "Update failed, resource no longer exists"); } - if (status.get() == EnvironmentStatus.AVAILABLE) { log("status is AVAILABLE, returning success"); return ProgressEvent.progress(model, callbackContext).then( progress -> getEnvironmentDetails("Update::PostUpdateRead", proxies, progress)); } - if (status.get() == EnvironmentStatus.UPDATE_FAILED) { - log("status is UPDATE_FAILED, returning failure"); + log("status is UPDATE_FAILED, returning failure"); return ProgressEvent.failed( model, null, @@ -88,8 +86,7 @@ protected ProgressEvent handleRequest( HandlerErrorCode.NotStabilized, String.format("Update failed, Environment unavailable. %s", errorMessage)); } - - log("status is %s, requesting a callback in %s", status, CALLBACK_DELAY); + log("status is {}, requesting a callback in {}", status, CALLBACK_DELAY); return ProgressEvent.builder() .resourceModel(model) .callbackContext(callbackContext) diff --git a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/CreateTranslator.java b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/CreateTranslator.java index a4c69d7..87aa9e0 100644 --- a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/CreateTranslator.java +++ b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/CreateTranslator.java @@ -39,11 +39,15 @@ public static CreateEnvironmentRequest translateToCreateRequest(final ResourceMo .requirementsS3Path(model.getRequirementsS3Path()) .requirementsS3ObjectVersion( model.getRequirementsS3ObjectVersion()) + .startupScriptS3Path(model.getStartupScriptS3Path()) + .startupScriptS3ObjectVersion( + model.getStartupScriptS3ObjectVersion()) .airflowConfigurationOptions(toStringToStringMap( model.getAirflowConfigurationOptions())) .environmentClass(model.getEnvironmentClass()) .maxWorkers(model.getMaxWorkers()) .minWorkers(model.getMinWorkers()) + .schedulers(model.getSchedulers()) .networkConfiguration(toApiNetworkConfiguration( model.getNetworkConfiguration())) .loggingConfiguration(toApiLoggingConfiguration( diff --git a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/ReadTranslator.java b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/ReadTranslator.java index fe4c255..34bc504 100644 --- a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/ReadTranslator.java +++ b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/ReadTranslator.java @@ -64,10 +64,13 @@ public static ResourceModel translateFromReadResponse(final GetEnvironmentRespon .pluginsS3ObjectVersion(env.pluginsS3ObjectVersion()) .requirementsS3Path(env.requirementsS3Path()) .requirementsS3ObjectVersion(env.requirementsS3ObjectVersion()) + .startupScriptS3Path(env.startupScriptS3Path()) + .startupScriptS3ObjectVersion(env.startupScriptS3ObjectVersion()) .airflowConfigurationOptions(toStringToObjectMap(env.airflowConfigurationOptions())) .environmentClass(env.environmentClass()) .maxWorkers(env.maxWorkers()) .minWorkers(env.minWorkers()) + .schedulers(env.schedulers()) .networkConfiguration(toCfnNetworkConfiguration(env.networkConfiguration())) .loggingConfiguration(toCfnLoggingConfiguration(env.loggingConfiguration())) .weeklyMaintenanceWindowStart(env.weeklyMaintenanceWindowStart()) diff --git a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/UpdateTranslator.java b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/UpdateTranslator.java index 4d8162f..260db29 100644 --- a/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/UpdateTranslator.java +++ b/aws-mwaa-environment/src/main/java/software/amazon/mwaa/translator/UpdateTranslator.java @@ -28,19 +28,23 @@ private UpdateTranslator() { public static UpdateEnvironmentRequest translateToUpdateRequest(final ResourceModel model) { return UpdateEnvironmentRequest.builder() .name(model.getName()) + .networkConfiguration(toApiUpdateNetworkConfiguration(model.getNetworkConfiguration())) .executionRoleArn(model.getExecutionRoleArn()) .airflowVersion(model.getAirflowVersion()) - .networkConfiguration(toApiUpdateNetworkConfiguration(model.getNetworkConfiguration())) .sourceBucketArn(model.getSourceBucketArn()) .dagS3Path(model.getDagS3Path()) .pluginsS3Path(model.getPluginsS3Path()) .pluginsS3ObjectVersion(model.getPluginsS3ObjectVersion()) .requirementsS3Path(model.getRequirementsS3Path()) .requirementsS3ObjectVersion(model.getRequirementsS3ObjectVersion()) + .startupScriptS3Path(model.getStartupScriptS3Path()) + .startupScriptS3ObjectVersion( + model.getStartupScriptS3ObjectVersion()) .airflowConfigurationOptions(toStringToStringMap(model.getAirflowConfigurationOptions())) .environmentClass(model.getEnvironmentClass()) .maxWorkers(model.getMaxWorkers()) .minWorkers(model.getMinWorkers()) + .schedulers(model.getSchedulers()) .loggingConfiguration(toApiLoggingConfiguration(model.getLoggingConfiguration())) .weeklyMaintenanceWindowStart(model.getWeeklyMaintenanceWindowStart()) .webserverAccessMode(model.getWebserverAccessMode()) diff --git a/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/CreateEnvironmentRetryListenerTest.java b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/CreateEnvironmentRetryListenerTest.java new file mode 100644 index 0000000..486793f --- /dev/null +++ b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/CreateEnvironmentRetryListenerTest.java @@ -0,0 +1,64 @@ +// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved. + +package software.amazon.mwaa.environment; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.github.rholder.retry.Attempt; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; +import software.amazon.awssdk.services.mwaa.model.CreateEnvironmentResponse; +import software.amazon.awssdk.services.mwaa.model.ValidationException; +import software.amazon.cloudformation.proxy.Logger; + +/** + * Tests for {@link CreateEnvironmentRetryListener}. + */ +public class CreateEnvironmentRetryListenerTest { + + private static final String SUCCESS_KEY_STRING = "successful"; + private static final String FAILURE_KEY_STRING = "failed"; + /** + * Test for non null logger on retry success. + */ + @Test + public void logNonNullRetrySuccess() { + // given + final int maxRetries = 8; + final String environmentName = "ENVIRONMENT_NAME"; + + final Logger logger = mock(Logger.class); + final CreateEnvironmentRetryListener listener = new CreateEnvironmentRetryListener( + logger, maxRetries, environmentName); + Attempt attempt = mock(Attempt.class); + // when + when(attempt.hasResult()).thenReturn(true); + // then + listener.onRetry(attempt); + verify(logger, times(1)).log(Mockito.argThat(s -> s.contains(SUCCESS_KEY_STRING))); + } + + /** + * Test for non null logger on retry attempt failure. + */ + @Test + public void logNonNullRetryFailure() { + // given + final int maxRetries = 8; + final String environmentName = "ENVIRONMENT_NAME"; + + final Logger logger = mock(Logger.class); + final CreateEnvironmentRetryListener listener = new CreateEnvironmentRetryListener( + logger, maxRetries, environmentName); + Attempt attempt = mock(Attempt.class); + // when + when(attempt.hasResult()).thenReturn(false); + when(attempt.getExceptionCause()).thenReturn(ValidationException.builder().build()); + // then + listener.onRetry(attempt); + verify(logger, times(1)).log(Mockito.argThat(s -> s.contains(FAILURE_KEY_STRING))); + } +} diff --git a/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/CreateHandlerTest.java b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/CreateHandlerTest.java index d7ba236..7ec206c 100644 --- a/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/CreateHandlerTest.java +++ b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/CreateHandlerTest.java @@ -3,9 +3,11 @@ package software.amazon.mwaa.environment; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @@ -14,13 +16,18 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.junit.jupiter.MockitoExtension; +import software.amazon.awssdk.services.mwaa.MwaaClient; +import software.amazon.awssdk.services.mwaa.model.AccessDeniedException; import software.amazon.awssdk.services.mwaa.model.CreateEnvironmentRequest; import software.amazon.awssdk.services.mwaa.model.CreateEnvironmentResponse; import software.amazon.awssdk.services.mwaa.model.Environment; import software.amazon.awssdk.services.mwaa.model.EnvironmentStatus; import software.amazon.awssdk.services.mwaa.model.GetEnvironmentRequest; import software.amazon.awssdk.services.mwaa.model.GetEnvironmentResponse; +import software.amazon.awssdk.services.mwaa.model.InternalServerException; import software.amazon.awssdk.services.mwaa.model.ResourceNotFoundException; import software.amazon.awssdk.services.mwaa.model.ValidationException; import software.amazon.cloudformation.exceptions.CfnAlreadyExistsException; @@ -28,7 +35,10 @@ import software.amazon.cloudformation.proxy.HandlerErrorCode; import software.amazon.cloudformation.proxy.OperationStatus; import software.amazon.cloudformation.proxy.ProgressEvent; +import software.amazon.cloudformation.proxy.ProxyClient; import software.amazon.cloudformation.proxy.ResourceHandlerRequest; +import software.amazon.mwaa.translator.CreateTranslator; +import software.amazon.mwaa.translator.ReadTranslator; /** * Tests for {@link CreateHandler}. @@ -194,6 +204,86 @@ public void handleRequestAlreadyExists() { } } + /** + * Asserts throwing {@link CfnInvalidRequestException} when request experiences transient error on retry. + */ + @Test + public void handleRequestInvalidInputNonRetryableException() { + // given + final CreateHandler handler = new CreateHandler(); + final ResourceModel model = ResourceModel.builder().name("NAME").kmsKey(INVALID_DATA).build(); + final ResourceHandlerRequest request = ResourceHandlerRequest.builder() + .desiredResourceState(model) + .build(); + final GetEnvironmentRequest awsGetEnvironmentRequest = ReadTranslator.translateToReadRequest(model); + final CreateEnvironmentRequest awsCreateEnvironmentRequest = CreateTranslator.translateToCreateRequest(model); + ProxyClient mwaaClientProxy = getProxies().getMwaaClientProxy(); + + // when + when(mwaaClientProxy.injectCredentialsAndInvokeV2(awsGetEnvironmentRequest, + mwaaClientProxy.client()::getEnvironment)) + .thenThrow(ResourceNotFoundException.class); + + when(mwaaClientProxy.injectCredentialsAndInvokeV2(awsCreateEnvironmentRequest, + mwaaClientProxy.client()::createEnvironment)) + .thenThrow(AccessDeniedException.builder().message(INVALID_DATA).build()); + + // then + assertThatThrownBy(() -> handler.handleRequest( + getProxies(), + request, + new CallbackContext()) + ).isInstanceOf(CfnInvalidRequestException.class); + + verify(getSdkClient(), times(1)).createEnvironment( + any(CreateEnvironmentRequest.class)); + } + + /** + * Asserts that environment is successfully created after recovering from both + * {@link ValidationException} and {@link InternalServerException}. + * @param mwaaCreateEnvironmentExceptionClass class of Exception to be thrown by + * mwaa client upon CreateEnvironment request. + */ + @ParameterizedTest + @ValueSource(classes = {ValidationException.class, InternalServerException.class}) + public void handleRequestInvalidInputRecovery(Class mwaaCreateEnvironmentExceptionClass) { + // given + final CreateHandler handler = new CreateHandler(); + final ResourceModel model = ResourceModel.builder().name("NAME").kmsKey(INVALID_DATA).build(); + final ResourceHandlerRequest request = ResourceHandlerRequest.builder() + .desiredResourceState(model) + .build(); + final GetEnvironmentRequest awsGetEnvironmentRequest = ReadTranslator.translateToReadRequest(model); + final CreateEnvironmentRequest awsCreateEnvironmentRequest = CreateTranslator.translateToCreateRequest(model); + final CreateEnvironmentResponse createEnvironmentResponse = CreateEnvironmentResponse.builder().build(); + + ProxyClient mwaaClientProxy = getProxies().getMwaaClientProxy(); + final int expectedNumberOfInvocations = 5; + + // when + when(mwaaClientProxy.injectCredentialsAndInvokeV2(awsGetEnvironmentRequest, + mwaaClientProxy.client()::getEnvironment)) + // to indicate this environment does not exists and allow creation + .thenThrow(ResourceNotFoundException.class); + + when(mwaaClientProxy.injectCredentialsAndInvokeV2(awsCreateEnvironmentRequest, + mwaaClientProxy.client()::createEnvironment)) + .thenThrow(mwaaCreateEnvironmentExceptionClass) + .thenThrow(mwaaCreateEnvironmentExceptionClass) + .thenThrow(mwaaCreateEnvironmentExceptionClass) + .thenThrow(mwaaCreateEnvironmentExceptionClass) + .thenReturn(createEnvironmentResponse); + // then + ProgressEvent response = handler.handleRequest( + getProxies(), + request, + new CallbackContext()); + checkResponseNeedsCallback(response); + verify(getSdkClient(), times(expectedNumberOfInvocations)).createEnvironment( + any(CreateEnvironmentRequest.class)); + } + /** * Asserts throwing {@link CfnInvalidRequestException} when given model has invalid data. */ @@ -205,15 +295,21 @@ public void handleRequestInvalidInput() { final ResourceHandlerRequest request = ResourceHandlerRequest.builder() .desiredResourceState(model) .build(); + final CreateEnvironmentRequest awsCreateEnvironmentRequest = CreateTranslator.translateToCreateRequest(model); + final GetEnvironmentRequest awsGetEnvironmentRequest = ReadTranslator.translateToReadRequest(model); - when(getSdkClient().getEnvironment(any(GetEnvironmentRequest.class))) + ProxyClient mwaaClientProxy = getProxies().getMwaaClientProxy(); + + // when + when(mwaaClientProxy.injectCredentialsAndInvokeV2(awsGetEnvironmentRequest, + mwaaClientProxy.client()::getEnvironment)) // to indicate this environment does not exists and allow creation .thenThrow(ResourceNotFoundException.class); - when(getSdkClient().createEnvironment(any(CreateEnvironmentRequest.class))) + when(mwaaClientProxy.injectCredentialsAndInvokeV2(awsCreateEnvironmentRequest, + mwaaClientProxy.client()::createEnvironment)) .thenThrow(ValidationException.builder().message(INVALID_DATA).build()); - // when try { handler.handleRequest( getProxies(), @@ -224,6 +320,8 @@ public void handleRequestInvalidInput() { } catch (CfnInvalidRequestException e) { // expect exception assertThat(e.getMessage().contains(INVALID_DATA)).isTrue(); + verify(getSdkClient(), times(CreateHandler.MAX_RETRIES)).createEnvironment( + any(CreateEnvironmentRequest.class)); } } diff --git a/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/HandlerTestBase.java b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/HandlerTestBase.java index ca70ce8..88f1d30 100644 --- a/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/HandlerTestBase.java +++ b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/HandlerTestBase.java @@ -43,10 +43,13 @@ public class HandlerTestBase { private static final String PLUGINS_S3_OBJECT_VERSION = "PLUGINS_S3_OBJECT_VERSION"; private static final String REQUIREMENTS_S3_PATH = "REQUIREMENTS_S3_PATH"; private static final String REQUIREMENTS_S3_OBJECT_VERSION = "REQUIREMENTS_S3_OBJECT_VERSION"; + private static final String STARTUP_SCRIPT_S3_PATH = "STARTUP_SCRIPT_S3_PATH"; + private static final String STARTUP_SCRIPT_S3_OBJECT_VERSION = "STARTUP_SCRIPT_S3_OBJECT_VERSION"; private static final String ENVIRONMENT_CLASS = "ENVIRONMENT_CLASS"; private static final String WEEKLY_MAINTENANCE_WINDOW_START = "WEEKLY_MAINTENANCE_WINDOW_START"; private static final Integer MAX_WORKERS = 3; private static final Integer MIN_WORKERS = 1; + private static final Integer SCHEDULERS = 2; private static final String KEY = "KEY"; private static final String VALUE = "VALUE"; private static final String KEY_INTERNAL = "aws:tag:domain"; @@ -181,9 +184,12 @@ ResourceModel createCfnModel() { .pluginsS3ObjectVersion(PLUGINS_S3_OBJECT_VERSION) .requirementsS3Path(REQUIREMENTS_S3_PATH) .requirementsS3ObjectVersion(REQUIREMENTS_S3_OBJECT_VERSION) + .startupScriptS3Path(STARTUP_SCRIPT_S3_PATH) + .startupScriptS3ObjectVersion(STARTUP_SCRIPT_S3_OBJECT_VERSION) .environmentClass(ENVIRONMENT_CLASS) .maxWorkers(MAX_WORKERS) .minWorkers(MIN_WORKERS) + .schedulers(SCHEDULERS) .weeklyMaintenanceWindowStart(WEEKLY_MAINTENANCE_WINDOW_START) .airflowConfigurationOptions(ImmutableMap.of(KEY, VALUE)) .networkConfiguration(new NetworkConfiguration( @@ -213,9 +219,12 @@ Environment createApiEnvironment(final EnvironmentStatus status) { .pluginsS3ObjectVersion(PLUGINS_S3_OBJECT_VERSION) .requirementsS3Path(REQUIREMENTS_S3_PATH) .requirementsS3ObjectVersion(REQUIREMENTS_S3_OBJECT_VERSION) + .startupScriptS3Path(STARTUP_SCRIPT_S3_PATH) + .startupScriptS3ObjectVersion(STARTUP_SCRIPT_S3_OBJECT_VERSION) .environmentClass(ENVIRONMENT_CLASS) .maxWorkers(MAX_WORKERS) .minWorkers(MIN_WORKERS) + .schedulers(SCHEDULERS) .weeklyMaintenanceWindowStart(WEEKLY_MAINTENANCE_WINDOW_START) .airflowConfigurationOptions(ImmutableMap.of(KEY, VALUE)) .networkConfiguration(createNetworkConfiguration()) diff --git a/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/UpdateHandlerTest.java b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/UpdateHandlerTest.java index ec46624..c7d1597 100644 --- a/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/UpdateHandlerTest.java +++ b/aws-mwaa-environment/src/test/java/software/amazon/mwaa/environment/UpdateHandlerTest.java @@ -38,23 +38,22 @@ import software.amazon.cloudformation.proxy.ProgressEvent; import software.amazon.cloudformation.proxy.ResourceHandlerRequest; - /** * Tests for {@link UpdateHandler}. */ @ExtendWith(MockitoExtension.class) public class UpdateHandlerTest extends HandlerTestBase { private static final Integer UPDATED_MAX_WORKERS = 5; + private static final Integer UPDATED_MIN_WORKERS = 2; + private static final Integer UPDATED_SCHEDULERS = 3; private static final String NEW_TAG_KEY = "NEW_KEY"; private static final String NEW_TAG_VALUE = "NEW_VALUE"; private static final String INVALID_DATA = "INVALID_DATA"; - private static final Integer UPDATED_MIN_WORKERS = 2; private static final String LAST_UPDATE_ERROR_MESSAGE = "SOME_ERROR_MESSAGE"; private UpdateError error = UpdateError.builder().errorMessage(LAST_UPDATE_ERROR_MESSAGE).build(); private LastUpdate lastUpdateFailed = LastUpdate.builder().status(UpdateStatus.FAILED).error(error).build(); private LastUpdate lastUpdateSuccess = LastUpdate.builder().status(UpdateStatus.SUCCESS).build(); - /** * Prepares mocks. */ @@ -122,6 +121,9 @@ public void handleRequestSimpleSuccess() { verify(getSdkClient(), times(1)).tagResource(any(TagResourceRequest.class)); } + /** + * Tests a sad path. + */ @Test public void handleRequestUpdateFailed() { // given @@ -219,7 +221,7 @@ public void handleRequestUpdateEnvironmentUnavailable() { assertThat(response.getResourceModel()).isEqualTo(request.getDesiredResourceState()); assertThat(response.getResourceModels()).isNull(); assertThat(response.getMessage()).isEqualTo(String.format("Update failed, Environment unavailable. %s", - LAST_UPDATE_ERROR_MESSAGE)); + LAST_UPDATE_ERROR_MESSAGE)); assertThat(response.getErrorCode()).isEqualTo(HandlerErrorCode.NotStabilized); assertThat(response.getCallbackContext()).isNull(); } @@ -236,7 +238,6 @@ public void handleResourceMissingDuringUpdate() { final ResourceHandlerRequest request = ResourceHandlerRequest.builder() .desiredResourceState(model) .build(); - final GetEnvironmentResponse existing = createGetExistingEnvironmentResponse(); final GetEnvironmentResponse updating = createGetUpdatingEnvironmentResponse(); @@ -384,6 +385,8 @@ private ResourceModel createUpdatedCfnModel() { final ResourceModel model = createCfnModel(); model.setMaxWorkers(UPDATED_MAX_WORKERS); model.setMinWorkers(UPDATED_MIN_WORKERS); + model.setSchedulers(UPDATED_SCHEDULERS); + return model; } @@ -402,6 +405,7 @@ private GetEnvironmentResponse createGetUpdatedEnvironmentResponse() { .toBuilder() .maxWorkers(UPDATED_MAX_WORKERS) .minWorkers(UPDATED_MIN_WORKERS) + .schedulers(UPDATED_SCHEDULERS) .tags(ImmutableMap.of(NEW_TAG_KEY, NEW_TAG_VALUE)) .lastUpdate(lastUpdateSuccess) .build(); @@ -423,4 +427,5 @@ private GetEnvironmentResponse createGetUnavailableEnvironmentResponse() { .build(); return GetEnvironmentResponse.builder().environment(environment).build(); } + } diff --git a/aws-mwaa-environment/template.yml b/aws-mwaa-environment/template.yml index c280b67..afa3164 100644 --- a/aws-mwaa-environment/template.yml +++ b/aws-mwaa-environment/template.yml @@ -4,7 +4,7 @@ Description: AWS SAM template for the AWS::MWAA::Environment resource type Globals: Function: - Timeout: 180 # docker start-up times can be long for SAM CLI + Timeout: 300 # docker start-up times can be long for SAM CLI. Retry delays can also be long MemorySize: 256 Resources: