Python wrapt.wrap_function_wrapper方法代码示例

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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,
    ) 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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,
    ) 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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,
    ) 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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
            ) 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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) 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
def patch():

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

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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
    ) 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
def patch():

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

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
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 

# 需要导入模块: import wrapt [as 别名]
# 或者: from wrapt import wrap_function_wrapper [as 别名]
def patch():

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

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