Merge pull request #114 from gridsmartercities/websocket_error_reporting

jamsidedown · web-flow · commit 1e3a327d2dc6 · 2020-03-16T12:56:46.000Z
Websocket response/error pushing
diff --git a/README.md b/README.md
@@ -44,6 +44,8 @@ The current list of AWS Lambda Python Decorators includes:
 * [__response_body_as_json__](#response_body_as_json): a decorator to transform a response dictionary body to a json string.
 * [__handle_all_exceptions__](#handle_all_exceptions): a decorator to handle all exceptions thrown by the lambda function.
 * [__cors__](#cors): a decorator to add cors headers to a lambda function.
+* [__push_ws_errors__](#push_ws_errors): a decorator to push unsuccessful responses back to the calling user via websockets with api gateway.
+* [__push_ws_responses__](#push_ws_response): a decorator to push all responses back to the calling user via websockets with api gateway.
 
 
 ### [Validators](https://github.com/gridsmartercities/aws-lambda-decorators/blob/master/aws_lambda_decorators/validators.py)
@@ -456,6 +458,45 @@ def cors_example():
 cors_example()  # returns {'statusCode': 200, 'headers': {'access-control-allow-origin': '*', 'access-control-allow-methods': 'POST', 'access-control-allow-headers': 'Content-Type', 'access-control-max-age': 86400}}
 ```
 
+### push_ws_errors
+
+This decorator pushes unsuccessful responses back to the calling client over websockets built on api gateway
+
+This decorator requires the client is connected to the websocket api gateway instance, and will therefore have a connection id
+
+Example:
+```py
+@push_ws_errors('https://api_id.execute_id.region.amazonaws.com/Prod')
+@handle_all_exceptions()
+def handler(event, context):
+    return {
+        'statusCode': 400,
+        'body': {
+            'message': 'Bad request'
+        }
+    }
+
+# will push {'type': 'error', 'statusCode': 400, 'message': 'Bad request'} back to the client via websockets
+```
+
+### push_ws_response
+
+This decorator pushes all responses back to the calling client over websockets built on api gateway
+
+This decorator requires the client is connected to the websocket api gateway instance, and will therefore have a connection id
+
+Example:
+```py
+@push_ws_response('https://api_id.execute_id.region.amazonaws.com/Prod')
+def handler(event, context):
+    return {
+        'statusCode': 200,
+        'body': 'Hello, world!'
+    }
+
+# will push {'statusCode': 200, 'body': 'Hello, world!'} back to the client via websockets
+```
+
 ## Writing your own validators
 
 You can create your own validators by inheriting from the Validator class.
diff --git a/aws_lambda_decorators/decorators.py b/aws_lambda_decorators/decorators.py
@@ -7,7 +7,8 @@
 import json
 from http import HTTPStatus
 import boto3
-from aws_lambda_decorators.utils import full_name, all_func_args, find_key_case_insensitive, failure, get_logger
+from aws_lambda_decorators.utils import (full_name, all_func_args, find_key_case_insensitive, failure, get_logger,
+                                         find_websocket_connection_id, get_websocket_endpoint)
 
 
 LOGGER = get_logger(__name__)
@@ -334,3 +335,50 @@ def update_header(headers, header_name, value, value_type):
                 return failure(CORS_NON_DICT_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR)
         return wrapper
     return decorator
+
+
+def push_ws_errors(websocket_endpoint_url: str):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            connection_id = find_websocket_connection_id(args)
+
+            response = func(*args, **kwargs)
+            success = response.get("statusCode", HTTPStatus.INTERNAL_SERVER_ERROR).value < 300
+
+            if connection_id and not success:
+                websocket_endpoint = get_websocket_endpoint(websocket_endpoint_url)
+
+                ws_response = {
+                    "type": "error",
+                    "statusCode": response.get("statusCode", HTTPStatus.INTERNAL_SERVER_ERROR),
+                    "message": json.loads(response.get("body", "{}")).get("message")
+                }
+
+                websocket_endpoint.post_to_connection(
+                    ConnectionId=connection_id,
+                    Data=json.dumps(ws_response)
+                )
+
+            return response
+        return wrapper
+    return decorator
+
+
+def push_ws_response(websocket_endpoint_url: str):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            connection_id = find_websocket_connection_id(args)
+
+            response = func(*args, **kwargs)
+
+            if connection_id:
+                websocket_endpoint = get_websocket_endpoint(websocket_endpoint_url)
+
+                websocket_endpoint.post_to_connection(
+                    ConnectionId=connection_id,
+                    Data=json.dumps(response)
+                )
+
+            return response
+        return wrapper
+    return decorator
diff --git a/aws_lambda_decorators/utils.py b/aws_lambda_decorators/utils.py
@@ -1,10 +1,13 @@
 """Utility functions."""
-import logging
-import os
-import json
+from functools import lru_cache
 from http import HTTPStatus
-import keyword
 import inspect
+import json
+import keyword
+import logging
+import os
+
+import boto3
 
 
 LOG_LEVEL = getattr(logging, os.getenv("LOG_LEVEL", "INFO"))
@@ -112,3 +115,39 @@ def failure(errors, status_code=HTTPStatus.BAD_REQUEST):
         "statusCode": status_code,
         "body": json.dumps({"message": errors})
     }
+
+
+def find_websocket_connection_id(args: list) -> str:
+    """
+    Finds an API Gateway connection id from the event dictionary in the
+    arguments of a lambda
+
+    Args:
+        args (list): a list of arguments from a lambda (*args)
+
+    Returns:
+        The connection id of a user as a string if found
+        None if not
+    """
+    for arg in args:
+        if isinstance(arg, dict) and "requestContext" in arg:
+            return arg["requestContext"].get("connectionId")
+    return None
+
+
+@lru_cache()
+def get_websocket_endpoint(endpoint_url: str) -> "ApiGatewayManagementApi":  # I can't find this in botocore.client
+    """
+    Gets an instance of ApiGatewayManagementApi for sending messages
+    through websockets
+
+    Args:
+        endpoint_url (str): an api gateway connection url (ish)
+
+    Returns:
+        The api gateway management client (cached)
+    """
+    return boto3.client(
+        "apigatewaymanagementapi",
+        endpoint_url=endpoint_url
+    )
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 LONG_DESCRIPTION = open("README.md").read()
 
 setup(name="aws-lambda-decorators",
-      version="0.47",
+      version="0.48",
       description="A set of python decorators to simplify aws python lambda development",
       long_description=LONG_DESCRIPTION,
       long_description_content_type="text/markdown",
diff --git a/tests/test_decorators.py b/tests/test_decorators.py
@@ -11,7 +11,9 @@
 
 from aws_lambda_decorators.classes import ExceptionHandler, Parameter, SSMParameter, ValidatedParameter
 from aws_lambda_decorators.decorators import extract, extract_from_event, extract_from_context, handle_exceptions, \
-    log, response_body_as_json, extract_from_ssm, validate, handle_all_exceptions, cors
+    log, response_body_as_json, extract_from_ssm, validate, handle_all_exceptions, cors, push_ws_errors, \
+    push_ws_response
+from aws_lambda_decorators.utils import get_websocket_endpoint
 from aws_lambda_decorators.validators import Mandatory, RegexValidator, SchemaValidator, Minimum, Maximum, MaxLength, \
     MinLength, Type, EnumValidator, NonEmpty, DateValidator, CurrencyValidator
 
@@ -1833,6 +1835,140 @@ def handler(event, a=None):  # noqa: pylint - unused-argument
                                                   "event",
                                                   "/a")
 
+    @patch("boto3.client")
+    def test_push_ws_errors_missing_parameter(self, mock_boto3_client):
+        get_websocket_endpoint.cache_clear()
+
+        event = {
+            "requestContext": {
+                "connectionId": "test_connection_id"
+            },
+            "body": json.dumps({
+                "invalid_property": "invalid_value"
+            })
+        }
+
+        @push_ws_errors(websocket_endpoint_url="https://api_id.execute_id.region.amazonaws.com/Prod")
+        @extract_from_event(parameters=[
+            Parameter(path="body[json]/valid_property", validators=[Mandatory])
+        ])
+        def lambda_func(event, context, valid_property=None):  # noqa: pylint - unused-argument
+            return {"statusCode": HTTPStatus.OK}
+
+        response = lambda_func(event, None)
+
+        expected_data = {
+            "type": "error",
+            "statusCode": 400,
+            "message": [{
+                "valid_property": ["Missing mandatory value"]
+            }]
+        }
+
+        self.assertEqual(response["statusCode"], HTTPStatus.BAD_REQUEST)
+
+        mock_boto3_client.assert_called_once_with(
+            "apigatewaymanagementapi",
+            endpoint_url="https://api_id.execute_id.region.amazonaws.com/Prod"
+        )
+
+        mock_boto3_client.return_value.post_to_connection.assert_called_once_with(
+            ConnectionId="test_connection_id",
+            Data=json.dumps(expected_data)
+        )
+
+    @patch("boto3.client")
+    def test_push_ws_errors_no_action_on_success(self, mock_boto3_client):
+        get_websocket_endpoint.cache_clear()
+
+        event = {
+            "requestContext": {
+                "connectionId": "test_connection_id"
+            }
+        }
+
+        @push_ws_errors(websocket_endpoint_url="https://api_id.execute_id.region.amazonaws.com/Prod")
+        def lambda_func(event, context):  # noqa: pylint - unused-argument
+            return {"statusCode": HTTPStatus.OK}
+
+        response = lambda_func(event, None)
+
+        self.assertEqual(response["statusCode"], 200)
+
+        mock_boto3_client.return_value.post_to_connection.assert_not_called()
+
+    @patch("boto3.client")
+    def test_push_ws_errors_no_connection_id(self, mock_boto3_client):
+        get_websocket_endpoint.cache_clear()
+
+        event = {
+            "body": {
+                "property": "value"
+            }
+        }
+
+        @push_ws_errors(websocket_endpoint_url="https://api_id.execute_id.region.amazonaws.com/Prod")
+        def lambda_func(event, context):  # noqa: pylint - unused-argument
+            return {"statusCode": HTTPStatus.BAD_REQUEST}
+
+        response = lambda_func(event, None)
+
+        self.assertEqual(response["statusCode"], 400)
+
+        mock_boto3_client.return_value.post_to_connection.assert_not_called()
+
+    @patch("boto3.client")
+    def test_push_ws_response(self, mock_boto3_client):
+        get_websocket_endpoint.cache_clear()
+
+        event = {
+            "requestContext": {
+                "connectionId": "test_connection_id"
+            }
+        }
+
+        @push_ws_response(websocket_endpoint_url="https://api_id.execute_id.region.amazonaws.com/Prod")
+        def lambda_func(event, context):  # noqa: pylint - unused-argument
+            return {
+                "statusCode": HTTPStatus.OK,
+                "body": "Hello, world!"
+            }
+
+        response = lambda_func(event, None)
+
+        self.assertEqual(response["statusCode"], 200)
+        self.assertEqual(response["body"], "Hello, world!")
+
+        mock_boto3_client.return_value.post_to_connection.assert_called_once_with(
+            ConnectionId="test_connection_id",
+            Data="{\"statusCode\": 200, \"body\": \"Hello, world!\"}"
+        )
+
+    @patch("boto3.client")
+    def test_push_ws_response_no_connection_id(self, mock_boto3_client):
+        get_websocket_endpoint.cache_clear()
+
+        event = {
+            "body": "Hello, world!"
+        }
+
+        @push_ws_response(websocket_endpoint_url="https://api_id.execute_id.region.amazonaws.com/Prod")
+        @extract_from_event(parameters=[
+            Parameter(path="body")
+        ])
+        def lambda_func(event, context, body=None):  # noqa: pylint - unused-argument
+            return {
+                "statusCode": HTTPStatus.OK,
+                "body": body
+            }
+
+        response = lambda_func(event, None)
+
+        self.assertEqual(response["statusCode"], 200)
+        self.assertEqual(response["body"], "Hello, world!")
+
+        mock_boto3_client.return_value.post_to_connection.assert_not_called()
+
 
 class IsolatedDecoderTests(unittest.TestCase):
     # Tests have been named so they run in a specific order
