Add SNS Topic ARN to CloudFormation Output

The ARN is now output as SNSTopicArn and export as <STACK_NAME>-SNSTopicArn

Author: pwrmiller

source/aws_lambda/s3_event/handler.py

@@ -10,9 +10,8 @@
 #  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for   #
 #  the specific language governing permissions and limitations under the License.                                      #
 # ######################################################################################################################
-
+import json
 import os
-from typing import List
 
 from aws_lambda_powertools import Logger, Tracer, Metrics
 from aws_lambda_powertools.metrics import MetricUnit
@@ -43,21 +42,50 @@ def solution_name() -> str:
     return os.environ["SOLUTION_NAME"]
 
 
-def send_configuration_error(errors: List[str]):
+def send_configuration_error(configuration: Configuration):
+    errors = configuration.errors
     sns = get_service_client("sns")
+    dataset_group = configuration.dataset_group
+
     subject = f"{solution_name()} Notifications"
 
-    message = "There were errors detected when reading a personalization job configuration file:\n\n"
-    for error in errors:
-        logger.error(f"Personalization job configuration error: {error}")
-        message += f"   - {error}\n"
-    message += "\nPlease correct these errors and upload the configuration again."
+    def build_default_message():
+        f"The personalization workflow for {configuration.dataset_group} completed with errors."
+
+    def build_json_message():
+        return json.dumps(
+            {
+                "datasetGroup": dataset_group,
+                "status": "UPDATE FAILED",
+                "summary": "There were errors detected when reading a personalization job configuration file",
+                "description": [error for error in errors],
+            }
+        )
+
+    def build_long_message():
+        message = "There were errors detected when reading a personalization job configuration file:\n\n"
+        for error in errors:
+            logger.error(f"Personalization job configuration error: {error}")
+            message += f"   - {error}\n"
+        message += "\nPlease correct these errors and upload the configuration again."
+        return message
 
+    logger.error("publishing configuration error to SQS")
     sns.publish(
         TopicArn=topic_arn(),
-        Message=message,
+        Message=json.dumps(
+            {
+                "default": build_default_message(),
+                "sms": build_default_message(),
+                "email": build_long_message(),
+                "email-json": build_json_message(),
+                "sqs": build_json_message(),
+            }
+        ),
+        MessageStructure="json",
         Subject=subject,
     )
+    logger.error("published configuration error to SQS")
 
 
 @metrics.log_metrics
@@ -86,7 +114,7 @@ def lambda_handler(event, context):
         configuration = Configuration()
         configuration.load(config_text)
         if configuration.errors:
-            send_configuration_error(configuration.errors)
+            send_configuration_error(configuration)
             metrics.add_metric(
                 "ConfigurationsProcessedFailures", unit=MetricUnit.Count, value=1
             )
@@ -98,7 +126,7 @@ def lambda_handler(event, context):
             metrics.add_metric(
                 "ConfigurationsProcessedFailures", unit=MetricUnit.Count, value=1
             )
-            send_configuration_error(configuration.errors)
+            send_configuration_error(configuration)
         else:
             config = configuration.config_dict
             config = set_bucket(config, bucket, key)

source/aws_lambda/shared/personalize_service.py

@@ -609,6 +609,7 @@ class Configuration:
     def __init__(self):
         self._configuration_errors = []
         self.config_dict = {}
+        self.dataset_group = "UNKNOWN"
 
     def load(self, content: Union[Path, str]):
         if isinstance(content, Path):
@@ -670,6 +671,8 @@ def _validate_dataset_group(self, path="datasetGroup.serviceConfig"):
             )
         else:
             self._validate_resource(DatasetGroup(), dataset_group)
+            if isinstance(dataset_group, dict):
+                self.dataset_group = dataset_group.get("name", self.dataset_group)
 
     def _validate_event_tracker(self, path="eventTracker.serviceConfig"):
         event_tracker = jmespath.search(path, self.config_dict)

source/aws_lambda/sns_notification/handler.py

@@ -67,6 +67,26 @@ def __init__(self, event: Dict, context: LambdaContext):
             metrics.add_metric("JobSuccess", unit=MetricUnit.Count, value=1)
             self.message = self._build_success_message()
 
