Creating Custom Framework Integrations with the Elastic APM Python Agent

By
Benjamin Wohlwend

As you might know, the Elastic APM Python agent currently has built-in support for two web frameworks: Django and Flask. But the ecosystem of Python web frameworks is teeming with dozens of frameworks. From the venerable elder statesmen like Zope, to the aforementioned big players like Django and Flask, and the new kids on the block like Falcon or Sanic, there's something for everybody.

While this rich assortment of options is a great sign for the health of the Python web ecosystem, it comes with a drawback for products like Elastic APM: it is virtually impossible to offer built-in support for all of them.

But I'm about to let you in on a little known secret: building your own custom integration isn't rocket science and can be done in less than 100 lines of code!

The anatomy of a framework integration

For an integration to be on par with our existing framework support like Django and Flask, there's a few things it should do:

  1. Call elasticapm.instrument() as early as possible during the startup of the framework.
  2. Instantiate an elasticapm.Client object.
  3. Begin and end transactions at the appropriate times during request handling.
  4. Figure out a name and "result" for the transaction (e.g. the route name for the request, and the HTTP status code as the result).
  5. Gather contextual information from the request and response objects.
  6. Hook into the exception handling of the framework, so we can report any application errors to the Elastic APM Server.

Building the integration

To make this a bit less abstract, let's randomly choose a web framework and build an integration for it. How about Pyramid? Cool? Cool.

Pyramid, like many other Python web frameworks, is built on WSGI and supports WSGI middleware. A WSGI middleware is a wrapper around the app, and as such it sees any request going into the app, and the response coming out of the app. This is great for timing the request, but it has the drawback of being somewhat generic. For deep framework integration, we try to find something more "native" to the framework. As it happens, Pyramid has a concept called "tweens" that looks promising. To quote from the Pyramid docs:

Tweens behave a bit like WSGI middleware, but they have the benefit of running in a context in which they have access to the Pyramid request, response, and application registry, as well as the Pyramid rendering machinery.

That sounds exactly like what we need! As a starting point, we'll take the class-based tween directly from the documentation:

class simple_tween_factory(object):
   def __init__(self, handler, registry):
        self.handler = handler
        self.registry = registry

        # one-time configuration code goes here

    def __call__(self, request):
        # code to be executed for each request before
        # the actual application code goes here

        response = self.handler(request)

        # code to be executed for each request after
        # the actual application code goes here

        return response

Checking our TODO list from above, it looks like step 2 can be done in __init__(), and steps 3 to 6 can be done in __call__. Let's start with step 1 though.

The call to elasticapm.instrument() has to happen as early as possible during startup of the app. The instrument() call instruments a lot of different modules, so we can trace SQL queries, external HTTP requests, and other shenanigans your code might be up to.

Many frameworks offer hooks which can be used to execute code during startup. For Pyramid, we can subscribe to the ApplicationCreated event:

from pyramid.events import ApplicationCreated, subscriber

import elasticapm

@subscriber(ApplicationCreated)
def elasticapm_instrument(event):
    elasticapm.instrument()

That looks about right. Back to our tween. We need to initialize an elasticapm.Client object and store it for later use on the tween. We simply instantiate the Client here, relying on any configuration being set via environment variables. A more feature-complete integration could try and get framework-specific configuration as well (in the case of Pyramid, this would probably be an .ini file).

import pkg_resources
import elasticapm

class elasticapm_tween_factory(object):
    def __init__(self):
        self.handler = handler
        self.registry = registry

        self.client = elasticapm.Client(
            framework_name="Pyramid",
            framework_version=pkg_resources.get_distribution("pyramid").version
        )

We also set the framework name and version here. This is helpful information to have on errors and transactions.

Step 3 is quick work as well: we wrap the call to the handler with the appropriate API calls to start and end the transaction:

def __call__(self, request):
    self.client.begin_transaction('request')
    response = self.handler(request)
    self.client.end_transaction()
    return response

Now, if you used that code as is, all transactions would be grouped together in the APM UI. That's because we didn't provide a name and result for the transaction. Conveniently, that's step 4, which comes next!

In a web framework, the convention is to use the parametrized route (e.g. /users/{id}) as the transaction name, and the HTTP status class as the result ("status class" in this case means that we bunch all 2xx requests together, all 3xx etc.). Avoid to use the full URL as the transaction name, as this would create a new entry in the transaction list for every variation of the same route (e.g. /users/1, /users/2, etc.).

Consulting the relevant documentation, it looks like Pyramid provides the route pattern on the request object. Unsurprisingly, the HTTP status code is provided on the response object. Looks like we're all set for step 4!

