Python wrapt.wrap_function_wrapper() Examples

def patch_botocore_session_send(context, http_filter, http_headers):
    """
    Monkey patches botocore's session, if available. Overloads the
    session class' send method to add tracing and metric collection.
    """

    def wrapper(wrapped, instance, args, kwargs):
        if not hasattr(context, "iopipe") or not hasattr(context.iopipe, "mark"):
            return wrapped(*args, **kwargs)
        id = str(uuid.uuid4())
        with context.iopipe.mark(id):
            response = wrapped(*args, **kwargs)
        trace = context.iopipe.mark.measure(id)
        context.iopipe.mark.delete(id)
        collect_metrics_for_response(
            args[0], response, context, trace, http_filter, http_headers
        )
        return response

    try:
        wrapt.wrap_function_wrapper(
            "botocore.httpsession", "URLLib3Session.send", wrapper
        )
    except Exception:  # pragma: no cover
        pass 

def patch_botocore_vendored_session_send(context, http_filter, http_headers):
    """
    Monkey patches botocore's vendored requests, if available. Overloads the
    session class' send method to add tracing and metric collection.
    """

    def wrapper(wrapped, instance, args, kwargs):
        if not hasattr(context, "iopipe") or not hasattr(context.iopipe, "mark"):
            return wrapped(*args, **kwargs)
        id = str(uuid.uuid4())
        with context.iopipe.mark(id):
            response = wrapped(*args, **kwargs)
        trace = context.iopipe.mark.measure(id)
        context.iopipe.mark.delete(id)
        collect_metrics_for_response(
            response.request, response, context, trace, http_filter, http_headers
        )
        return response

    try:
        wrapt.wrap_function_wrapper(
            "botocore.vendored.requests.sessions", "Session.send", wrapper
        )
    except Exception:  # pragma: no cover
        pass 

def patch_requests_session_send(context, http_filter, http_headers):
    """
    Monkey patches requests' session class, if available. Overloads the
    send method to add tracing and metrics collection.
    """

    def wrapper(wrapped, instance, args, kwargs):
        if not hasattr(context, "iopipe") or not hasattr(context.iopipe, "mark"):
            return wrapped(*args, **kwargs)
        id = ensure_utf8(str(uuid.uuid4()))
        with context.iopipe.mark(id):
            response = wrapped(*args, **kwargs)
        trace = context.iopipe.mark.measure(id)
        context.iopipe.mark.delete(id)
        collect_metrics_for_response(
            response.request, response, context, trace, http_filter, http_headers
        )
        return response

    try:
        wrapt.wrap_function_wrapper("requests.sessions", "Session.send", wrapper)
    except Exception:  # pragma: no cover
        pass 

def trace_integration(tracer=None):
    """Wrap the requests library to trace it."""
    log.info('Integrated module: {}'.format(MODULE_NAME))

    if tracer is not None:
        # The execution_context tracer should never be None - if it has not
        # been set it returns a no-op tracer. Most code in this library does
        # not handle None being used in the execution context.
        execution_context.set_opencensus_tracer(tracer)

    # Wrap the requests functions
    for func in REQUESTS_WRAP_METHODS:
        requests_func = getattr(requests, func)
        wrapped = wrap_requests(requests_func)
        setattr(requests, requests_func.__name__, wrapped)

    # Wrap Session class
    wrapt.wrap_function_wrapper(
        MODULE_NAME, 'Session.request', wrap_session_request) 

def patch():
    """
    Patch aiobotocore client so it generates subsegments
    when calling AWS services.
    """
    if hasattr(aiobotocore.client, '_xray_enabled'):
        return
    setattr(aiobotocore.client, '_xray_enabled', True)

    wrapt.wrap_function_wrapper(
        'aiobotocore.client',
        'AioBaseClient._make_api_call',
        _xray_traced_aiobotocore,
    )

    wrapt.wrap_function_wrapper(
        'aiobotocore.endpoint',
        'AioEndpoint.prepare_request',
        inject_header,
    ) 

