def return_response(self, method, path, data, headers, response):
        if path.startswith('/shell') or method == 'GET':
            return

        data = json.loads(to_str(data))

        # update table definitions
        if data and 'TableName' in data and 'KeySchema' in data:
            TABLE_DEFINITIONS[data['TableName']] = data

        if response._content:
            # fix the table and latest stream ARNs (DynamoDBLocal hardcodes "ddblocal" as the region)
            content_replaced = re.sub(
                r'("TableArn"|"LatestStreamArn"|"StreamArn")\s*:\s*"arn:aws:dynamodb:ddblocal:([^"]+)"',
                r'\1: "arn:aws:dynamodb:%s:\2"' % aws_stack.get_region(),
                to_str(response._content))
            if content_replaced != response._content:
                response._content = content_replaced
                fix_headers_for_updated_response(response)

        action = headers.get('X-Amz-Target', '')
        action = action.replace(ACTION_PREFIX, '')
        if not action:
            return
        # upgrade event version to 1.1
        record = {
            'eventID': '1',
            'eventVersion': '1.1',
            'dynamodb': {
                'ApproximateCreationDateTime': time.time(),
                'StreamViewType': 'NEW_AND_OLD_IMAGES',
                'SizeBytes': -1
            },
            'awsRegion': aws_stack.get_region(),
            'eventSource': 'aws:dynamodb'
        }
        records = [record]

        streams_enabled_cache = {}
        table_name = data.get('TableName')
        event_sources_or_streams_enabled = has_event_sources_or_streams_enabled(
            table_name, streams_enabled_cache)
        if action == 'UpdateItem':
            if response.status_code == 200 and event_sources_or_streams_enabled:
                existing_item = self._thread_local('existing_item')
                record[
                    'eventName'] = 'INSERT' if not existing_item else 'MODIFY'
                record['eventID'] = short_uid()
                updated_item = find_existing_item(data)
                if not updated_item:
                    return
                record['dynamodb']['Keys'] = data['Key']
                if existing_item:
                    record['dynamodb']['OldImage'] = existing_item
                record['dynamodb']['NewImage'] = updated_item
                record['dynamodb']['SizeBytes'] = len(json.dumps(updated_item))
                stream_spec = dynamodb_get_table_stream_specification(
                    table_name=table_name)
                if stream_spec:
                    record['dynamodb']['StreamViewType'] = stream_spec[
                        'StreamViewType']
        elif action == 'BatchWriteItem':
            records = self.prepare_batch_write_item_records(record, data)
            for record in records:
                event_sources_or_streams_enabled = (
                    event_sources_or_streams_enabled
                    or has_event_sources_or_streams_enabled(
                        record['eventSourceARN'], streams_enabled_cache))

        elif action == 'TransactWriteItems':
            records = self.prepare_transact_write_item_records(record, data)
            for record in records:
                event_sources_or_streams_enabled = (
                    event_sources_or_streams_enabled
                    or has_event_sources_or_streams_enabled(
                        record['eventSourceARN'], streams_enabled_cache))

        elif action == 'PutItem':
            if response.status_code == 200:
                keys = dynamodb_extract_keys(item=data['Item'],
                                             table_name=table_name)
                if isinstance(keys, Response):
                    return keys
                # fix response
                if response._content == '{}':
                    response._content = update_put_item_response_content(
                        data, response._content)
                    fix_headers_for_updated_response(response)
                if event_sources_or_streams_enabled:
                    existing_item = self._thread_local('existing_item')
                    # Get stream specifications details for the table
                    stream_spec = dynamodb_get_table_stream_specification(
                        table_name=table_name)
                    record[
                        'eventName'] = 'INSERT' if not existing_item else 'MODIFY'
                    # prepare record keys
                    record['dynamodb']['Keys'] = keys
                    record['dynamodb']['NewImage'] = data['Item']
                    record['dynamodb']['SizeBytes'] = len(
                        json.dumps(data['Item']))
                    record['eventID'] = short_uid()
                    if stream_spec:
                        record['dynamodb']['StreamViewType'] = stream_spec[
                            'StreamViewType']
                    if existing_item:
                        record['dynamodb']['OldImage'] = existing_item

        elif action in ('GetItem', 'Query'):
            if response.status_code == 200:
                content = json.loads(to_str(response.content))
                # make sure we append 'ConsumedCapacity', which is properly
                # returned by dynalite, but not by AWS's DynamoDBLocal
                if 'ConsumedCapacity' not in content and data.get(
                        'ReturnConsumedCapacity') in ['TOTAL', 'INDEXES']:
                    content['ConsumedCapacity'] = {
                        'TableName': table_name,
                        'CapacityUnits': 5,  # TODO hardcoded
                        'ReadCapacityUnits': 2,
                        'WriteCapacityUnits': 3
                    }
                    response._content = json.dumps(content)
                    fix_headers_for_updated_response(response)

        elif action == 'DeleteItem':
            if response.status_code == 200 and event_sources_or_streams_enabled:
                old_item = self._thread_local('existing_item')
                record['eventName'] = 'REMOVE'
                record['dynamodb']['Keys'] = data['Key']
                record['dynamodb']['OldImage'] = old_item

        elif action == 'CreateTable':
            if 'StreamSpecification' in data:
                if response.status_code == 200:
                    content = json.loads(to_str(response._content))
                    create_dynamodb_stream(
                        data,
                        content['TableDescription'].get('LatestStreamLabel'))

            event_publisher.fire_event(
                event_publisher.EVENT_DYNAMODB_CREATE_TABLE,
                payload={'n': event_publisher.get_hash(table_name)})

            if data.get('Tags') and response.status_code == 200:
                table_arn = json.loads(
                    response._content)['TableDescription']['TableArn']
                TABLE_TAGS[table_arn] = {
                    tag['Key']: tag['Value']
                    for tag in data['Tags']
                }

            return

        elif action == 'DeleteTable':
            if response.status_code == 200:
                table_arn = json.loads(response._content).get(
                    'TableDescription', {}).get('TableArn')
                event_publisher.fire_event(
                    event_publisher.EVENT_DYNAMODB_DELETE_TABLE,
                    payload={'n': event_publisher.get_hash(table_name)})
                self.delete_all_event_source_mappings(table_arn)
                dynamodbstreams_api.delete_streams(table_arn)
                TABLE_TAGS.pop(table_arn, None)
            return

        elif action == 'UpdateTable':
            if 'StreamSpecification' in data:
                if response.status_code == 200:
                    content = json.loads(to_str(response._content))
                    create_dynamodb_stream(
                        data,
                        content['TableDescription'].get('LatestStreamLabel'))
            return

        elif action == 'TagResource':
            table_arn = data['ResourceArn']
            if table_arn not in TABLE_TAGS:
                TABLE_TAGS[table_arn] = {}
            TABLE_TAGS[table_arn].update(
                {tag['Key']: tag['Value']
                 for tag in data.get('Tags', [])})
            return

        elif action == 'UntagResource':
            table_arn = data['ResourceArn']
            for tag_key in data.get('TagKeys', []):
                TABLE_TAGS.get(table_arn, {}).pop(tag_key, None)
            return

        else:
            # nothing to do
            return

        if event_sources_or_streams_enabled and records and 'eventName' in records[
                0]:
            if 'TableName' in data:
                records[0]['eventSourceARN'] = aws_stack.dynamodb_table_arn(
                    table_name)

            forward_to_lambda(records)
            records = self.prepare_records_to_forward_to_ddb_stream(records)
            forward_to_ddb_stream(records)
