Repository: Miserlou/Zappa

Issue:
Ability to capture the responses from async tasks
## Context

The async tasks functionality is really useful, but wouldn't it be more useful if there was a way to capture the responses from functions called asynchronously?

## Expected Behavior

Functions run asynchronously, and their responses are queryable by the calling context and other future requests.

## Actual Behavior

Functions run, but their output is unavailable to the caller, and also unavailable to future app requests.

## Possible Fix

Assign asynchronous tasks an ID, and store the responses in a more-persistent datastore for later retrieval using this ID.



Discussion:

Resolving patch:
diff --git a/zappa/async.py b/zappa/async.py
--- a/zappa/async.py
+++ b/zappa/async.py
@@ -92,15 +92,23 @@ def my_async_func(*args, **kwargs):
 import inspect
 import json
 import os
+import uuid
+import time
 
 from .utilities import get_topic_name
 
+try:
+    from zappa_settings import ASYNC_RESPONSE_TABLE
+except ImportError:
+    ASYNC_RESPONSE_TABLE = None
+
 # Declare these here so they're kept warm.
 try:
     aws_session = boto3.Session()
     LAMBDA_CLIENT = aws_session.client('lambda')
     SNS_CLIENT = aws_session.client('sns')
     STS_CLIENT = aws_session.client('sts')
+    DYNAMODB_CLIENT = aws_session.client('dynamodb')
 except botocore.exceptions.NoRegionError as e: # pragma: no cover
     # This can happen while testing on Travis, but it's taken care  of
     # during class initialization.
@@ -121,7 +129,7 @@ class LambdaAsyncResponse(object):
     Base Response Dispatcher class
     Can be used directly or subclassed if the method to send the message is changed.
     """
-    def __init__(self, lambda_function_name=None, aws_region=None, **kwargs):
+    def __init__(self, lambda_function_name=None, aws_region=None, capture_response=False, **kwargs):
         """ """
         if kwargs.get('boto_session'):
             self.client = kwargs.get('boto_session').client('lambda')
@@ -130,6 +138,21 @@ def __init__(self, lambda_function_name=None, aws_region=None, **kwargs):
 
         self.lambda_function_name = lambda_function_name
         self.aws_region = aws_region
+        if capture_response:
+            if ASYNC_RESPONSE_TABLE is None:
+                print(
+                    "Warning! Attempted to capture a response without "
+                    "async_response_table configured in settings (you won't "
+                    "capture async responses)."
+                )
+                capture_response = False
+                self.response_id = "MISCONFIGURED"
+
+            else:
+                self.response_id = str(uuid.uuid4())
+
+        self.capture_response = capture_response
+
 
     def send(self, task_path, args, kwargs):
         """
@@ -137,6 +160,8 @@ def send(self, task_path, args, kwargs):
         """
         message = {
                 'task_path': task_path,
+                'capture_response': self.capture_response,
+                'response_id': self.response_id,
                 'args': args,
                 'kwargs': kwargs
             }
@@ -163,16 +188,17 @@ class SnsAsyncResponse(LambdaAsyncResponse):
     Send a SNS message to a specified SNS topic
     Serialise the func path and arguments
     """
-    def __init__(self, lambda_function_name=None, aws_region=None, **kwargs):
+    def __init__(self, lambda_function_name=None, aws_region=None, capture_response=False, **kwargs):
 
         self.lambda_function_name = lambda_function_name
-        self.aws_region = aws_region
+        self.aws_region=aws_region
 
         if kwargs.get('boto_session'):
             self.client = kwargs.get('boto_session').client('sns')
         else: # pragma: no cover
             self.client = SNS_CLIENT
 
+
         if kwargs.get('arn'):
             self.arn = kwargs.get('arn')
         else:
@@ -187,6 +213,11 @@ def __init__(self, lambda_function_name=None, aws_region=None, **kwargs):
                                     topic_name=get_topic_name(self.lambda_function_name)
                                 )
 
+        self.capture_response = capture_response
+        if capture_response:
+            self.response_id = str(uuid.uuid4())
+
+
     def _send(self, message):
         """
         Given a message, publish to this topic.
@@ -238,24 +269,48 @@ def run_message(message):
     'task_path', 'args', and 'kwargs' used by lambda routing
     and a 'command' in handler.py
     """