def _instrument(self, **kwargs):
        tracer_provider = kwargs.get(
            "tracer_provider", trace.get_tracer_provider()
        )
        setattr(
            asyncpg,
            _APPLIED,
            tracer_provider.get_tracer("asyncpg", __version__),
        )
        for method in [
            "Connection.execute",
            "Connection.executemany",
            "Connection.fetch",
            "Connection.fetchval",
            "Connection.fetchrow",
        ]:
            wrapt.wrap_function_wrapper(
                "asyncpg.connection", method, _do_execute
            ) 

def patch():
    """
    Patch botocore client so it generates subsegments
    when calling AWS services.
    """
    if hasattr(botocore.client, '_xray_enabled'):
        return
    setattr(botocore.client, '_xray_enabled', True)

    wrapt.wrap_function_wrapper(
        'botocore.client',
        'BaseClient._make_api_call',
        _xray_traced_botocore,
    )

    wrapt.wrap_function_wrapper(
        'botocore.endpoint',
        'Endpoint.prepare_request',
        inject_header,
    ) 

def patch():
    """Patch PynamoDB so it generates subsegements when calling DynamoDB."""

    if PYNAMODB4:
        if hasattr(botocore.httpsession, '_xray_enabled'):
            return
        setattr(botocore.httpsession, '_xray_enabled', True)

        module = 'botocore.httpsession'
        name = 'URLLib3Session.send'
    else:
        if hasattr(botocore.vendored.requests.sessions, '_xray_enabled'):
            return
        setattr(botocore.vendored.requests.sessions, '_xray_enabled', True)

        module = 'botocore.vendored.requests.sessions'
        name = 'Session.send'

    wrapt.wrap_function_wrapper(
        module, name, _xray_traced_pynamodb,
    ) 

def patch():
    wrapt.wrap_function_wrapper('requests', 'Session.request', _wrapped_request)
    mark_patched('requests') 

def monkey_patch_requests_for_xray():
    """Explicit way to monkey-patch requests to trace HTTP requests."""
    wrapt.wrap_function_wrapper(
        "requests.sessions", "Session.send", xray_requests_send,
    ) 

def monkey_patch_botocore_for_xray():
    """Explicit way to monkey-patch botocore to trace AWS API calls."""
    wrapt.wrap_function_wrapper(
        "botocore.client", "BaseClient._make_api_call", xray_botocore_api_call,
    ) 

def patch():

    wrapt.wrap_function_wrapper(
        'requests',
        'Session.request',
        _xray_traced_requests
    )

    wrapt.wrap_function_wrapper(
        'requests',
        'Session.prepare_request',
        _inject_header
    ) 

def patch():

    wrapt.wrap_function_wrapper(
        'pymysql',
        'connect',
        _xray_traced_connect
    )

    # patch alias
    if hasattr(pymysql, 'Connect'):
        pymysql.Connect = pymysql.connect 

def patch():

    wrapt.wrap_function_wrapper(
        'mysql.connector',
        'connect',
        _xray_traced_connect
    )

    # patch alias
    if hasattr(mysql.connector, 'Connect'):
        mysql.connector.Connect = mysql.connector.connect 

def patch():

    wrapt.wrap_function_wrapper(
        'pg8000',
        'connect',
        _xray_traced_connect
    ) 

def patch():
    """
    patch the built-in `urllib/httplib/httplib.client` methods for tracing.
    """
    if getattr(httplib, PATCH_FLAG, False):
        return
    # we set an attribute to avoid multiple wrapping
    setattr(httplib, PATCH_FLAG, True)

    wrapt.wrap_function_wrapper(
        httplib_client_module,
        'HTTPConnection._send_request',
        _send_request
    )

    wrapt.wrap_function_wrapper(
        httplib_client_module,
        'HTTPConnection.getresponse',
        _xray_traced_http_getresponse
    )

    wrapt.wrap_function_wrapper(
        httplib_client_module,
        'HTTPResponse.read',
        _xray_traced_http_client_read
    ) 

def patch():

    wrapt.wrap_function_wrapper(
        'sqlite3',
        'connect',
        _xray_traced_connect
    ) 