def __call__(self, request):
    self.client.begin_transaction('request')
    response = self.handler(request)
    transaction_name = request.matched_route.pattern if request.matched_route else request.view_name
    # prepend request method
    transaction_name = " ".join((request.method, transaction_name)) if transaction_name else ""
    transaction_result = response.status[0] + "xx"
    self.client.end_transaction(transaction_name, transaction_result)
    return response

Sweet!

Readers who haven't nodded off quite yet will have noticed that we're falling back to request.view_name if request.matched_route is not set. Unfortunately, the real world is messy and we have to cope with that.

We are also prepending the request method to the transaction name. This ensures that for example GET /users/{id} and DELETE /users/{id} get their own entry in the transaction list, as the code paths used by different HTTP methods can vary greatly.

This brings us to step 5: contextual information. This can include information from the request (full URL, headers, cookies), response (full status code, content type), the logged in user (e.g. the user ID) and more. In this blog post, we'll handle request and response as an example.

We will put the code to gather the request / response data into reusable functions, because, SPOILER ALERT, we will be using them again in step 6.

from elasticapm.utils import compat, get_url_dict

def get_data_from_request(request):
    data = {
        "headers": dict(**request.headers),
        "method": request.method,
        "socket": {
            "remote_address": request.remote_addr,
            "encrypted": request.scheme == 'https'
        },
        "cookies": dict(**request.cookies),
        "url": get_url_dict(request.url)
    }
    # remove Cookie header since the same data is in request["cookies"] as well
    data["headers"].pop("Cookie", None)
    return data


def get_data_from_response(response):
    data = {"status_code": response.status_int}
    if response.headers:
        data["headers"] = {
            key: ";".join(response.headers.getall(key))
            for key in compat.iterkeys(response.headers)
        }
    return data

You might have noticed that we need to do some special handling for cookies and headers. This is because request.cookies, request.headers and response.headers might look like dictionaries, but in fact are special-purpose dict-like objects. If we just passed those objects along, our JSON serializer would throw a tantrum when trying to send the data to the APM Server.

Having done that, we need to call these functions before ending the transaction, and use elasticapm.set_context() to attach it to the current transaction. To ensure that we only do all this work if the transaction is actually sampled, set_context() accepts a callable, which it will only call if the transaction is sampled. Let's modify __call__ a bit by adding the necessary calls before self.client.end_transaction():

elasticapm.set_context(lambda: get_data_from_request(request), "request")
elasticapm.set_context(lambda: get_data_from_response(response), "response")
self.client.end_transaction(transaction_name, transaction_result)
return response

And there we have it! A Pyramid tween to instrument and measure HTTP requests, all in a couple dozen lines of code. If we lived in a perfect world and always wrote code without bugs, we would be done now. Alas…

Capturing exceptions

Capturing exceptions is an important part of APM. Your app might be super quick because half the requests raise exceptions, and you'll never be the wiser if you don't monitor it.

The good news is that most Python web frameworks have some kind of hook or callback you can register for getting notified of exceptions. In the case of Pyramid, we can wrap the call to the handler in a try/except block, and capture the exception (if any). A bit simplified, it looks like this:

import sys
from pyramid.compat import reraise

#...
# in __call__():

try:
    response = self.handler(request)
    return response
except Exception:
    self.client.capture_exception(
        context={
                "request": get_data_from_request(request)
        },
        handled=False,  # indicate that this exception bubbled all the way up to the user
    )
    reraise(*sys.exc_info())
finally:
    self.client.end_transaction(transaction_name, transaction_result)

And that's it! We're done! High fives for everybody, and let's call it a day.

Step 7

Oh. What now? Right… we need to hook all of this up with our app! This part is very framework-specific, and there is often more than one way to do it. Looking at Pyramid, one option seems particularly nice: config.include. Using include, we can define an includeme function in our "integration" module, and then simply call config.include("elasticapm_integration") in the app config. Our includeme does two things: register the tween, and advising Pyramid to scan our module, which in turn should pick up our ApplicationCreated subscriber from way back in step 1.

def includeme(config):
    config.add_tween('elasticapm_integration.elasticapm_tween_factory')
    config.scan('elasticapm_integration')

Now that's really it! I hope this example illustrates that creating a framework integration is no black magic. If you want to give it a try with your favorite framework and hit any troubles, come talk to us in our forum. We also run a survey for the Python agent, in which you can tell us which framework we should add next to our list of officially supported frameworks.

You can find the complete code of this example integration on GitHub. That repository also includes a sample TODO-list application, originating from the Pyramid Community Cookbook.