test_asyncpg_wrapper.py

import asyncio
from unittest import mock

import pytest
from asyncpg import Connection, Record, cursor
from asyncpg.prepared_stmt import PreparedStatement

try:
    # wrapt 2.0.0+
    from wrapt import BaseObjectProxy  # pylint: disable=no-name-in-module
except ImportError:
    from wrapt import ObjectProxy as BaseObjectProxy

from opentelemetry import trace as trace_api
from opentelemetry.instrumentation.asyncpg import (
    _PREPARED_STMT_METHODS,
    AsyncPGInstrumentor,
)
from opentelemetry.test.test_base import TestBase


class TestAsyncPGInstrumentation(TestBase):
    def tearDown(self):
        super().tearDown()
        AsyncPGInstrumentor().uninstrument()

    def test_duplicated_instrumentation_can_be_uninstrumented(self):
        AsyncPGInstrumentor().instrument()
        AsyncPGInstrumentor().instrument()
        AsyncPGInstrumentor().instrument()
        AsyncPGInstrumentor().uninstrument()
        for method_name in ["execute", "fetch"]:
            method = getattr(Connection, method_name, None)
            self.assertFalse(
                hasattr(method, "_opentelemetry_ext_asyncpg_applied")
            )

    def test_duplicated_instrumentation_works(self):
        first = AsyncPGInstrumentor()
        first.instrument()
        second = AsyncPGInstrumentor()
        second.instrument()
        self.assertIsNotNone(first._tracer)
        self.assertIsNotNone(second._tracer)

    def test_duplicated_uninstrumentation(self):
        AsyncPGInstrumentor().instrument()
        AsyncPGInstrumentor().uninstrument()
        AsyncPGInstrumentor().uninstrument()
        AsyncPGInstrumentor().uninstrument()
        for method_name in ["execute", "fetch"]:
            method = getattr(Connection, method_name, None)
            self.assertFalse(
                hasattr(method, "_opentelemetry_ext_asyncpg_applied")
            )

    def test_cursor_instrumentation(self):
        def assert_wrapped(assert_fnc):
            for cls, methods in [
                (cursor.Cursor, ("forward", "fetch", "fetchrow")),
                (cursor.CursorIterator, ("__anext__",)),
            ]:
                for method_name in methods:
                    method = getattr(cls, method_name, None)
                    assert_fnc(
                        isinstance(method, BaseObjectProxy),
                        f"{method} isinstance {type(method)}",
                    )

        assert_wrapped(self.assertFalse)
        AsyncPGInstrumentor().instrument()
        assert_wrapped(self.assertTrue)
        AsyncPGInstrumentor().uninstrument()
        assert_wrapped(self.assertFalse)

    def test_cursor_span_creation(self):
        """Test the cursor wrapper if it creates spans correctly."""

        # Mock out all interaction with postgres
        async def bind_mock(*args, **kwargs):
            return []

        async def exec_mock(*args, **kwargs):
            return [], None, True

        conn = mock.Mock()
        conn.is_closed = lambda: False

        conn._protocol = mock.Mock()
        conn._protocol.bind = bind_mock
        conn._protocol.execute = exec_mock
        conn._protocol.bind_execute = exec_mock
        conn._protocol.close_portal = bind_mock

        state = mock.Mock()
        state.closed = False

        apg = AsyncPGInstrumentor()
        apg.instrument(tracer_provider=self.tracer_provider)

        # init the cursor and fetch a single record
        crs = cursor.Cursor(conn, "SELECT * FROM test", state, [], Record)
        asyncio.run(crs._init(1))
        asyncio.run(crs.fetch(1))

        spans = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans), 1)
        self.assertEqual(spans[0].name, "CURSOR: SELECT")
        self.assertTrue(spans[0].status.is_ok)

        # Now test that the StopAsyncIteration of the cursor does not get recorded as an ERROR
        crs_iter = cursor.CursorIterator(
            conn, "SELECT * FROM test", state, [], Record, 1, 1
        )

        with pytest.raises(StopAsyncIteration):
            asyncio.run(anext(crs_iter))

        spans = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans), 2)
        self.assertEqual([span.status.is_ok for span in spans], [True, True])

    def test_no_op_tracer_provider(self):
        AsyncPGInstrumentor().uninstrument()
        AsyncPGInstrumentor().instrument(
            tracer_provider=trace_api.NoOpTracerProvider()
        )

        # Mock out all interaction with postgres
        async def bind_mock(*args, **kwargs):
            return []

        async def exec_mock(*args, **kwargs):
            return [], None, True

        conn = mock.Mock()
        conn.is_closed = lambda: False

        conn._protocol = mock.Mock()
        conn._protocol.bind = bind_mock
        conn._protocol.execute = exec_mock
        conn._protocol.bind_execute = exec_mock
        conn._protocol.close_portal = bind_mock

        state = mock.Mock()
        state.closed = False

        # init the cursor and fetch a single record
        crs = cursor.Cursor(conn, "SELECT * FROM test", state, [], Record)
        asyncio.run(crs._init(1))
        asyncio.run(crs.fetch(1))

        spans = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans), 0)

    def test_prepared_statement_instrumentation(self):
        methods = [
            m for m in _PREPARED_STMT_METHODS if hasattr(PreparedStatement, m)
        ]

        for method_name in methods:
            with self.subTest(method=method_name, phase="before"):
                self.assertFalse(
                    isinstance(
                        getattr(PreparedStatement, method_name),
                        BaseObjectProxy,
                    )
                )

        AsyncPGInstrumentor().instrument()

        for method_name in methods:
            with self.subTest(method=method_name, phase="instrumented"):
                self.assertTrue(
                    isinstance(
                        getattr(PreparedStatement, method_name),
                        BaseObjectProxy,
                    )
                )

        AsyncPGInstrumentor().uninstrument()

        for method_name in methods:
            with self.subTest(method=method_name, phase="uninstrumented"):
                self.assertFalse(
                    isinstance(
                        getattr(PreparedStatement, method_name),
                        BaseObjectProxy,
                    )
                )

    @staticmethod
    def _make_prepared_stmt_conn():
        async def bind_execute_mock(*args, **kwargs):
            return [], b"SELECT 1", True

        async def bind_execute_many_mock(*args, **kwargs):
            return None

        conn = mock.Mock()
        conn._pool_release_ctr = 0
        conn.is_closed = lambda: False
        conn._protocol = mock.Mock()
        conn._protocol.bind_execute = bind_execute_mock
        conn._protocol.bind_execute_many = bind_execute_many_mock

        state = mock.Mock()
        state.closed = False
        return conn, state

    def test_prepared_statement_span(self):
        # Per-method: (query, call_args, expected_span_name)
        method_cases = {
            "fetch": ("SELECT * FROM users", (), "SELECT"),
            "fetchval": ("SELECT id FROM users WHERE id=$1", (1,), "SELECT"),
            "fetchrow": ("SELECT * FROM t WHERE v=$1", ("x",), "SELECT"),
            "executemany": (
                "INSERT INTO t (v) VALUES ($1)",
                ([("a",), ("b",)],),
                "INSERT",
            ),
            "fetchmany": ("SELECT * FROM t", ([],), "SELECT"),
        }

        for method_name in _PREPARED_STMT_METHODS:
            if not hasattr(PreparedStatement, method_name):
                continue
            query, call_args, expected_name = method_cases[method_name]
            with self.subTest(method=method_name):
                self.memory_exporter.clear()
                conn, state = self._make_prepared_stmt_conn()
                apg = AsyncPGInstrumentor()
                apg.instrument(tracer_provider=self.tracer_provider)

                stmt = PreparedStatement(conn, query, state)
                asyncio.run(getattr(stmt, method_name)(*call_args))

                spans = self.memory_exporter.get_finished_spans()
                self.assertEqual(len(spans), 1)
                self.assertEqual(spans[0].name, expected_name)
                self.assertTrue(spans[0].status.is_ok)
                self.assertEqual(
                    spans[0].attributes.get("db.statement"), query
                )
                self.assertEqual(
                    spans[0].attributes.get("db.system"), "postgresql"
                )

                apg.uninstrument()

    def test_prepared_statement_error_span(self):
        async def bind_execute_error(*args, **kwargs):
            raise RuntimeError("db error")

        conn = mock.Mock()
        conn._pool_release_ctr = 0
        conn.is_closed = lambda: False
        conn._protocol = mock.Mock()
        conn._protocol.bind_execute = bind_execute_error

        state = mock.Mock()
        state.closed = False

        apg = AsyncPGInstrumentor()
        apg.instrument(tracer_provider=self.tracer_provider)

        stmt = PreparedStatement(conn, "SELECT 1", state)
        with self.assertRaises(RuntimeError):
            asyncio.run(stmt.fetch())

        spans = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans), 1)
        self.assertFalse(spans[0].status.is_ok)