def patch_pymysql(context):
    """
    Monkey patches pymysql client, if available. Overloads the
    execute method to add tracing and metrics collection.
    """

    class _CursorProxy(CursorProxy):
        def execute(self, *args, **kwargs):
            if not hasattr(context, "iopipe") or not hasattr(
                context.iopipe, "mark"
            ):  # pragma: no cover
                self.__wrapped__.execute(*args, **kwargs)
                return

            id = ensure_utf8(str(uuid.uuid4()))
            with context.iopipe.mark(id):
                self.__wrapped__.execute(*args, **kwargs)
            trace = context.iopipe.mark.measure(id)
            context.iopipe.mark.delete(id)
            collect_mysql_metrics(context, trace, self, args)

    class _ConnectionProxy(ConnectionProxy):
        def cursor(self, *args, **kwargs):
            cursor = self.__wrapped__.cursor(*args, **kwargs)
            return _CursorProxy(cursor, self)

    def connect_wrapper(wrapped, instance, args, kwargs):
        connection = wrapped(*args, **kwargs)
        return _ConnectionProxy(connection, args, kwargs)

    try:
        wrapt.wrap_function_wrapper("pymysql", "connect", connect_wrapper)
    except Exception:  # pragma: no cover
        pass 

def patch_mysqldb(context):
    """
    Monkey patches mysqldb client, if available. Overloads the
    execute method to add tracing and metrics collection.
    """

    class _CursorProxy(CursorProxy):
        def execute(self, *args, **kwargs):
            if not hasattr(context, "iopipe") or not hasattr(
                context.iopipe, "mark"
            ):  # pragma: no cover
                self.__wrapped__.execute(*args, **kwargs)
                return

            id = ensure_utf8(str(uuid.uuid4()))
            with context.iopipe.mark(id):
                self.__wrapped__.execute(*args, **kwargs)
            trace = context.iopipe.mark.measure(id)
            context.iopipe.mark.delete(id)
            collect_mysql_metrics(context, trace, self, args)

    class _ConnectionProxy(ConnectionProxy):
        def cursor(self, *args, **kwargs):
            cursor = self.__wrapped__.cursor(*args, **kwargs)
            return _CursorProxy(cursor, self)

    def connect_wrapper(wrapped, instance, args, kwargs):
        connection = wrapped(*args, **kwargs)
        return _ConnectionProxy(connection, args, kwargs)

    for module, attr, wrapper in [
        ("MySQLdb", "connect", connect_wrapper),
        ("MySQLdb", "Connection", connect_wrapper),
        ("MySQLdb", "Connect", connect_wrapper),
    ]:
        try:
            wrapt.wrap_function_wrapper(module, attr, wrapper)
        except Exception:  # pragma: no cover
            pass 

def wrap_connect(
    tracer: Tracer,
    connect_module: typing.Callable[..., typing.Any],
    connect_method_name: str,
    database_component: str,
    database_type: str = "",
    connection_attributes: typing.Dict = None,
):
    """Integrate with DB API library.
        https://www.python.org/dev/peps/pep-0249/

        Args:
            tracer: The :class:`opentelemetry.trace.Tracer` to use.
            connect_module: Module name where connect method is available.
            connect_method_name: The connect method name.
            database_component: Database driver name or database name "JDBI",
                "jdbc", "odbc", "postgreSQL".
            database_type: The Database type. For any SQL database, "sql".
            connection_attributes: Attribute names for database, port, host and
                user in Connection object.
    """

    # pylint: disable=unused-argument
    def wrap_connect_(
        wrapped: typing.Callable[..., typing.Any],
        instance: typing.Any,
        args: typing.Tuple[typing.Any, typing.Any],
        kwargs: typing.Dict[typing.Any, typing.Any],
    ):
        db_integration = DatabaseApiIntegration(
            tracer,
            database_component,
            database_type=database_type,
            connection_attributes=connection_attributes,
        )
        return db_integration.wrapped_connection(wrapped, args, kwargs)

    try:
        wrapt.wrap_function_wrapper(
            connect_module, connect_method_name, wrap_connect_
        )
    except Exception as ex:  # pylint: disable=broad-except
        logger.warning("Failed to integrate with DB API. %s", str(ex)) 