+        self.default = self._build_default_message()
+        self.sms = self._build_sms_message()
+        self.json = self._build_json_message()
+
+    def _build_json_message(self):
+        return json.dumps(
+            {
+                "datasetGroup": self.dataset_group,
+                "status": "UPDATE FAILED" if self.error else "UPDATE COMPLETE",
+                "summary": self._build_default_message(),
+                "description": self.message,
+            }
+        )
+
+    def _build_default_message(self) -> str:
+        return f"The personalization workflow for {self.dataset_group} completed {'with errors' if self.error else 'successfully'}"
+
+    def _build_sms_message(self) -> str:
+        return self._build_default_message()
+
     def _build_error_message(self) -> str:
         """
         Build the error message
@@ -116,12 +136,21 @@ def lambda_handler(event, context):
     :return: None
     """
     sns = get_service_client("sns")
-    message = MessageBuilder(event, context).message
+    message_builder = MessageBuilder(event, context)
     subject = f"{solution_name()} Notifications"
 
     logger.info("publishing message for event", extra={"event": event})
     sns.publish(
         TopicArn=topic_arn(),
-        Message=message,
+        Message=json.dumps(
+            {
+                "default": message_builder.default,
+                "sms": message_builder.sms,
+                "email": message_builder.message,
+                "email-json": message_builder.json,
+                "sqs": message_builder.json,
+            }
+        ),
+        MessageStructure="json",
         Subject=subject,
     )

source/infrastructure/personalize/stack.py

