Instrument a Python function with New Relic being run in multiprocessing or subprocess

instrument.py

"""
This utility lets you instrument a function with New Relic even when it's in a third party library and called as a subprocess.

This is necessary because New Relic refuses to send data when the agent is initialized in a parent PID.
If the subprocess function is directly part of your code, you can simply wrap your code in something that initializes
and shuts down the agent.

Usage:

from instrument import subprocess_wrapper

subprocess_wrapper("path.to.third.party.module", "function_to_wrap")
"""
import newrelic.agent
from newrelic.agent import FunctionWrapper, wrap_background_task, wrap_object


def agent_wrapper(wrapped):
    def wrapper(wrapped, instance, args, kwargs):
        newrelic.agent.register_application(timeout=10.0)
        try:
            return wrapped(*args, **kwargs)
        finally:
            newrelic.agent.shutdown_agent()

    return FunctionWrapper(wrapped, wrapper)


def subprocess_wrapper(module, object_path, name=None, group=None):
    """
    Instrument a function with New Relic being run in a subprocess
    """
    wrap_background_task(module, object_path, application=None, name=name, group=group)
    wrap_object(module, object_path, agent_wrapper)