Python wrapt.ObjectProxy() Examples

def test_onetomany_should_handle_not_being_lazy():
    MockArticle = ObjectProxy(Article)
    MockArticle.get = MagicMock()
    MockArticle.batch_get = MagicMock(return_value=Article.batch_get([1, 3]))
    relationship = OneToMany(MockArticle, lazy=False)
    articles = list(relationship.deserialize([1, 3]))
    MockArticle.batch_get.assert_called_once()
    MockArticle.get.assert_not_called()
    # Access fields that would normally trigger laxy loading
    # order is not guaranteed in batch_get
    if articles[0].id == 1:
        assert articles[0].headline == "Hi!"
        assert articles[1].id == 3
        assert articles[1].headline == "My Article"
    else:
        assert articles[0].id == 3
        assert articles[0].headline == "My Article"
        assert articles[1].id == 1
        assert articles[1].headline == "Hi!"
    # make sure our call count is still 1
    MockArticle.batch_get.assert_called_once()
    MockArticle.get.assert_not_called() 

def get_traced_cursor_proxy(cursor, db_api_integration, *args, **kwargs):
    _traced_cursor = TracedCursor(db_api_integration)
    # pylint: disable=abstract-method
    class TracedCursorProxy(wrapt.ObjectProxy):

        # pylint: disable=unused-argument
        def __init__(self, cursor, *args, **kwargs):
            wrapt.ObjectProxy.__init__(self, cursor)

        def execute(self, *args, **kwargs):
            return _traced_cursor.traced_execution(
                self.__wrapped__.execute, *args, **kwargs
            )

        def executemany(self, *args, **kwargs):
            return _traced_cursor.traced_execution(
                self.__wrapped__.executemany, *args, **kwargs
            )

        def callproc(self, *args, **kwargs):
            return _traced_cursor.traced_execution(
                self.__wrapped__.callproc, *args, **kwargs
            )

    return TracedCursorProxy(cursor, *args, **kwargs) 

def test_patch_db_requests(mock_context_wrapper,):
    """Asserts that monkey patching occurs if iopipe present"""
    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__")

    patch_db_requests(mock_context_wrapper)

    assert hasattr(MySQLdb.connect, "__wrapped__")
    assert hasattr(pymysql.connect, "__wrapped__")

    assert isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert hasattr(Redis.execute_command, "__wrapped__")
    assert hasattr(Pipeline.execute, "__wrapped__")
    assert hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert hasattr(getattr(Collection, class_method), "__wrapped__")

    restore_db_requests()

    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__") 

def test_patch_db_requests_no_iopipe(mock_context,):
    """Asserts that monkey patching does not occur if IOpipe not present"""
    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__")

    delattr(mock_context, "iopipe")

    patch_db_requests(mock_context)

    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__") 

def prepare(self, *args, **kwargs):  # pragma: no cover
        if not args:
            return self.__wrapped__.prepare(*args, **kwargs)

        connection = args[0]

        if isinstance(connection, wrapt.ObjectProxy):
            connection = connection.__wrapped__

        return self.__wrapped__.prepare(connection, *args[1:], **kwargs) 

def test_result_should_be_lazy():
    MockArticle = ObjectProxy(Article)
    MockArticle.get = MagicMock(return_value=Article.get(1))
    relationship = RelationshipResult('id', 1, MockArticle)

    # test before db call (lazy)
    assert relationship.id == 1
    MockArticle.get.assert_not_called()
    # test db call only once
    assert relationship.headline == "Hi!"
    assert relationship.reporter.id == 1
    assert relationship.headline == "Hi!"
    MockArticle.get.assert_called_once_with(1) 

def test_onetoone_should_handle_not_being_lazy():
    MockArticle = ObjectProxy(Article)
    MockArticle.get = MagicMock(return_value=Article.get(1))
    relationship = OneToOne(MockArticle, lazy=False)
    article = relationship.deserialize(1)
    MockArticle.get.assert_called_once_with(1)

    # Access fields that would normally trigger laxy loading
    assert article.id == 1
    assert article.headline == "Hi!"
    assert article.reporter.id == 1
    assert article.headline == "Hi!"
    # make sure our call count is still 1
    MockArticle.get.assert_called_once_with(1) 

def unwrap(obj, attr):
    """
    Will unwrap a `wrapt` attribute
    :param obj: base object
    :param attr: attribute on `obj` to unwrap
    """
    f = getattr(obj, attr, None)
    if f and isinstance(f, wrapt.ObjectProxy) and hasattr(f, '__wrapped__'):
        setattr(obj, attr, f.__wrapped__) 

def unwrap(obj, attr: str):
    """Given a function that was wrapped by wrapt.wrap_function_wrapper, unwrap it

    Args:
        obj: Object that holds a reference to the wrapped function
        attr (str): Name of the wrapped function
    """
    func = getattr(obj, attr, None)
    if func and isinstance(func, ObjectProxy) and hasattr(func, "__wrapped__"):
        setattr(obj, attr, func.__wrapped__) 

def unwrap(obj, attr):
    function = getattr(obj, attr, None)
    if (
        function
        and isinstance(function, ObjectProxy)
        and hasattr(function, "__wrapped__")
    ):
        setattr(obj, attr, function.__wrapped__) 

def uninstrument_connection(connection):
    """Disable instrumentation in a database connection.

    Args:
        connection: The connection to uninstrument.

    Returns:
        An uninstrumented connection.
    """
    if isinstance(connection, wrapt.ObjectProxy):
        return connection.__wrapped__

    logger.warning("Connection is not instrumented")
    return connection 

def get_traced_connection_proxy(
    connection, db_api_integration, *args, **kwargs
):
    # pylint: disable=abstract-method
    class TracedConnectionProxy(wrapt.ObjectProxy):
        # pylint: disable=unused-argument
        def __init__(self, connection, *args, **kwargs):
            wrapt.ObjectProxy.__init__(self, connection)

        def cursor(self, *args, **kwargs):
            return get_traced_cursor_proxy(
                self.__wrapped__.cursor(*args, **kwargs), db_api_integration
            )

    return TracedConnectionProxy(connection, *args, **kwargs) 

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