AWS Lambda Integration (#1107)

* Initial lambda integration

* feature(lambda): Add lambda tests and client

* Add D107 to ignore flake8 docs

* Initial lambda integration

* feature(lambda): Add lambda tests and client

* Add reference to raven-python-lambda

7 files changed, 417 insertions, 1 deletions

diff --git a/.travis.yml b/.travis.yml
index 87fbf99..f5c28e3 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -115,6 +115,14 @@ jobs:
       python: 2.7
       env: TOXENV=py27-celery-4
 
+    - stage: contrib
+      python: 2.7
+      env: TOXENV=py27-lambda
+
+    - stage: contrib
+      python: 3.6
+      env: TOXENV=py36-lambda
+
 #    - stage: deploy
 #      script: skip
 #      deploy:
diff --git a/conftest.py b/conftest.py
index 53a8b88..70d582d 100644
--- a/conftest.py
+++ b/conftest.py
@@ -4,7 +4,10 @@ import os.path
 import pytest
 import sys
 
-collect_ignore = []
+collect_ignore = [
+    'tests/contrib/awslambda'
+]
+
 if sys.version_info[0] > 2:
     if sys.version_info[1] < 3:
         collect_ignore.append('tests/contrib/flask')
diff --git a/docs/integrations/awslambda.rst b/docs/integrations/awslambda.rst
new file mode 100644
index 0000000..07ac0db
--- /dev/null
+++ b/docs/integrations/awslambda.rst
@@ -0,0 +1,54 @@
+Amazon Web Services Lambda
+==========================
+
+.. default-domain:: py
+
+
+
+Installation
+------------
+
+To use `Sentry`_ with `AWS Lambda`_, you have to install `raven` as an external
+dependency. This involves creating a `Deployment package`_ and uploading it
+to AWS.
+
+To install raven into your current project directory:
+
+.. code-block:: console
+
+    pip install raven -t /path/to/project-dir
+
+Setup
+-----
+
+Create a `LambdaClient` instance and wrap your lambda handler with
+the `capture_exeptions` decorator:
+
+
+.. sourcecode:: python
+
+    from raven.contrib.awslambda import LambdaClient
+
+
+    client = LambdaClient()
+
+    @client.capture_exceptions
+    def handler(event, context):
+        ...
+        raise Exception('I will be sent to sentry!')
+
+
+By default this will report unhandled exceptions and errors to Sentry.
+
+Additional settings for the client are configured using environment variables or
+subclassing `LambdaClient`.
+
+
+The integration was inspired by `raven python lambda`_, another implementation that
+also integrates with Serverless Framework and has SQS transport support. 
+
+
+.. _Sentry: https://getsentry.com/
+.. _AWS Lambda: https://aws.amazon.com/lambda
+.. _Deployment package: https://docs.aws.amazon.com/lambda/latest/dg/lambda-python-how-to-create-deployment-package.html
+.. _raven python lambda: https://github.com/Netflix-Skunkworks/raven-python-lambda
diff --git a/raven/contrib/awslambda/__init__.py b/raven/contrib/awslambda/__init__.py
new file mode 100644
index 0000000..7d6f2b8
--- /dev/null
+++ b/raven/contrib/awslambda/__init__.py
@@ -0,0 +1,173 @@
+"""
+raven.contrib.awslambda
+~~~~~~~~~~~~~~~~~~~~
+
+Raven wrapper for AWS Lambda handlers.
+
+:copyright: (c) 2010-2012 by the Sentry Team, see AUTHORS for more details.
+:license: BSD, see LICENSE for more details.
+"""
+# flake8: noqa
+
+from __future__ import absolute_import
+
+import os
+import logging
+import functools
+from types import FunctionType
+
+from raven.base import Client
+from raven.transport.http import HTTPTransport
+
+logger = logging.getLogger('sentry.errors.client')
+
+
+def get_default_tags():
+    return {
+        'lambda': 'AWS_LAMBDA_FUNCTION_NAME',
+        'version': 'AWS_LAMBDA_FUNCTION_VERSION',
+        'memory_size': 'AWS_LAMBDA_FUNCTION_MEMORY_SIZE',
+        'log_group': 'AWS_LAMBDA_LOG_GROUP_NAME',
+        'log_stream': 'AWS_LAMBDA_LOG_STREAM_NAME',
+        'region': 'AWS_REGION'
+    }
+
+
+class LambdaClient(Client):
+    """
+    Raven decorator for AWS Lambda.
+
+    By default, the lambda integration will capture unhandled exceptions and instrument logging.
+
+    Usage:
+
+    >>> from raven.contrib.awslambda import LambdaClient
+    >>>
+    >>>
+    >>> client = LambdaClient()
+    >>>
+    >>> @client.capture_exceptions
+    >>> def handler(event, context):
+    >>>    ...
+    >>>    raise Exception('I will be sent to sentry!')
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        transport = kwargs.get('transport', HTTPTransport)
+        super(LambdaClient, self).__init__(*args, transport=transport, **kwargs)
+
+    def capture(self, *args, **kwargs):
+        if 'data' not in kwargs:
+            kwargs['data'] = data = {}
+        else:
+            data = kwargs['data']
+        event = kwargs.get('event', None)
+        context = kwargs.get('context', None)
+        user_info = self._get_user_interface(event)
+        if user_info:
+            data.update(user_info)
+        if event:
+            http_info = self._get_http_interface(event)
+            if http_info:
+                data.update(http_info)
+            data['extra'] = self._get_extra_data(event, context)
+        return super(LambdaClient, self).capture(*args, **kwargs)
+
+    def build_msg(self, *args, **kwargs):
+
+        data = super(LambdaClient, self).build_msg(*args, **kwargs)
+        for option, default in get_default_tags().items():
+            data['tags'].setdefault(option, os.environ.get(default))
+        data.setdefault('release', os.environ.get('SENTRY_RELEASE'))
+        data.setdefault('environment', os.environ.get('SENTRY_ENVIRONMENT'))
+        return data
+
+    def capture_exceptions(self, f=None, exceptions=None):  # TODO: Ash fix kwargs in base
+        """
+        Wrap a function or code block in try/except and automatically call
+        ``.captureException`` if it raises an exception, then the exception
+        is reraised.
+
+        By default, it will capture ``Exception``
+
+        >>> @client.capture_exceptions
+        >>> def foo():
+        >>>     raise Exception()
+
+        >>> with client.capture_exceptions():
+        >>>    raise Exception()
+
+        You can also specify exceptions to be caught specifically
+
+        >>> @client.capture_exceptions((IOError, LookupError))
+        >>> def bar():
+        >>>     ...
+
+        ``kwargs`` are passed through to ``.captureException``.
+        """
+        if not isinstance(f, FunctionType):
+            # when the decorator has args which is not a function we except
+            # f to be the exceptions tuple
+            return functools.partial(self.capture_exceptions, exceptions=f)
+
+        exceptions = exceptions or (Exception,)
+
+        @functools.wraps(f)
+        def wrapped(event, context, *args, **kwargs):
+            try:
+                return f(event, context, *args, **kwargs)
+            except exceptions:
+                self.captureException(event=event, context=context, **kwargs)
+                self.context.clear()
+                raise
+        return wrapped
+
+    @staticmethod
+    def _get_user_interface(event):
+        if event.get('requestContext'):
+            identity = event['requestContext']['identity']
+            if identity:
+                user = {
+                    'id': identity.get('cognitoIdentityId', None) or identity.get('user', None),
+                    'username': identity.get('user', None),
+                    'ip_address': identity.get('sourceIp', None),
+                    'cognito_identity_pool_id': identity.get('cognitoIdentityPoolId', None),
+                    'cognito_authentication_type': identity.get('cognitoAuthenticationType', None),
+                    'user_agent': identity.get('userAgent')
+                }
+                return {'user': user}
+
+    @staticmethod
+    def _get_http_interface(event):
+        if event.get('path') and event.get('httpMethod'):
+            request = {
+                "url": event.get('path'),
+                "method": event.get('httpMethod'),
+                "query_string": event.get('queryStringParameters', None),
+                "headers": event.get('headers', None) or [],
+            }
+            return {'request': request}
+
+    @staticmethod
+    def _get_extra_data(event, context):
+        extra_context = {
+            'event': event,
+            'aws_request_id': context.aws_request_id,
+            'context': vars(context),
+        }
+
+        if context.client_context:
+            extra_context['client_context'] = {
+                'client.installation_id': context.client_context.client.installation_id,
+                'client.app_title': context.client_context.client.app_title,
+                'client.app_version_name': context.client_context.client.app_version_name,
+                'client.app_version_code': context.client_context.client.app_version_code,
+                'client.app_package_name': context.client_context.client.app_package_name,
+                'custom': context.client_context.custom,
+                'env': context.client_context.env,
+            }
+        return extra_context
+
+
+
diff --git a/tests/contrib/awslambda/conftest.py b/tests/contrib/awslambda/conftest.py
new file mode 100644
index 0000000..c1a4812
--- /dev/null
+++ b/tests/contrib/awslambda/conftest.py
@@ -0,0 +1,106 @@
+
+import pytest
+from raven.contrib.awslambda import LambdaClient
+import uuid
+import time
+
+
+class MockClient(LambdaClient):
+    def __init__(self, *args, **kwargs):
+        self.events = []
+        super(MockClient, self).__init__(*args, **kwargs)
+
+    def send(self, **kwargs):
+        self.events.append(kwargs)
+
+    def is_enabled(self, **kwargs):
+        return True
+
+
+class LambdaIndentityStub(object):
+    def __init__(self, id=1, pool_id=1):
+        self.cognito_identity_id = id
+        self.cognito_identity_pool_id = pool_id
+
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+    def get(self, name, default=None):
+        return getattr(self, name, default)
+
+
+class LambdaContextStub(object):
+
+    def __init__(self, function_name, memory_limit_in_mb=128, timeout=300, function_version='$LATEST'):
+        self.function_name = function_name
+        self.memory_limit_in_mb = memory_limit_in_mb
+        self.timeout = timeout
+        self.function_version = function_version
+        self.timeout = timeout
+        self.invoked_function_arn = 'invoked_function_arn'
+        self.log_group_name = 'log_group_name'
+        self.log_stream_name = 'log_stream_name'
+        self.identity = LambdaIndentityStub(id=0, pool_id=0)
+        self.client_context = None
+        self.aws_request_id = str(uuid.uuid4())
+        self.start_time = time.time() * 1000
+
+    def __getitem__(self, item, default):
+        return getattr(self, item, default)
+
+    def get(self, name, default=None):
+        return getattr(self, name, default)
+
+    def get_remaining_time_in_millis(self):
+        return max(self.timeout * 1000 - int((time.time() * 1000) - self.start_time), 0)
+
+
+class LambdaEventStub(object):
+    def __init__(self, body=None, headers=None, http_method='GET', path='/test', query_string=None):
+        self.body = body
+        self.headers = headers
+        self.httpMethod = http_method
+        self.isBase64Encoded = False
+        self.path = path
+        self.queryStringParameters = query_string
+        self.resource = path
+        self.stageVariables = None
+        self.requestContext = {
+            'accountId': '0000000',
+            'apiId': 'AAAAAAAA',
+            'httpMethod': http_method,
+            'identity': LambdaIndentityStub(),
+            'path': path,
+            'requestId': 'test-request',
+            'resourceId': 'bbzeyv',
+            'resourcePath': '/test',
+            'stage': 'test-stage'
+        }
+
+    def __getitem__(self, name):
+        return getattr(self, name)
+
+    def get(self, name, default=None):
+        return getattr(self, name, default)
+
+
+@pytest.fixture
+def lambda_env(monkeypatch):
+    monkeypatch.setenv('SENTRY_DSN', 'http://public:secret@example.com/1')
+    monkeypatch.setenv('AWS_LAMBDA_FUNCTION_NAME', 'test_func')
+    monkeypatch.setenv('AWS_LAMBDA_FUNCTION_VERSION', '$LATEST')
+    monkeypatch.setenv('SENTRY_RELEASE', '$LATEST')
+    monkeypatch.setenv('SENTRY_ENVIRONMENT', 'testing')
+
+
+@pytest.fixture
+def mock_client():
+    return MockClient
+
+@pytest.fixture
+def lambda_event():
+    return LambdaEventStub
+
+@pytest.fixture
+def lambda_context():
+    return LambdaContextStub
diff --git a/tests/contrib/awslambda/test_lambda.py b/tests/contrib/awslambda/test_lambda.py
new file mode 100644
index 0000000..39cfc95
--- /dev/null
+++ b/tests/contrib/awslambda/test_lambda.py
@@ -0,0 +1,69 @@
+import pytest
+from raven.transport.http import HTTPTransport
+
+
+class MyException(Exception):
+    pass
+
+
+def test_decorator_exception(lambda_env, mock_client, lambda_event, lambda_context):
+
+    client = mock_client()
+
+    @client.capture_exceptions
+    def test_func(event, context):
+        raise MyException('There was an error.')
+
+    with pytest.raises(MyException):
+        test_func(event=lambda_event(), context=lambda_context(function_name='test_func'))
+
+    assert client.events
+    assert isinstance(client.remote.get_transport(), HTTPTransport)
+    assert 'user' in client.events[0].keys()
+    assert 'request' in client.events[0].keys()
+
+
+def test_decorator_with_args(lambda_env, mock_client, lambda_event, lambda_context):
+    client = mock_client()
+
+    @client.capture_exceptions((MyException,))
+    def test_func(event, context):
+        raise Exception
+
+    with pytest.raises(Exception):
+        test_func(event=lambda_event(), context=lambda_context(function_name='test_func'))
+
+    assert not client.events
+
+    @client.capture_exceptions((MyException,))
+    def test_func(event, context):
+        raise MyException
+
+    with pytest.raises(Exception):
+        test_func(event=lambda_event(), context=lambda_context(function_name='test_func'))
+
+    assert client.events
+
+
+def test_decorator_without_exceptions(lambda_env, mock_client, lambda_event, lambda_context):
+    client = mock_client()
+
+    @client.capture_exceptions((MyException,))
+    def test_func(event, context):
+        return 0
+
+    assert test_func(event=lambda_event(), context=lambda_context(function_name='test_func')) == 0
+
+
+def test_decorator_without_kwargs(lambda_env, mock_client, lambda_event, lambda_context):
+
+    client = mock_client()
+
+    @client.capture_exceptions((MyException,))
+    def test_func(event, context):
+        raise MyException
+
+    with pytest.raises(Exception):
+        test_func(lambda_event(), lambda_context(function_name='test_func'))
+
+    assert client.events
diff --git a/tox.ini b/tox.ini
index 4dbda27..ddb68ff 100644
--- a/tox.ini
+++ b/tox.ini
@@ -20,6 +20,7 @@ envlist =
     py35-flask-12
     py35-flask-dev
     py27-celery-{3,4}
+    py{27-36}-lambda
 
 
 [testenv]
@@ -49,6 +50,7 @@ setenv =
     django: TESTPATH=tests/contrib/django
     flask: TESTPATH=tests/contrib/flask
     celery: TESTPATH=tests/contrib/test_celery.py
+    lambda: TESTPATH=tests/contrib/awslambda
 usedevelop = true
 extras = tests
 
@@ -58,6 +60,7 @@ basepython =
     py33: python3.3
     py34: python3.4
     py35: python3.5
+    py36: python3.6
     pypy: pypy
 
 commands =