+    if message.get('capture_response', False):
+        DYNAMODB_CLIENT.put_item(
+            TableName=ASYNC_RESPONSE_TABLE,
+            Item={
+                'id': {'S': str(message['response_id'])},
+                'ttl': {'N': str(int(time.time()+600))},
+                'async_status': {'S': 'in progress'},
+                'async_response': {'S': str(json.dumps('N/A'))},
+            }
+        )
+
     func = import_and_get_task(message['task_path'])
     if hasattr(func, 'sync'):
-        return func.sync(
+        response = func.sync(
             *message['args'],
             **message['kwargs']
         )
     else:
-        return func(
+        response = func(
             *message['args'],
             **message['kwargs']
         )
 
+    if message.get('capture_response', False):
+        DYNAMODB_CLIENT.update_item(
+            TableName=ASYNC_RESPONSE_TABLE,
+            Key={'id': {'S': str(message['response_id'])}},
+            UpdateExpression="SET async_response = :r, async_status = :s",
+            ExpressionAttributeValues={
+                ':r': {'S': str(json.dumps(response))},
+                ':s': {'S': 'complete'},
+            },
+        )
+
+    return response
+
 ##
 # Execution interfaces and classes
 ##
 
 
-def run(func, args=[], kwargs={}, service='lambda',
+def run(func, args=[], kwargs={}, service='lambda', capture_response=False,
         remote_aws_lambda_function_name=None, remote_aws_region=None, **task_kwargs):
     """
     Instead of decorating a function with @task, you can just run it directly.
@@ -276,6 +331,7 @@ def run(func, args=[], kwargs={}, service='lambda',
     task_path = get_func_task_path(func)
     return ASYNC_CLASSES[service](lambda_function_name=lambda_function_name,
                                   aws_region=aws_region,
+                                  capture_response=capture_response,
                                   **task_kwargs).send(task_path, args, kwargs)
 
 
@@ -315,6 +371,8 @@ def task(*args, **kwargs):
         lambda_function_name = kwargs.get('remote_aws_lambda_function_name') or os.environ.get('AWS_LAMBDA_FUNCTION_NAME')
         aws_region = kwargs.get('remote_aws_region') or os.environ.get('AWS_REGION')
 
+    capture_response = kwargs.get('capture_response', False)
+
     def func_wrapper(func):
 
         task_path = get_func_task_path(func)
@@ -341,7 +399,8 @@ def _run_async(*args, **kwargs):
             """
             if (service in ASYNC_CLASSES) and (lambda_function_name):
                 send_result = ASYNC_CLASSES[service](lambda_function_name=lambda_function_name,
-                                                     aws_region=aws_region).send(task_path, args, kwargs)
+                                                     aws_region=aws_region,
+                                                     capture_response=capture_response).send(task_path, args, kwargs)
                 return send_result
             else:
                 return func(*args, **kwargs)
@@ -388,3 +447,17 @@ def get_func_task_path(func):
                                         func_name=func.__name__
                                     )
     return task_path
+
+
+def get_async_response(response_id):
+    response = DYNAMODB_CLIENT.get_item(
+        TableName=ASYNC_RESPONSE_TABLE,
+        Key={'id': {'S': str(response_id)}}
+    )
+    if 'Item' not in response:
+        return None
+
+    return {
+        'status': response['Item']['async_status']['S'],
+        'response': json.loads(response['Item']['async_response']['S']),
+    }
diff --git a/zappa/cli.py b/zappa/cli.py
--- a/zappa/cli.py
+++ b/zappa/cli.py
@@ -1079,6 +1079,25 @@ def schedule(self):
             )
             click.echo('SNS Topic created: %s' % topic_arn)
 
