Psycopg3 sync and async instrumentation (#2146)

Author: reiktar

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- `opentelemetry-instrumentation-psycopg` Async Instrumentation for psycopg 3.x
+  ([#2146](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2146))
+
+### Fixed
 - `opentelemetry-instrumentation-celery` Allow Celery instrumentation to be installed multiple times
   ([#2342](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2342))
 - Align gRPC span status codes to OTEL specification

instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py

@@ -105,8 +105,13 @@
 import typing
 from typing import Collection
 
-import psycopg
-from psycopg import Cursor as pg_cursor  # pylint: disable=no-name-in-module
+import psycopg  # pylint: disable=import-self
+from psycopg import (
+    AsyncCursor as pg_async_cursor,  # pylint: disable=import-self,no-name-in-module
+)
+from psycopg import (
+    Cursor as pg_cursor,  # pylint: disable=no-name-in-module,import-self
+)
 from psycopg.sql import Composed  # pylint: disable=no-name-in-module
 
 from opentelemetry.instrumentation import dbapi
@@ -151,9 +156,40 @@ def _instrument(self, **kwargs):
             commenter_options=commenter_options,
         )
 
+        dbapi.wrap_connect(
+            __name__,
+            psycopg.Connection,  # pylint: disable=no-member
+            "connect",
+            self._DATABASE_SYSTEM,
+            self._CONNECTION_ATTRIBUTES,
+            version=__version__,
+            tracer_provider=tracer_provider,
+            db_api_integration_factory=DatabaseApiIntegration,
+            enable_commenter=enable_sqlcommenter,
+            commenter_options=commenter_options,
+        )
+        dbapi.wrap_connect(
+            __name__,
+            psycopg.AsyncConnection,  # pylint: disable=no-member
+            "connect",
+            self._DATABASE_SYSTEM,
+            self._CONNECTION_ATTRIBUTES,
+            version=__version__,
+            tracer_provider=tracer_provider,
+            db_api_integration_factory=DatabaseApiAsyncIntegration,
+            enable_commenter=enable_sqlcommenter,
+            commenter_options=commenter_options,
+        )
+
     def _uninstrument(self, **kwargs):
         """ "Disable Psycopg instrumentation"""
-        dbapi.unwrap_connect(psycopg, "connect")
+        dbapi.unwrap_connect(psycopg, "connect")  # pylint: disable=no-member
+        dbapi.unwrap_connect(
+            psycopg.Connection, "connect"  # pylint: disable=no-member
+        )
+        dbapi.unwrap_connect(
+            psycopg.AsyncConnection, "connect"  # pylint: disable=no-member
+        )
 
     # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
     @staticmethod
@@ -204,6 +240,26 @@ def wrapped_connection(
         return connection
 
 
+class DatabaseApiAsyncIntegration(dbapi.DatabaseApiIntegration):
+    async def wrapped_connection(
+        self,
+        connect_method: typing.Callable[..., typing.Any],
+        args: typing.Tuple[typing.Any, typing.Any],
+        kwargs: typing.Dict[typing.Any, typing.Any],
+    ):
+        """Add object proxy to connection object."""
+        base_cursor_factory = kwargs.pop("cursor_factory", None)
+        new_factory_kwargs = {"db_api": self}
+        if base_cursor_factory:
+            new_factory_kwargs["base_factory"] = base_cursor_factory
+        kwargs["cursor_factory"] = _new_cursor_async_factory(
+            **new_factory_kwargs
+        )
+        connection = await connect_method(*args, **kwargs)
+        self.get_connection_attributes(connection)
+        return connection
+
+
 class CursorTracer(dbapi.CursorTracer):
     def get_operation_name(self, cursor, args):
         if not args:
@@ -259,3 +315,36 @@ def callproc(self, *args, **kwargs):
             )
 
     return TracedCursorFactory
+
+
+def _new_cursor_async_factory(
+    db_api=None, base_factory=None, tracer_provider=None
+):
+    if not db_api:
+        db_api = DatabaseApiAsyncIntegration(
+            __name__,
+            PsycopgInstrumentor._DATABASE_SYSTEM,
+            connection_attributes=PsycopgInstrumentor._CONNECTION_ATTRIBUTES,
+            version=__version__,
+            tracer_provider=tracer_provider,
+        )
+    base_factory = base_factory or pg_async_cursor
+    _cursor_tracer = CursorTracer(db_api)
+
+    class TracedCursorAsyncFactory(base_factory):
+        async def execute(self, *args, **kwargs):
+            return await _cursor_tracer.traced_execution(
+                self, super().execute, *args, **kwargs
+            )
+
+        async def executemany(self, *args, **kwargs):
+            return await _cursor_tracer.traced_execution(
+                self, super().executemany, *args, **kwargs
+            )
+
+        async def callproc(self, *args, **kwargs):
+            return await _cursor_tracer.traced_execution(
+                self, super().callproc, *args, **kwargs
+            )
+
+    return TracedCursorAsyncFactory

instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import types
 from unittest import mock
 
@@ -45,6 +46,35 @@ def __exit__(self, *args):
         return self
 
 
+class MockAsyncCursor:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    # pylint: disable=unused-argument, no-self-use
+    async def execute(self, query, params=None, throw_exception=False):
+        if throw_exception:
+            raise Exception("Test Exception")
+
+    # pylint: disable=unused-argument, no-self-use
+    async def executemany(self, query, params=None, throw_exception=False):
+        if throw_exception:
+            raise Exception("Test Exception")
+
+    # pylint: disable=unused-argument, no-self-use
+    async def callproc(self, query, params=None, throw_exception=False):
+        if throw_exception:
+            raise Exception("Test Exception")
+
+    async def __aenter__(self, *args, **kwargs):
+        return self
+
+    async def __aexit__(self, *args, **kwargs):
+        pass
+
+    def close(self):
+        pass
+
+
 class MockConnection:
     commit = mock.MagicMock(spec=types.MethodType)
     commit.__name__ = "commit"
@@ -64,22 +94,68 @@ def get_dsn_parameters(self):  # pylint: disable=no-self-use
         return {"dbname": "test"}
 
 
+class MockAsyncConnection:
+    commit = mock.MagicMock(spec=types.MethodType)
+    commit.__name__ = "commit"
+
+    rollback = mock.MagicMock(spec=types.MethodType)
+    rollback.__name__ = "rollback"
+
+    def __init__(self, *args, **kwargs):
+        self.cursor_factory = kwargs.pop("cursor_factory", None)
+
+    @staticmethod
+    async def connect(*args, **kwargs):
+        return MockAsyncConnection(**kwargs)
+
+    def cursor(self):
+        if self.cursor_factory:
+            cur = self.cursor_factory(self)
+            return cur
+        return MockAsyncCursor()
+
+    def get_dsn_parameters(self):  # pylint: disable=no-self-use
+        return {"dbname": "test"}
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        return mock.MagicMock(spec=types.MethodType)
+
+
 class TestPostgresqlIntegration(TestBase):
     def setUp(self):
         super().setUp()
         self.cursor_mock = mock.patch(
             "opentelemetry.instrumentation.psycopg.pg_cursor", MockCursor
         )
+        self.cursor_async_mock = mock.patch(
+            "opentelemetry.instrumentation.psycopg.pg_async_cursor",
+            MockAsyncCursor,
+        )
         self.connection_mock = mock.patch("psycopg.connect", MockConnection)
+        self.connection_sync_mock = mock.patch(
+            "psycopg.Connection.connect", MockConnection
+        )
+        self.connection_async_mock = mock.patch(
+            "psycopg.AsyncConnection.connect", MockAsyncConnection.connect
+        )
 
         self.cursor_mock.start()
+        self.cursor_async_mock.start()
         self.connection_mock.start()
+        self.connection_sync_mock.start()
+        self.connection_async_mock.start()
 
     def tearDown(self):
         super().tearDown()
         self.memory_exporter.clear()
         self.cursor_mock.stop()
+        self.cursor_async_mock.stop()
         self.connection_mock.stop()
+        self.connection_sync_mock.stop()
+        self.connection_async_mock.stop()
         with self.disable_logging():
             PsycopgInstrumentor().uninstrument()
 
@@ -114,6 +190,91 @@ def test_instrumentor(self):
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
 
+    # pylint: disable=unused-argument
+    def test_instrumentor_with_connection_class(self):
+        PsycopgInstrumentor().instrument()
+
+        cnx = psycopg.Connection.connect(database="test")
+
+        cursor = cnx.cursor()
+
+        query = "SELECT * FROM test"
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+
+        # Check version and name in span's instrumentation info
+        self.assertEqualSpanInstrumentationInfo(
+            span, opentelemetry.instrumentation.psycopg
+        )
+
+        # check that no spans are generated after uninstrument
+        PsycopgInstrumentor().uninstrument()
+
+        cnx = psycopg.Connection.connect(database="test")
+        cursor = cnx.cursor()
+        query = "SELECT * FROM test"
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+
+    async def test_wrap_async_connection_class_with_cursor(self):
+        PsycopgInstrumentor().instrument()
+
+        async def test_async_connection():
+            acnx = await psycopg.AsyncConnection.connect(database="test")
+            async with acnx as cnx:
+                async with cnx.cursor() as cursor:
+                    await cursor.execute("SELECT * FROM test")
+
+        asyncio.run(test_async_connection())
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+
+        # Check version and name in span's instrumentation info
+        self.assertEqualSpanInstrumentationInfo(
+            span, opentelemetry.instrumentation.psycopg
+        )
+
+        # check that no spans are generated after uninstrument
+        PsycopgInstrumentor().uninstrument()
+
+        asyncio.run(test_async_connection())
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+
+    # pylint: disable=unused-argument
+    async def test_instrumentor_with_async_connection_class(self):
+        PsycopgInstrumentor().instrument()
+
+        async def test_async_connection():
+            acnx = await psycopg.AsyncConnection.connect(database="test")
+            async with acnx as cnx:
+                await cnx.execute("SELECT * FROM test")
+
+        asyncio.run(test_async_connection())
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+
+        # Check version and name in span's instrumentation info
+        self.assertEqualSpanInstrumentationInfo(
+            span, opentelemetry.instrumentation.psycopg
+        )
+
+        # check that no spans are generated after uninstrument
+        PsycopgInstrumentor().uninstrument()
+        asyncio.run(test_async_connection())
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+
     def test_span_name(self):
         PsycopgInstrumentor().instrument()
 
@@ -140,6 +301,33 @@ def test_span_name(self):
         self.assertEqual(spans_list[4].name, "query")
         self.assertEqual(spans_list[5].name, "query")
 
+    async def test_span_name_async(self):
+        PsycopgInstrumentor().instrument()
+
+        cnx = psycopg.AsyncConnection.connect(database="test")
+        async with cnx.cursor() as cursor:
+            await cursor.execute("Test query", ("param1Value", False))
+            await cursor.execute(
+                """multi
+    line
+    query"""
+            )
+            await cursor.execute("tab\tseparated query")
+            await cursor.execute("/* leading comment */ query")
+            await cursor.execute(
+                "/* leading comment */ query /* trailing comment */"
+            )
+            await cursor.execute("query /* trailing comment */")
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 6)
+        self.assertEqual(spans_list[0].name, "Test")
+        self.assertEqual(spans_list[1].name, "multi")
+        self.assertEqual(spans_list[2].name, "tab")
+        self.assertEqual(spans_list[3].name, "query")
+        self.assertEqual(spans_list[4].name, "query")
+        self.assertEqual(spans_list[5].name, "query")
+
     # pylint: disable=unused-argument
     def test_not_recording(self):
         mock_tracer = mock.Mock()
@@ -160,6 +348,26 @@ def test_not_recording(self):
 
         PsycopgInstrumentor().uninstrument()
 
+    # pylint: disable=unused-argument
+    async def test_not_recording_async(self):
+        mock_tracer = mock.Mock()
+        mock_span = mock.Mock()
+        mock_span.is_recording.return_value = False
+        mock_tracer.start_span.return_value = mock_span
+        PsycopgInstrumentor().instrument()
+        with mock.patch("opentelemetry.trace.get_tracer") as tracer:
+            tracer.return_value = mock_tracer
+            cnx = psycopg.AsyncConnection.connect(database="test")
+            async with cnx.cursor() as cursor:
+                query = "SELECT * FROM test"
+                cursor.execute(query)
+            self.assertFalse(mock_span.is_recording())
+            self.assertTrue(mock_span.is_recording.called)
+            self.assertFalse(mock_span.set_attribute.called)
+            self.assertFalse(mock_span.set_status.called)
+
+        PsycopgInstrumentor().uninstrument()
+
     # pylint: disable=unused-argument
     def test_custom_tracer_provider(self):
         resource = resources.Resource.create({})