@@ -68,6 +68,25 @@ def __init__(
         super().__init__(scope, construct_id, *args, **kwargs)
 
         # CloudFormation Parameters
+        self.email = cdk.CfnParameter(
+            self,
+            id="Email",
+            type="String",
+            description="Email to notify with personalize workflow results",
+            default="",
+            max_length=50,
+            allowed_pattern=r"(^[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+$|^$)",
+            constraint_description="Must be a valid email address or blank",
+        )
+        self.solutions_template_options.add_parameter(
+            self.email, "Email", "Solution Configuration"
+        )
+        self.email_provided = CfnCondition(
+            self,
+            "EmailProvided",
+            expression=Fn.condition_not(Fn.condition_equals(self.email, "")),
+        )
+
         self.personalize_kms_key_arn = cdk.CfnParameter(
             self,
             id="PersonalizeKmsKeyArn",
@@ -88,25 +107,6 @@ def __init__(
             ),
         )
 
-        self.email = cdk.CfnParameter(
-            self,
-            id="Email",
-            type="String",
-            description="Email to notify with personalize workflow results",
-            default="",
-            max_length=50,
-            allowed_pattern=r"(^[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+$|^$)",
-            constraint_description="Must be a valid email address or blank",
-        )
-        self.solutions_template_options.add_parameter(
-            self.email, "Email", "Solution Configuration"
-        )
-        self.email_provided = CfnCondition(
-            self,
-            "EmailProvided",
-            expression=Fn.condition_not(Fn.condition_equals(self.email, "")),
-        )
-
         # layers
         layer_powertools = PowertoolsLayer.get_or_create(self)
         layer_solutions = SolutionsLayer.get_or_create(self)
@@ -413,3 +413,9 @@ def __init__(
             value=self.dashboard.name,
             export_name=f"{Aws.STACK_NAME}-Dashboard",
         )
+        cdk.CfnOutput(
+            self,
+            "SNSTopicArn",
+            value=notifications.topic.topic_arn,
+            export_name=f"{Aws.STACK_NAME}-SNSTopicArn",
+        )

source/tests/aws_lambda/sns_notification/test_sns_notification.py

@@ -10,23 +10,40 @@
 #  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for   #
 #  the specific language governing permissions and limitations under the License.                                      #
 # ######################################################################################################################
+import json
 import os
 from collections import namedtuple
 
+import boto3
 import pytest
-from botocore.stub import Stubber
+from moto import mock_sns, mock_sqs
 
 from aws_lambda.sns_notification.handler import lambda_handler
-from aws_solutions.core import get_service_client
 
 TRACE_ID = "1-57f5498f-d91047849216d0f2ea3b6442"
 
 
 @pytest.fixture
-def sns_stubber():
-    sns_client = get_service_client("sns")
-    with Stubber(sns_client) as stubber:
-        yield stubber
+def sqs_mock():
+    topic_arn = os.environ.get("SNS_TOPIC_ARN")
+    topic_name = topic_arn.split(":")[-1]
+
+    with mock_sqs():
+        with mock_sns():
+
+            cli = boto3.client("sns")
+            cli.create_topic(Name=topic_name)
+
+            sqs = boto3.client("sqs")
+            sqs.create_queue(QueueName="TestQueue")
+
+            cli.subscribe(
+                TopicArn=topic_arn,
+                Protocol="sqs",
+                Endpoint=f"arn:aws:sqs:us-east-1:{'1'*12}:TestQueue",
+            )
+
+            yield sqs
 
 
 @pytest.fixture
@@ -37,23 +54,6 @@ def trace_enabled():
 
 
 DATASET_GROUP_NAME = "DATASET_GROUP_NAME"
-EXPECTED_MESSAGE = """
-There was an error running the personalization job for dataset group DATASET_GROUP_NAME
-
-Message: ERROR_MESSAGE
-
-""".lstrip(
-    "\n"
-)
-EXPECTED_MESSAGE_TRACE = f"""
-There was an error running the personalization job for dataset group DATASET_GROUP_NAME
-
-Message: ERROR_MESSAGE
-
-Traces: https://console.aws.amazon.com/xray/home?region=us-east-1#/traces/{TRACE_ID}
-""".strip(
-    "\n"
-)
 
 
 @pytest.fixture
@@ -62,18 +62,8 @@ def context():
     return ctx(f"arn:aws:lambda:us-east-1:{'1' * 12}:function:my-function:1")
 
 
-def test_sns_notification(sns_stubber, context):
+def test_sns_notification(context, sqs_mock):
     """Test without traces"""
-    sns_stubber.add_response(
-        "publish",
-        {},
-        expected_params={
-            "TopicArn": os.environ.get("SNS_TOPIC_ARN"),
-            "Subject": "Maintaining Personalized Experiences with Machine Learning Notifications",
-            "Message": EXPECTED_MESSAGE,
-        },
-    )
-
     lambda_handler(
         {
             "datasetGroup": DATASET_GROUP_NAME,
@@ -84,19 +74,32 @@ def test_sns_notification(sns_stubber, context):
         context,
     )
 
+    url = sqs_mock.get_queue_url(QueueName="TestQueue")["QueueUrl"]
+    msg = json.loads(
+        json.loads(
+            sqs_mock.receive_message(QueueUrl=url, MaxNumberOfMessages=1,)["Messages"][
+                0
+            ]["Body"]
+        )["Message"]
+    )
 
-def test_sns_notification_trace(sns_stubber, trace_enabled, context):
-    """Test with traces"""
-    sns_stubber.add_response(
-        "publish",
-        {},
-        expected_params={
-            "TopicArn": os.environ.get("SNS_TOPIC_ARN"),
-            "Subject": "Maintaining Personalized Experiences with Machine Learning Notifications",
-            "Message": EXPECTED_MESSAGE_TRACE,
-        },
+    error_default = (
+        f"The personalization workflow for {DATASET_GROUP_NAME} completed with errors"
     )
+    error_json = {
+        "datasetGroup": DATASET_GROUP_NAME,
+        "status": "UPDATE FAILED",
+        "summary": f"The personalization workflow for {DATASET_GROUP_NAME} completed with errors",
+        "description": f"There was an error running the personalization job for dataset group {DATASET_GROUP_NAME}\n\nMessage: ERROR_MESSAGE\n\n",
+    }
+
+    assert msg["default"] == error_default
+    assert msg["sms"] == error_default
+    assert json.loads(msg["sqs"]) == error_json
+
 
+def test_sns_notification_trace(sqs_mock, trace_enabled, context):
+    """Test with traces"""
     lambda_handler(
         {
             "datasetGroup": DATASET_GROUP_NAME,
@@ -106,3 +109,26 @@ def test_sns_notification_trace(sns_stubber, trace_enabled, context):
         },
         context,
     )
+
+    url = sqs_mock.get_queue_url(QueueName="TestQueue")["QueueUrl"]
+    msg = json.loads(
+        json.loads(
+            sqs_mock.receive_message(QueueUrl=url, MaxNumberOfMessages=1,)["Messages"][
+                0
+            ]["Body"]
+        )["Message"]
+    )
+
+    error_default = (
+        f"The personalization workflow for {DATASET_GROUP_NAME} completed with errors"
+    )
+    error_json = {
+        "datasetGroup": f"{DATASET_GROUP_NAME}",
+        "status": "UPDATE FAILED",
+        "summary": f"The personalization workflow for {DATASET_GROUP_NAME} completed with errors",
+        "description": f"There was an error running the personalization job for dataset group {DATASET_GROUP_NAME}\n\nMessage: ERROR_MESSAGE\n\nTraces: https://console.aws.amazon.com/xray/home?region=us-east-1#/traces/1-57f5498f-d91047849216d0f2ea3b6442",
+    }
+
+    assert msg["default"] == error_default
+    assert msg["sms"] == error_default
+    assert json.loads(msg["sqs"]) == error_json