+        # Add async tasks DynamoDB
+        table_name = self.stage_config.get('async_response_table', False)
+        read_capacity = self.stage_config.get('async_response_table_read_capacity', 1)
+        write_capacity = self.stage_config.get('async_response_table_write_capacity', 1)
+        if table_name and self.stage_config.get('async_resources', True):
+            created, response_table = self.zappa.create_async_dynamodb_table(
+                table_name, read_capacity, write_capacity)
+            if created:
+                click.echo('DynamoDB table created: %s' % table_name)
+            else:
+                click.echo('DynamoDB table exists: %s' % table_name)
+                provisioned_throughput = response_table['Table']['ProvisionedThroughput']
+                if provisioned_throughput['ReadCapacityUnits'] != read_capacity or \
+                    provisioned_throughput['WriteCapacityUnits'] != write_capacity:
+                        click.echo(click.style(
+                            "\nWarning! Existing DynamoDB table ({}) does not match configured capacity.\n".format(table_name),
+                            fg='red'
+                        ))
+
     def unschedule(self):
         """
         Given a a list of scheduled functions,
@@ -2208,6 +2227,10 @@ def create_package(self, output=None):
                 django_py = ''.join(os.path.join(base, 'ext', 'django_zappa.py'))
                 lambda_zip.write(django_py, 'django_zappa_app.py')
 
+            # async response
+            async_response_table = self.stage_config.get('async_response_table', '')
+            settings_s += "ASYNC_RESPONSE_TABLE='{0!s}'\n".format(async_response_table)
+
             # Lambda requires a specific chmod
             temp_settings = tempfile.NamedTemporaryFile(delete=False)
             os.chmod(temp_settings.name, 0o644)
diff --git a/zappa/core.py b/zappa/core.py
--- a/zappa/core.py
+++ b/zappa/core.py
@@ -266,6 +266,7 @@ def __init__(self,
             self.route53 = self.boto_session.client('route53')
             self.sns_client = self.boto_session.client('sns')
             self.cf_client = self.boto_session.client('cloudformation')
+            self.dynamodb_client = self.boto_session.client('dynamodb')
 
         self.cf_template = troposphere.Template()
         self.cf_api_resources = []
@@ -2310,6 +2311,74 @@ def remove_async_sns_topic(self, lambda_name):
                 removed_arns.append(sub['TopicArn'])
         return removed_arns
 
+
+    ###
+    # Async / DynamoDB
+    ##
+
+    def _set_async_dynamodb_table_ttl(self, table_name):
+        self.dynamodb_client.update_time_to_live(
+            TableName=table_name,
+            TimeToLiveSpecification={
+                'Enabled': True,
+                'AttributeName': 'ttl'
+            }
+        )
+
+
+    def create_async_dynamodb_table(self, table_name, read_capacity, write_capacity):
+        """
+        Create the DynamoDB table for async task return values
+        """
+        try:
+            dynamodb_table = self.dynamodb_client.describe_table(TableName=table_name)
+            return False, dynamodb_table
+
+        # catch this exception (triggered if the table doesn't exist)
+        except botocore.exceptions.ClientError:
+            dynamodb_table = self.dynamodb_client.create_table(
+                AttributeDefinitions=[
+                    {
+                        'AttributeName': 'id',
+                        'AttributeType': 'S'
+                    }
+                ],
+                TableName=table_name,
+                KeySchema=[
+                    {
+                        'AttributeName': 'id',
+                        'KeyType': 'HASH'
+                    },
+                ],
+                ProvisionedThroughput = {
+                    'ReadCapacityUnits': read_capacity,
+                    'WriteCapacityUnits': write_capacity
+                }
+            )
+            if dynamodb_table:
+                try:
+                    self._set_async_dynamodb_table_ttl(table_name)
+                except botocore.exceptions.ClientError:
+                    # this fails because the operation is async, so retry
+                    time.sleep(10)
+                    self._set_async_dynamodb_table_ttl(table_name)
+
+        return True, dynamodb_table
+
+
+    def remove_async_dynamodb_table(self, table_name):
+        """
+        Remove the DynamoDB Table used for async return values
+        """
+
+        topic_name = get_topic_name(lambda_name)
+        removed_arns = []
+        for sub in self.sns_client.list_subscriptions()['Subscriptions']:
+            if topic_name in sub['TopicArn']:
+                self.sns_client.delete_topic(TopicArn=sub['TopicArn'])
+                removed_arns.append(sub['TopicArn'])
+        return removed_arns
+
     ##
     # CloudWatch Logging
     ##


Test patch: