Response Event Info, Step Function Meta Data (#341)

* Response Event Info, Step Function Meta Data

Closes #340

* Adding tests

* Add LambdaProxy response type

* Add test for http response collection

* Update README

.circleci/config.yml

@@ -47,9 +47,9 @@ jobs:
       - run:
           name: Check code style
           command: |
-            pip install black==18.6b2
-            black --check --line-length=88 --safe iopipe
-            black --check --line-length=88 --safe tests
+            pip install black==19.3b0
+            black iopipe
+            black tests
 
   coverage:
     working_directory: ~/iopipe-python

.pre-commit-config.yaml

@@ -1,8 +1,5 @@
-repos:
--   repo: https://github.com/ambv/black
-    rev: 18.6b2
-
+-   repo: https://github.com/psf/black
+    rev: stable
     hooks:
     - id: black
-      args: [--line-length=88, --safe]
-      python_version: python3.6
+      language_version: python3.7

README.md

@@ -23,6 +23,7 @@ This package provides analytics and distributed tracing for event-driven applica
   - [Labels](https://github.com/iopipe/iopipe-python#labels)
   - [Core Agent](https://github.com/iopipe/iopipe-python#core-agent)
   - [Disabling Reporting](https://github.com/iopipe/iopipe-python#disabling-reporting)
+  - [Step Functions](https://github.com/iopipe/iopipe-python#step-functions)
 - [Plugins](https://github.com/iopipe/iopipe-python#plugins)
   - [Event Info Plugin](https://github.com/iopipe/iopipe-python#event-info-plugin)
   - [Logger Plugin](https://github.com/iopipe/iopipe-python#logger-plugin)
@@ -205,6 +206,22 @@ def handler(event, context):
 
 Reporting will be re-enabled on the next invocation.
 
+### Step Functions
+
+IOpipe is compatible with AWS Lambda step functions. To enable step function tracing:
+
+```python
+from iopipe import IOpipe
+
+iopipe = IOpipe()
+
+@iopipe.step
+def handler(event, context):
+    pass
+```
+
+The `@iopipe.step` decorator will enable step function mode, which will collect additional meta data about your step functions.
+
 ## Plugins
 
 IOpipe's functionality can be extended through plugins. Plugins hook into the agent lifecycle to allow you to perform additional analytics.
@@ -497,6 +514,9 @@ class MyPlugin(Plugin):
     def post_invoke(self, event, context):
         pass
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         pass
 
@@ -519,12 +539,13 @@ A plugin has the following properties defined:
 
 A plugin has the following methods defined:
 
-- `pre_setup`: Is called once prior to the agent initialization. Is passed the `iopipe` instance.
-- `post_setup`: Is called once after the agent is initialized, is passed the `iopipe` instance.
-- `pre_invoke`: Is called prior to each invocation, is passed the `event` and `context` of the invocation.
-- `post_invoke`: Is called after each invocation, is passed the `event` and `context` of the invocation.
-- `pre_report`: Is called prior to each report being sent, is passed the `report` instance.
-- `post_report`: Is called after each report is sent, is passed the `report` instance.
+- `pre_setup`: Is called once prior to the agent initialization; is passed the `iopipe` instance.
+- `post_setup`: Is called once after the agent is initialized; is passed the `iopipe` instance.
+- `pre_invoke`: Is called prior to each invocation; is passed the `event` and `context` of the invocation.
+- `post_invoke`: Is called after each invocation; is passed the `event` and `context` of the invocation.
+- `post_response`: Is called after the invocation response; is passed the `response`value.
+- `pre_report`: Is called prior to each report being sent; is passed the `report` instance.
+- `post_report`: Is called after each report is sent; is passed the `report` instance.
 
 ## Supported Python Versions
 

iopipe/agent.py

@@ -84,7 +84,7 @@ def error(self, error):
 
     err = error
 
-    def __call__(self, func):
+    def __call__(self, func, **kwargs):
         @functools.wraps(func)
         def wrapped(event, context):
             # Skip if function is already instrumented
@@ -96,7 +96,7 @@ def wrapped(event, context):
 
             logger.debug("Wrapping %s with IOpipe decorator" % repr(func))
 
-            self.context = context = ContextWrapper(context, self)
+            self.context = context = ContextWrapper(context, self, **kwargs)
 
             # if env var IOPIPE_ENABLED is set to False skip reporting
             if self.config["enabled"] is False:
@@ -197,6 +197,9 @@ def wrapped(event, context):
 
     decorator = __call__
 
+    def step(self, func):
+        return self(func, step_function=True)
+
     def load_plugins(self, plugins):
         """
         Loads plugins that match the `Plugin` interface and are instantiated.

iopipe/context.py

@@ -1,6 +1,7 @@
 import decimal
 import logging
 import numbers
+import uuid
 import warnings
 
 from . import constants
@@ -25,20 +26,22 @@ def __getattr__(self, name):
 
 
 class ContextWrapper(object):
-    def __init__(self, base_context, instance):
+    def __init__(self, base_context, instance, **kwargs):
         self.base_context = base_context
         self.instance = instance
-        self.iopipe = IOpipeContext(self.instance)
+        self.iopipe = IOpipeContext(self.instance, **kwargs)
 
     def __getattr__(self, name):
         return getattr(self.base_context, name)
 
 
 class IOpipeContext(object):
-    def __init__(self, instance):
+    def __init__(self, instance, **kwargs):
         self.instance = instance
         self.log = LogWrapper(self)
         self.disabled = False
+        self.is_step_function = kwargs.pop("step_function", False)
+        self.step_meta = None
 
     def metric(self, key, value):
         if self.instance.report is None:
@@ -129,3 +132,13 @@ def unregister(self, name):
 
     def disable(self):
         self.disabled = True
+
+    def collect_step_meta(self, event):
+        if self.is_step_function:
+            self.step_meta = event.get("iopipe", {"id": str(uuid.uuid4()), "step": 0})
+
+    def inject_step_meta(self, response):
+        if self.step_meta and isinstance(response, dict):
+            step_meta = self.step_meta.copy()
+            step_meta["step"] += 1
+            response["iopipe"] = step_meta

iopipe/contrib/eventinfo/event_types.py

@@ -1,14 +1,19 @@
+from .response_types import LambdaProxy
 from .util import collect_all_keys, get_value, has_key, slugify
 
 
 class EventType(object):
     keys = []
+    response_keys = []
     exclude_keys = []
     required_keys = []
+    response_type = None
     source = None
 
     def __init__(self, event):
         self.event = event
+        if self.response_type:
+            self.response_type = self.response_type(self.type)
 
     @property
     def slug(self):
@@ -83,6 +88,7 @@ class ApiGateway(EventType):
         "resource",
     ]
     required_keys = ["headers", "httpMethod", "path", "requestContext", "resource"]
+    response_type = LambdaProxy
 
 
 class CloudFront(EventType):
@@ -172,6 +178,7 @@ class ServerlessLambda(EventType):
         ("stage", "requestContext.stage"),
     ]
     required_keys = ["identity.userAgent", "identity.sourceIp", "identity.accountId"]
+    response_type = LambdaProxy
 
 
 class SES(EventType):
@@ -266,3 +273,17 @@ def metrics_for_event_type(event, context):
             event_info = event_type.collect()
             [context.iopipe.metric(k, v) for k, v in event_info.items()]
             break
+
+    if context.iopipe.is_step_function:
+        context.iopipe.collect_step_meta(event)
+        if context.iopipe.step_meta:
+            for key, value in context.iopipe.step_meta.items():
+                context.iopipe.metric("@iopipe/event-info.stepFunction.%s" % key, value)
+
+    return event_type
+
+
+def metrics_for_response_type(event_type, context, response):
+    if event_type.response_type:
+        response_info = event_type.response_type.collect(response)
+        [context.iopipe.metric(k, v) for k, v in response_info.items()]

iopipe/contrib/eventinfo/plugin.py

@@ -3,7 +3,7 @@
 
 from iopipe.plugins import Plugin
 
-from .event_types import metrics_for_event_type
+from .event_types import metrics_for_event_type, metrics_for_response_type
 
 
 class EventInfoPlugin(Plugin):
@@ -35,14 +35,16 @@ def post_setup(self, iopipe):
         pass
 
     def pre_invoke(self, event, context):
-        pass
+        self.context = context
 
     def post_invoke(self, event, context):
         if self.enabled:
-            metrics_for_event_type(event, context)
+            self.event_type = metrics_for_event_type(event, context)
 
     def post_response(self, response):
-        pass
+        if self.enabled:
+            metrics_for_response_type(self.event_type, self.context, response)
+            self.context.iopipe.inject_step_meta(response)
 
     def pre_report(self, report):
         pass

iopipe/contrib/eventinfo/response_types.py

@@ -0,0 +1,28 @@
+from .util import get_value
+
+
+class ResponseType(object):
+    keys = []
+
+    def __init__(self, event_type):
+        self.event_type = event_type
+
+    def collect(self, response):
+        response_info = {}
+
+        for key in self.keys:
+            if isinstance(key, tuple):
+                old_key, new_key = key
+            else:
+                old_key = new_key = key
+            value = get_value(response, old_key)
+            if value is not None:
+                response_info[
+                    "@iopipe/event-info.%s.response.%s" % (self.event_type, new_key)
+                ] = value
+
+        return response_info
+
+
+class LambdaProxy(ResponseType):
+    keys = ["statusCode"]

pyproject.toml

@@ -0,0 +1,18 @@
+[tool.black]
+line-length = 88
+target-version = ['py27', 'py36', 'py37', 'py38']
+include = '\.pyi?$'
+exclude = '''
+/(
+    \.eggs
+  | \.git
+  | \.hg
+  | \.mypy_cache
+  | \.tox
+  | \.venv
+  | _build
+  | buck-out
+  | build
+  | dist
+)/
+'''

setup.py

@@ -39,7 +39,7 @@
     packages=find_packages(exclude=("tests", "tests.*")),
     extras_require={
         "coverage": coverage_requires,
-        "dev": tests_require + ["black==18.6b2", "pre-commit"],
+        "dev": tests_require + ["black==19.3b0", "pre-commit"],
     },
     install_requires=install_requires,
     setup_requires=["pytest-runner==4.2"],

tests/conftest.py

@@ -208,3 +208,16 @@ def _handler_that_disables_reporting_with_error(event, context):
         raise Exception("An error happened")
 
     return iopipe, _handler_that_disables_reporting_with_error
+
+
+@pytest.fixture
+def handler_step_function(iopipe):
+    @iopipe
+    def _handler_not_step_function(event, context):
+        assert context.iopipe.is_step_function is False
+
+    @iopipe.step
+    def _handler_step_function(event, context):
+        assert context.iopipe.is_step_function is True
+
+    return iopipe, _handler_step_function

tests/contrib/eventinfo/conftest.py

@@ -26,6 +26,25 @@ def _handler(event, context):
     return iopipe_with_eventinfo, _handler
 
 
+@pytest.fixture
+def handler_step_function_with_eventinfo(iopipe_with_eventinfo):
+    @iopipe_with_eventinfo.step
+    def _handler(event, context):
+        assert context.iopipe.is_step_function is True
+        return {}
+
+    return iopipe_with_eventinfo, _handler
+
+
+@pytest.fixture
+def handler_http_response_with_eventinfo(iopipe_with_eventinfo):
+    @iopipe_with_eventinfo
+    def _handler(event, context):
+        return {"statusCode": 200, "body": "success"}
+
+    return iopipe_with_eventinfo, _handler
+
+
 def _load_event(name):
     json_file = os.path.join(os.path.dirname(__file__), "events", "%s.json" % name)
     with open(json_file) as f:

tests/contrib/eventinfo/test_plugin.py

@@ -112,3 +112,52 @@ def test__eventinfo_plugin__enabled(monkeypatch):
     plugin = EventInfoPlugin(enabled=False)
 
     assert plugin.enabled is True
+
+
+@mock.patch("iopipe.report.send_report", autospec=True)
+def test__eventinfo_plugin__step_function(
+    mock_send_report, handler_step_function_with_eventinfo, event_apigw, mock_context
+):
+    iopipe, handler = handler_step_function_with_eventinfo
+
+    plugins = iopipe.config["plugins"]
+    assert len(plugins) == 1
+    assert plugins[0].enabled is True
+    assert plugins[0].name == "event-info"
+
+    response1 = handler(event_apigw, mock_context)
+    assert "iopipe" in response1
+    assert "id" in response1["iopipe"]
+    assert "step" in response1["iopipe"]
+
+    response2 = handler(response1, mock_context)
+
+    assert "iopipe" in response2
+    assert response1["iopipe"]["id"] == response2["iopipe"]["id"]
+    assert response2["iopipe"]["step"] > response1["iopipe"]["step"]
+
+
+@mock.patch("iopipe.report.send_report", autospec=True)
+def test__eventinfo_plugin__http_response(
+    mock_send_report, handler_http_response_with_eventinfo, event_apigw, mock_context
+):
+    iopipe, handler = handler_http_response_with_eventinfo
+
+    handler(event_apigw, mock_context)
+    metrics = iopipe.report.custom_metrics
+
+    assert any(
+        (
+            m["name"] == "@iopipe/event-info.apiGateway.response.statusCode"
+            for m in metrics
+        )
+    )
+
+    metric = next(
+        (
+            m
+            for m in metrics
+            if m["name"] == "@iopipe/event-info.apiGateway.response.statusCode"
+        )
+    )
+    assert metric["n"] == 200