Esempio n. 2
0
    def return_response(self, method, path, data, headers, response):
        if path.startswith("/shell") or method == "GET":
            return

        data = json.loads(to_str(data))

        # update table definitions
        if data and "TableName" in data and "KeySchema" in data:
            table_definitions = DynamoDBRegion.get().table_definitions
            table_definitions[data["TableName"]] = data
        if response._content:
            # fix the table and latest stream ARNs (DynamoDBLocal hardcodes "ddblocal" as the region)
            content_replaced = re.sub(
                r'("TableArn"|"LatestStreamArn"|"StreamArn")\s*:\s*"arn:aws:dynamodb:ddblocal:([^"]+)"',
                r'\1: "arn:aws:dynamodb:%s:\2"' % aws_stack.get_region(),
                to_str(response._content),
            )
            if content_replaced != response._content:
                response._content = content_replaced
                fix_headers_for_updated_response(response)

        action = headers.get("X-Amz-Target", "")
        action = action.replace(ACTION_PREFIX, "")
        if not action:
            return
        # upgrade event version to 1.1
        record = {
            "eventID": "1",
            "eventVersion": "1.1",
            "dynamodb": {
                "ApproximateCreationDateTime": time.time(),
                # 'StreamViewType': 'NEW_AND_OLD_IMAGES',
                "SizeBytes": -1,
            },
            "awsRegion": aws_stack.get_region(),
            "eventSource": "aws:dynamodb",
        }
        records = [record]

        streams_enabled_cache = {}
        table_name = data.get("TableName")
        event_sources_or_streams_enabled = has_event_sources_or_streams_enabled(
            table_name, streams_enabled_cache)

        if action == "UpdateItem":
            if response.status_code == 200 and event_sources_or_streams_enabled:
                existing_item = self._thread_local("existing_item")
                record[
                    "eventName"] = "INSERT" if not existing_item else "MODIFY"
                record["eventID"] = short_uid()
                updated_item = find_existing_item(data)
                if not updated_item:
                    return
                record["dynamodb"]["Keys"] = data["Key"]
                if existing_item:
                    record["dynamodb"]["OldImage"] = existing_item
                record["dynamodb"]["NewImage"] = updated_item
                record["dynamodb"]["SizeBytes"] = len(json.dumps(updated_item))
                stream_spec = dynamodb_get_table_stream_specification(
                    table_name=table_name)
                if stream_spec:
                    record["dynamodb"]["StreamViewType"] = stream_spec[
                        "StreamViewType"]

        elif action == "BatchWriteItem":
            records, unprocessed_items = self.prepare_batch_write_item_records(
                record, data)
            for record in records:
                event_sources_or_streams_enabled = (
                    event_sources_or_streams_enabled
                    or has_event_sources_or_streams_enabled(
                        record["eventSourceARN"], streams_enabled_cache))
            if response.status_code == 200 and any(unprocessed_items):
                content = json.loads(to_str(response.content))
                table_name = list(data["RequestItems"].keys())[0]
                if table_name not in content["UnprocessedItems"]:
                    content["UnprocessedItems"][table_name] = []
                for key in ["PutRequest", "DeleteRequest"]:
                    if any(unprocessed_items[key]):
                        content["UnprocessedItems"][table_name].append(
                            {key: unprocessed_items[key]})
                unprocessed = content["UnprocessedItems"]
                for key in list(unprocessed.keys()):
                    if not unprocessed.get(key):
                        del unprocessed[key]

                response._content = json.dumps(content)
                fix_headers_for_updated_response(response)

        elif action == "TransactWriteItems":
            records = self.prepare_transact_write_item_records(record, data)
            for record in records:
                event_sources_or_streams_enabled = (
                    event_sources_or_streams_enabled
                    or has_event_sources_or_streams_enabled(
                        record["eventSourceARN"], streams_enabled_cache))

        elif action == "PutItem":
            if response.status_code == 200:
                keys = dynamodb_extract_keys(item=data["Item"],
                                             table_name=table_name)
                if isinstance(keys, Response):
                    return keys
                # fix response
                if response._content == "{}":
                    response._content = update_put_item_response_content(
                        data, response._content)
                    fix_headers_for_updated_response(response)
                if event_sources_or_streams_enabled:
                    existing_item = self._thread_local("existing_item")
                    # Get stream specifications details for the table
                    stream_spec = dynamodb_get_table_stream_specification(
                        table_name=table_name)
                    record[
                        "eventName"] = "INSERT" if not existing_item else "MODIFY"
                    # prepare record keys
                    record["dynamodb"]["Keys"] = keys
                    record["dynamodb"]["NewImage"] = data["Item"]
                    record["dynamodb"]["SizeBytes"] = len(
                        json.dumps(data["Item"]))
                    record["eventID"] = short_uid()
                    if stream_spec:
                        record["dynamodb"]["StreamViewType"] = stream_spec[
                            "StreamViewType"]
                    if existing_item:
                        record["dynamodb"]["OldImage"] = existing_item

        elif action in ("GetItem", "Query"):
            if response.status_code == 200:
                content = json.loads(to_str(response.content))
                # make sure we append 'ConsumedCapacity', which is properly
                # returned by dynalite, but not by AWS's DynamoDBLocal
                if "ConsumedCapacity" not in content and data.get(
                        "ReturnConsumedCapacity") in [
                            "TOTAL",
                            "INDEXES",
                        ]:
                    content["ConsumedCapacity"] = {
                        "TableName": table_name,
                        "CapacityUnits": 5,  # TODO hardcoded
                        "ReadCapacityUnits": 2,
                        "WriteCapacityUnits": 3,
                    }
                    response._content = json.dumps(content)
                    fix_headers_for_updated_response(response)

        elif action == "DeleteItem":
            if response.status_code == 200 and event_sources_or_streams_enabled:
                old_item = self._thread_local("existing_item")
                record["eventID"] = short_uid()
                record["eventName"] = "REMOVE"
                record["dynamodb"]["Keys"] = data["Key"]
                record["dynamodb"]["OldImage"] = old_item
                record["dynamodb"]["SizeBytes"] = len(json.dumps(old_item))
                # Get stream specifications details for the table
                stream_spec = dynamodb_get_table_stream_specification(
                    table_name=table_name)
                if stream_spec:
                    record["dynamodb"]["StreamViewType"] = stream_spec[
                        "StreamViewType"]

        elif action == "CreateTable":
            if "StreamSpecification" in data:
                if response.status_code == 200:
                    content = json.loads(to_str(response._content))
                    create_dynamodb_stream(
                        data,
                        content["TableDescription"].get("LatestStreamLabel"))

            event_publisher.fire_event(
                event_publisher.EVENT_DYNAMODB_CREATE_TABLE,
                payload={"n": event_publisher.get_hash(table_name)},
            )

            if data.get("Tags") and response.status_code == 200:
                table_arn = json.loads(
                    response._content)["TableDescription"]["TableArn"]
                DynamoDBRegion.TABLE_TAGS[table_arn] = {
                    tag["Key"]: tag["Value"]
                    for tag in data["Tags"]
                }

            return

        elif action == "DeleteTable":
            if response.status_code == 200:
                table_arn = (json.loads(response._content).get(
                    "TableDescription", {}).get("TableArn"))
                event_publisher.fire_event(
                    event_publisher.EVENT_DYNAMODB_DELETE_TABLE,
                    payload={"n": event_publisher.get_hash(table_name)},
                )
                self.delete_all_event_source_mappings(table_arn)
                dynamodbstreams_api.delete_streams(table_arn)
                DynamoDBRegion.TABLE_TAGS.pop(table_arn, None)
            return

        elif action == "UpdateTable":
            content_str = to_str(response._content or "")
            if response.status_code == 200 and "StreamSpecification" in data:
                content = json.loads(content_str)
                create_dynamodb_stream(
                    data, content["TableDescription"].get("LatestStreamLabel"))
            if (response.status_code >= 400 and data.get("ReplicaUpdates")
                    and "Nothing to update" in content_str):
                table_name = data.get("TableName")
                # update local table props (replicas)
                table_properties = DynamoDBRegion.get().table_properties
                table_properties[
                    table_name] = table_props = table_properties.get(
                        table_name) or {}
                table_props["Replicas"] = replicas = table_props.get(
                    "Replicas") or []
                for repl_update in data["ReplicaUpdates"]:
                    for key, details in repl_update.items():
                        region = details.get("RegionName")
                        if key == "Create":
                            details["ReplicaStatus"] = details.get(
                                "ReplicaStatus") or "ACTIVE"
                            replicas.append(details)
                        if key == "Update":
                            replica = [
                                r for r in replicas
                                if r.get("RegionName") == region
                            ]
                            if replica:
                                replica[0].update(details)
                        if key == "Delete":
                            table_props["Replicas"] = [
                                r for r in replicas
                                if r.get("RegionName") != region
                            ]
                # update response content
                schema = get_table_schema(table_name)
                result = {"TableDescription": schema["Table"]}
                update_response_content(response, json_safe(result), 200)
            return

        elif action == "DescribeTable":
            table_name = data.get("TableName")
            table_props = DynamoDBRegion.get().table_properties.get(table_name)
            if table_props:
                content = json.loads(to_str(response.content))
                content.get("Table", {}).update(table_props)
                update_response_content(response, content)

        elif action == "TagResource":
            table_arn = data["ResourceArn"]
            table_tags = DynamoDBRegion.TABLE_TAGS
            if table_arn not in table_tags:
                table_tags[table_arn] = {}
            table_tags[table_arn].update(
                {tag["Key"]: tag["Value"]
                 for tag in data.get("Tags", [])})
            return

        elif action == "UntagResource":
            table_arn = data["ResourceArn"]
            for tag_key in data.get("TagKeys", []):
                DynamoDBRegion.TABLE_TAGS.get(table_arn, {}).pop(tag_key, None)
            return

        else:
            # nothing to do
            return
        if event_sources_or_streams_enabled and records and "eventName" in records[
                0]:
            if "TableName" in data:
                records[0]["eventSourceARN"] = aws_stack.dynamodb_table_arn(
                    table_name)
            # forward to kinesis stream
            forward_to_kinesis_stream(records)
            # forward to lambda and ddb_streams
            forward_to_lambda(records)
            records = self.prepare_records_to_forward_to_ddb_stream(records)
            forward_to_ddb_stream(records)