added SQS schema & tests and sns skeleton

Author: Ran Isenberg

aws_lambda_powertools/utilities/validation/__init__.py

@@ -1,6 +1,6 @@
 """Validation utility
 """
-from .envelopes import DynamoDBEnvelope, EventBridgeEnvelope, UserEnvelope
+from .envelopes import DynamoDBEnvelope, EventBridgeEnvelope, SnsEnvelope, SqsEnvelope, UserEnvelope
 from .validator import validator
 
-__all__ = ["UserEnvelope", "DynamoDBEnvelope", "EventBridgeEnvelope", "validator"]
+__all__ = ["UserEnvelope", "DynamoDBEnvelope", "EventBridgeEnvelope", "SnsEnvelope", "SqsEnvelope", "validator"]

aws_lambda_powertools/utilities/validation/envelopes/__init__.py

@@ -1,9 +1,14 @@
 from .base import UserEnvelope
 from .dynamodb import DynamoDBEnvelope
 from .event_bridge import EventBridgeEnvelope
+from .sns import SnsEnvelope
+from .sqs import SqsEnvelope
+
 
 __all__ = [
     "UserEnvelope",
     "DynamoDBEnvelope",
     "EventBridgeEnvelope",
+    "SqsEnvelope",
+    "SnsEnvelope"
 ]

aws_lambda_powertools/utilities/validation/envelopes/base.py

@@ -13,15 +13,19 @@ def _parse_user_dict_schema(self, user_event: Dict[str, Any], inbound_schema_mod
         try:
             return inbound_schema_model(**user_event)
         except (ValidationError, TypeError):
-            logger.exception("Valdation exception while extracting user custom schema")
+            logger.exception("Validation exception while extracting user custom schema")
             raise
 
     def _parse_user_json_string_schema(self, user_event: str, inbound_schema_model: BaseModel) -> Any:
         logger.debug("parsing user dictionary schema")
+        if inbound_schema_model == str:
+            logger.debug("input is string, returning")
+            return user_event
+        logger.debug("trying to parse as json encoded string")
         try:
             return inbound_schema_model.parse_raw(user_event)
         except (ValidationError, TypeError):
-            logger.exception("Valdation exception while extracting user custom schema")
+            logger.exception("Validation exception while extracting user custom schema")
             raise
 
     @abstractmethod
@@ -34,5 +38,5 @@ def parse(self, event: Dict[str, Any], inbound_schema_model: BaseModel) -> Any:
         try:
             return inbound_schema_model(**event)
         except (ValidationError, TypeError):
-            logger.exception("Valdation exception received from input user custom envelopes event")
+            logger.exception("Validation exception received from input user custom envelopes event")
             raise

aws_lambda_powertools/utilities/validation/envelopes/dynamodb.py

@@ -3,8 +3,8 @@
 
 from pydantic import BaseModel, ValidationError
 
-from ..schemas import DynamoDBSchema
-from .base import BaseEnvelope
+from aws_lambda_powertools.utilities.validation.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.validation.schemas import DynamoDBSchema
 
 logger = logging.getLogger(__name__)
 
@@ -14,7 +14,7 @@ def parse(self, event: Dict[str, Any], inbound_schema_model: BaseModel) -> Any:
         try:
             parsed_envelope = DynamoDBSchema(**event)
         except (ValidationError, TypeError):
-            logger.exception("Valdation exception received from input dynamodb stream event")
+            logger.exception("Validation exception received from input dynamodb stream event")
             raise
         output = []
         for record in parsed_envelope.Records:

aws_lambda_powertools/utilities/validation/envelopes/event_bridge.py

@@ -3,8 +3,9 @@
 
 from pydantic import BaseModel, ValidationError
 
-from ..schemas import EventBridgeSchema
-from .base import BaseEnvelope
+from aws_lambda_powertools.utilities.validation.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.validation.schemas import EventBridgeSchema
+
 
 logger = logging.getLogger(__name__)
 
@@ -14,6 +15,6 @@ def parse(self, event: Dict[str, Any], inbound_schema_model: BaseModel) -> Any:
         try:
             parsed_envelope = EventBridgeSchema(**event)
         except (ValidationError, TypeError):
-            logger.exception("Valdation exception received from input eventbridge event")
+            logger.exception("Validation exception received from input eventbridge event")
             raise
         return self._parse_user_dict_schema(parsed_envelope.detail, inbound_schema_model)

aws_lambda_powertools/utilities/validation/envelopes/sns.py

@@ -0,0 +1,20 @@
+import logging
+from typing import Any, Dict
+
+from pydantic import BaseModel, ValidationError
+
+from aws_lambda_powertools.utilities.validation.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.validation.schemas import SnsSchema
+
+logger = logging.getLogger(__name__)
+
+
+class SnsEnvelope(BaseEnvelope):
+    def parse(self, event: Dict[str, Any], inbound_schema_model: BaseModel) -> Any:
+        try:
+            parsed_envelope = SnsSchema(**event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception received from input sqs event")
+            raise
+        ## TODO
+        return None

aws_lambda_powertools/utilities/validation/envelopes/sqs.py

@@ -0,0 +1,23 @@
+import logging
+from typing import Any, Dict
+
+from pydantic import BaseModel, ValidationError
+
+from aws_lambda_powertools.utilities.validation.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.validation.schemas import SqsSchema
+
+logger = logging.getLogger(__name__)
+
+
+class SqsEnvelope(BaseEnvelope):
+    def parse(self, event: Dict[str, Any], inbound_schema_model: BaseModel) -> Any:
+        try:
+            parsed_envelope = SqsSchema(**event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception received from input sqs event")
+            raise
+        output = []
+        for record in parsed_envelope.Records:
+            parsed_msg = self._parse_user_json_string_schema(record.body, inbound_schema_model)
+            output.append({"body": parsed_msg, "attributes": record.messageAttributes})
+        return output

aws_lambda_powertools/utilities/validation/schemas/sqs.py

@@ -1,5 +1,60 @@
-from pydantic import BaseModel
+import re
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, root_validator, validator
+from typing_extensions import Literal
+
+
+class SqsAttributesSchema(BaseModel):
+    ApproximateReceiveCount: str
+    ApproximateFirstReceiveTimestamp: datetime
+    MessageDeduplicationId: Optional[str]
+    MessageGroupId: Optional[str]
+    SenderId: str
+    SentTimestamp: datetime
+    SequenceNumber: Optional[str]
+
+
+class SqsMsgAttributeSchema(BaseModel):
+    stringValue: Optional[str]
+    binaryValue: Optional[str]
+    stringListValues: List[str] = []
+    binaryListValues: List[str] = []
+    dataType: str
+
+    @validator("dataType")
+    def valid_type(cls, v):  # noqa: VNE001
+        pattern = re.compile("Number.*|String.*|Binary.*")
+        if not pattern.match(v):
+            raise TypeError("data type is invalid")
+        return v
+
+    @root_validator
+    def check_str_and_binary_values(cls, values):
+        binary_val, str_val = values.get("binaryValue", ""), values.get("stringValue", "")
+        dataType = values.get("dataType")
+        if not str_val and not binary_val:
+            raise TypeError("both binaryValue and stringValue are missing")
+        if dataType.startswith("Binary") and not binary_val:
+            raise TypeError("binaryValue is missing")
+        if (dataType.startswith("String") or dataType.startswith("Number")) and not str_val:
+            raise TypeError("stringValue is missing")
+        return values
+
+
+class SqsRecordSchema(BaseModel):
+    messageId: str
+    receiptHandle: str
+    body: str
+    attributes: SqsAttributesSchema
+    messageAttributes: Dict[str, SqsMsgAttributeSchema]
+    md5OfBody: str
+    md5OfMessageAttributes: str
+    eventSource: Literal["aws:sqs"]
+    eventSourceARN: str
+    awsRegion: str
 
 
 class SqsSchema(BaseModel):
-    todo: str
+    Records: List[SqsRecordSchema]

aws_lambda_powertools/utilities/validation/validator.py

@@ -4,8 +4,7 @@
 from pydantic import BaseModel, ValidationError
 
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
-
-from .envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.validation.envelopes.base import BaseEnvelope
 
 logger = logging.getLogger(__name__)
 

tests/functional/test_validator.py

@@ -7,7 +7,13 @@
 from pydantic import BaseModel
 from pydantic.error_wrappers import ValidationError
 
-from aws_lambda_powertools.utilities.validation import DynamoDBEnvelope, EventBridgeEnvelope, UserEnvelope, validator
+from aws_lambda_powertools.utilities.validation import (
+    DynamoDBEnvelope,
+    EventBridgeEnvelope,
+    SqsEnvelope,
+    UserEnvelope,
+    validator
+)
 
 
 class OutboundSchema(BaseModel):
@@ -60,6 +66,7 @@ class MyMessage(BaseModel):
 @validator(inbound_schema_model=MyMessage, outbound_schema_model=OutboundSchema, envelope=DynamoDBEnvelope())
 def dynamodb_handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Optional[Any]]:
     assert event["custom"]
+    assert len(event["custom"]) == 3
     #  first record
     assert not event["custom"][0]["old"]
     assert event["custom"][0]["new"].message == "hello"
@@ -169,3 +176,95 @@ def test_eventbridge_fail_inbound_validation():
     event = {"greeting": "hello"}
     with pytest.raises(ValidationError):
         eventbridge_handler(event, LambdaContext())
+
+
+sqs_event_attribs = {
+    "test4": {
+        "stringValue": "dfgdfgfd",
+        "stringListValues": [],
+        "binaryListValues": [],
+        "dataType": "String.custom_type",
+    },
+    "test5": {"stringValue": "a,b,c,d", "stringListValues": [], "binaryListValues": [], "dataType": "String"},
+    "tes6": {"stringValue": "112.1", "stringListValues": [], "binaryListValues": [], "dataType": "Number.mytype"},
+    "test2": {"stringValue": "111", "stringListValues": [], "binaryListValues": [], "dataType": "Number"},
+    "test3": {"binaryValue": "w5NNNcOXXXU=", "stringListValues": [], "binaryListValues": [], "dataType": "Binary"},
+    "test": {"stringValue": "gfgf", "stringListValues": [], "binaryListValues": [], "dataType": "String"},
+}
+
+
+@validator(inbound_schema_model=MyMessage, outbound_schema_model=OutboundSchema, envelope=SqsEnvelope())
+def sqs_json_body_handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Optional[Any]]:
+    assert event["custom"]
+    assert len(event["custom"]) == 1
+    assert event["orig"]
+    assert event["custom"][0]["body"].message == "hello"
+    assert event["custom"][0]["body"].messageId == 8
+    assert len(event["custom"][0]["attributes"]) == len(sqs_event_attribs)
+    return {"response_code": 200, "message": "working"}
+
+
+def test_sqs_ok_json_string_body_validation():
+    event = {
+        "Records": [
+            {
+                "messageId": "1743e893-cc24-1234-88f8-f80c37dcd923",
+                "receiptHandle": "AKhXK7azPaZHY0zjmTsdfsdfdsfOgcVob",
+                "body": '{"message": "hello", "messageId": 8}',
+                "attributes": {
+                    "ApproximateReceiveCount": "1",
+                    "SentTimestamp": "1598117108660",
+                    "SenderId": "43434dsdfd:sdfds",
+                    "ApproximateFirstReceiveTimestamp": "1598117108667",
+                },
+                "messageAttributes": sqs_event_attribs,
+                "md5OfBody": "4db76498a982d84c188927c585076a6c",
+                "md5OfMessageAttributes": "7186428dc148b402947274e0bb41e7ee",
+                "eventSource": "aws:sqs",
+                "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:mytest",
+                "awsRegion": "us-west-1",
+            }
+        ]
+    }
+    sqs_json_body_handler(event, LambdaContext())
+
+
+@validator(inbound_schema_model=str, outbound_schema_model=OutboundSchema, envelope=SqsEnvelope())
+def sqs_string_body_handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Optional[Any]]:
+    assert event["custom"]
+    assert len(event["custom"]) == 1
+    assert event["orig"]
+    assert event["custom"][0]["body"] == "hello how are you"
+    assert len(event["custom"][0]["attributes"]) == len(sqs_event_attribs)
+    return {"response_code": 200, "message": "working"}
+
+
+def test_sqs_ok_json_string_validation():
+    event = {
+        "Records": [
+            {
+                "messageId": "1743e893-cc24-1234-88f8-f80c37dcd923",
+                "receiptHandle": "AKhXK7azPaZHY0zjmTsdfsdfdsfOgcVob",
+                "body": "hello how are you",
+                "attributes": {
+                    "ApproximateReceiveCount": "1",
+                    "SentTimestamp": "1598117108660",
+                    "SenderId": "43434dsdfd:sdfds",
+                    "ApproximateFirstReceiveTimestamp": "1598117108667",
+                },
+                "messageAttributes": sqs_event_attribs,
+                "md5OfBody": "4db76498a982d84c188927c585076a6c",
+                "md5OfMessageAttributes": "7186428dc148b402947274e0bb41e7ee",
+                "eventSource": "aws:sqs",
+                "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:mytest",
+                "awsRegion": "us-west-1",
+            }
+        ]
+    }
+    sqs_string_body_handler(event, LambdaContext())
+
+
+def test_sqs_fail_inbound_validation():
+    event = {"greeting": "hello"}
+    with pytest.raises(ValidationError):
+        sqs_string_body_handler(event, LambdaContext())