def patch_redis(context):
    """
    Monkey patches redis client, if available. Overloads the
    execute methods to add tracing and metrics collection.
    """

    def wrapper(wrapped, instance, args, kwargs):
        if not hasattr(context, "iopipe") or not hasattr(
            context.iopipe, "mark"
        ):  # pragma: no cover
            return wrapped(*args, **kwargs)

        id = ensure_utf8(str(uuid.uuid4()))
        with context.iopipe.mark(id):
            response = wrapped(*args, **kwargs)
        trace = context.iopipe.mark.measure(id)
        context.iopipe.mark.delete(id)
        collect_redis_metrics(
            context, trace, args, instance.connection_pool.connection_kwargs
        )
        return response

    def pipeline_wrapper(wrapped, instance, args, kwargs):  # pragma: no cover
        if not hasattr(context, "iopipe") or not hasattr(
            context.iopipe, "mark"
        ):  # pragma: no cover
            return wrapped(*args, **kwargs)

        # We don't need the entire command stack, just collect a stack count
        pipeline_args = ("PIPELINE", ensure_utf8(len(instance.command_stack)))

        id = ensure_utf8(str(uuid.uuid4()))
        with context.iopipe.mark(id):
            response = wrapped(*args, **kwargs)
        trace = context.iopipe.mark.measure(id)
        context.iopipe.mark.delete(id)
        collect_redis_metrics(
            context, trace, pipeline_args, instance.connection_pool.connection_kwargs
        )
        return response

    for module, attr, _wrapper in [
        ("redis.client", "Redis.execute_command", wrapper),
        ("redis.client", "Pipeline.execute", wrapper),
        ("redis.client", "Pipeline.immediate_execute_command", wrapper),
    ]:
        try:
            wrapt.wrap_function_wrapper(module, attr, _wrapper)
        except Exception:  # pragma: no cover
            pass 

def patch_psycopg2(context):
    """
    Monkey patches psycopg2 client, if available. Overloads the
    execute method to add tracing and metrics collection.
    """

    class _CursorProxy(CursorProxy):
        def execute(self, *args, **kwargs):
            if not hasattr(context, "iopipe") or not hasattr(
                context.iopipe, "mark"
            ):  # pragma: no cover
                self.__wrapped__.execute(*args, **kwargs)
                return

            id = ensure_utf8(str(uuid.uuid4()))
            with context.iopipe.mark(id):
                self.__wrapped__.execute(*args, **kwargs)
            trace = context.iopipe.mark.measure(id)
            context.iopipe.mark.delete(id)
            collect_psycopg2_metrics(context, trace, self, args)

    class _ConnectionProxy(ConnectionProxy):
        def cursor(self, *args, **kwargs):
            cursor = self.__wrapped__.cursor(*args, **kwargs)
            return _CursorProxy(cursor, self)

    def adapt_wrapper(wrapped, instance, args, kwargs):
        adapter = wrapped(*args, **kwargs)
        return AdapterProxy(adapter) if hasattr(adapter, "prepare") else adapter

    def connect_wrapper(wrapped, instance, args, kwargs):
        connection = wrapped(*args, **kwargs)
        return _ConnectionProxy(connection, args, kwargs)

    def register_type_wrapper(wrapped, instance, args, kwargs):
        def _extract_arguments(obj, scope=None):
            return obj, scope

        obj, scope = _extract_arguments(*args, **kwargs)

        if scope is not None:
            if isinstance(scope, wrapt.ObjectProxy):
                scope = scope.__wrapped__
            return wrapped(obj, scope)

        return wrapped(obj)

    for module, attr, wrapper in [
        ("psycopg2", "connect", connect_wrapper),
        ("psycopg2.extensions", "adapt", adapt_wrapper),
        ("psycopg2.extensions", "register_type", register_type_wrapper),
        ("psycopg2._psycopg", "register_type", register_type_wrapper),
        ("psycopg2._json", "register_type", register_type_wrapper),
    ]:
        try:
            wrapt.wrap_function_wrapper(module, attr, wrapper)
        except Exception:  # pragma: no cover
            pass