diff --git a/CHANGELOG.md b/CHANGELOG.md index d1f8850790..0fae8763b7 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py index ab473c2fe4..5d7054151a 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py +++ b/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 diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py index d5e4bc65f3..4fbcac6042 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py +++ b/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({})