From 66a1747f0fbad7b948a318ffdc4d681f6b63cded Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Wed, 4 Dec 2024 18:25:12 +0100 Subject: [PATCH 1/8] Add type hints to Psycopg --- .../instrumentation/psycopg/__init__.py | 67 ++++++++++--------- .../instrumentation/psycopg/package.py | 4 +- 2 files changed, 39 insertions(+), 32 deletions(-) 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 e986ec0d46..7668fa806e 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py @@ -101,27 +101,26 @@ --- """ +from __future__ import annotations + import logging -import typing -from typing import Collection +from typing import Any, Callable, Collection, TypeVar 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 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.psycopg.package import _instruments from opentelemetry.instrumentation.psycopg.version import __version__ +from opentelemetry.trace import TracerProvider _logger = logging.getLogger(__name__) _OTEL_CURSOR_FACTORY_KEY = "_otel_orig_cursor_factory" +Connection = TypeVar("Connection", psycopg.Connection, psycopg.AsyncConnection) +Cursor = TypeVar("Cursor", psycopg.Cursor, psycopg.AsyncCursor) + class PsycopgInstrumentor(BaseInstrumentor): _CONNECTION_ATTRIBUTES = { @@ -136,7 +135,7 @@ class PsycopgInstrumentor(BaseInstrumentor): def instrumentation_dependencies(self) -> Collection[str]: return _instruments - def _instrument(self, **kwargs): + def _instrument(self, **kwargs: Any): """Integrate with PostgreSQL Psycopg library. Psycopg: http://initd.org/psycopg/ """ @@ -181,7 +180,7 @@ def _instrument(self, **kwargs): commenter_options=commenter_options, ) - def _uninstrument(self, **kwargs): + def _uninstrument(self, **kwargs: Any): """ "Disable Psycopg instrumentation""" dbapi.unwrap_connect(psycopg, "connect") # pylint: disable=no-member dbapi.unwrap_connect( @@ -195,7 +194,9 @@ def _uninstrument(self, **kwargs): # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql @staticmethod - def instrument_connection(connection, tracer_provider=None): + def instrument_connection( + connection: Connection, tracer_provider: TracerProvider | None = None + ) -> Connection: if not hasattr(connection, "_is_instrumented_by_opentelemetry"): connection._is_instrumented_by_opentelemetry = False @@ -215,7 +216,7 @@ def instrument_connection(connection, tracer_provider=None): # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql @staticmethod - def uninstrument_connection(connection): + def uninstrument_connection(connection: Connection) -> Connection: connection.cursor_factory = getattr( connection, _OTEL_CURSOR_FACTORY_KEY, None ) @@ -227,9 +228,9 @@ def uninstrument_connection(connection): class DatabaseApiIntegration(dbapi.DatabaseApiIntegration): def wrapped_connection( self, - connect_method: typing.Callable[..., typing.Any], - args: typing.Tuple[typing.Any, typing.Any], - kwargs: typing.Dict[typing.Any, typing.Any], + connect_method: Callable[..., Any], + args: tuple[Any, Any], + kwargs: dict[Any, Any], ): """Add object proxy to connection object.""" base_cursor_factory = kwargs.pop("cursor_factory", None) @@ -245,9 +246,9 @@ def wrapped_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], + connect_method: Callable[..., Any], + args: tuple[Any, Any], + kwargs: dict[Any, Any], ): """Add object proxy to connection object.""" base_cursor_factory = kwargs.pop("cursor_factory", None) @@ -263,7 +264,7 @@ async def wrapped_connection( class CursorTracer(dbapi.CursorTracer): - def get_operation_name(self, cursor, args): + def get_operation_name(self, cursor: Cursor, args: list[Any]) -> str: if not args: return "" @@ -278,7 +279,7 @@ def get_operation_name(self, cursor, args): return "" - def get_statement(self, cursor, args): + def get_statement(self, cursor: Cursor, args: list[Any]) -> str: if not args: return "" @@ -288,7 +289,11 @@ def get_statement(self, cursor, args): return statement -def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None): +def _new_cursor_factory( + db_api: DatabaseApiIntegration | None = None, + base_factory: type[psycopg.Cursor] | None = None, + tracer_provider: TracerProvider | None = None, +): if not db_api: db_api = DatabaseApiIntegration( __name__, @@ -298,21 +303,21 @@ def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None): tracer_provider=tracer_provider, ) - base_factory = base_factory or pg_cursor + base_factory = base_factory or psycopg.Cursor _cursor_tracer = CursorTracer(db_api) class TracedCursorFactory(base_factory): - def execute(self, *args, **kwargs): + def execute(self, *args: Any, **kwargs: Any): return _cursor_tracer.traced_execution( self, super().execute, *args, **kwargs ) - def executemany(self, *args, **kwargs): + def executemany(self, *args: Any, **kwargs: Any): return _cursor_tracer.traced_execution( self, super().executemany, *args, **kwargs ) - def callproc(self, *args, **kwargs): + def callproc(self, *args: Any, **kwargs: Any): return _cursor_tracer.traced_execution( self, super().callproc, *args, **kwargs ) @@ -321,7 +326,9 @@ def callproc(self, *args, **kwargs): def _new_cursor_async_factory( - db_api=None, base_factory=None, tracer_provider=None + db_api: DatabaseApiAsyncIntegration | None = None, + base_factory: type[psycopg.AsyncCursor] | None = None, + tracer_provider: TracerProvider | None = None, ): if not db_api: db_api = DatabaseApiAsyncIntegration( @@ -331,21 +338,21 @@ def _new_cursor_async_factory( version=__version__, tracer_provider=tracer_provider, ) - base_factory = base_factory or pg_async_cursor + base_factory = base_factory or psycopg.AsyncCursor _cursor_tracer = CursorTracer(db_api) class TracedCursorAsyncFactory(base_factory): - async def execute(self, *args, **kwargs): + async def execute(self, *args: Any, **kwargs: Any): return await _cursor_tracer.traced_execution( self, super().execute, *args, **kwargs ) - async def executemany(self, *args, **kwargs): + async def executemany(self, *args: Any, **kwargs: Any): return await _cursor_tracer.traced_execution( self, super().executemany, *args, **kwargs ) - async def callproc(self, *args, **kwargs): + async def callproc(self, *args: Any, **kwargs: Any): return await _cursor_tracer.traced_execution( self, super().callproc, *args, **kwargs ) diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/package.py b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/package.py index 635edfb4db..a3ee72d1ae 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/package.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/package.py @@ -11,6 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations - -_instruments = ("psycopg >= 3.1.0",) +_instruments: tuple[str, ...] = ("psycopg >= 3.1.0",) From 8f154d88db97f0205d8b2da0745081181a76257b Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Wed, 4 Dec 2024 23:11:13 +0100 Subject: [PATCH 2/8] fix tests --- .../tests/test_psycopg_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py index 4ddaad9174..24d062dc67 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py @@ -132,10 +132,10 @@ class PostgresqlIntegrationTestMixin: def setUp(self): super().setUp() self.cursor_mock = mock.patch( - "opentelemetry.instrumentation.psycopg.pg_cursor", MockCursor + "opentelemetry.instrumentation.psycopg.Cursor", MockCursor ) self.cursor_async_mock = mock.patch( - "opentelemetry.instrumentation.psycopg.pg_async_cursor", + "opentelemetry.instrumentation.psycopg.AsyncCursor", MockAsyncCursor, ) self.connection_mock = mock.patch("psycopg.connect", MockConnection) From 47835aa4fcc3df5c5f79b018009412ec6d68976d Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Wed, 4 Dec 2024 23:21:31 +0100 Subject: [PATCH 3/8] fix --- .../tests/test_psycopg_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py index 24d062dc67..6c9bcf2d4b 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py @@ -132,10 +132,10 @@ class PostgresqlIntegrationTestMixin: def setUp(self): super().setUp() self.cursor_mock = mock.patch( - "opentelemetry.instrumentation.psycopg.Cursor", MockCursor + "opentelemetry.instrumentation.psycopg.psycopg.Cursor", MockCursor ) self.cursor_async_mock = mock.patch( - "opentelemetry.instrumentation.psycopg.AsyncCursor", + "opentelemetry.instrumentation.psycopg.psycopg.AsyncCursor", MockAsyncCursor, ) self.connection_mock = mock.patch("psycopg.connect", MockConnection) From d78970c137ec32d82bef06d2bbbe095d1f39fb2b Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 12 Dec 2024 10:10:25 +0100 Subject: [PATCH 4/8] Add psycopg.Connection to nitpick --- docs/nitpick-exceptions.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/nitpick-exceptions.ini b/docs/nitpick-exceptions.ini index 4b1b06f95b..479cf84d16 100644 --- a/docs/nitpick-exceptions.ini +++ b/docs/nitpick-exceptions.ini @@ -39,6 +39,7 @@ py-class= callable Consumer confluent_kafka.Message + psycopg.Connection any= ; API From ab3790f7b4fee8577740e6ffcdd7eee6a9c7bb7e Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 12 Dec 2024 10:15:58 +0100 Subject: [PATCH 5/8] Add py.typed --- docs/nitpick-exceptions.ini | 1 - .../instrumentation/psycopg/__init__.py | 16 +++++++++------- .../src/opentelemetry/instrumentation/py.typed | 0 3 files changed, 9 insertions(+), 8 deletions(-) create mode 100644 instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/py.typed diff --git a/docs/nitpick-exceptions.ini b/docs/nitpick-exceptions.ini index 479cf84d16..4b1b06f95b 100644 --- a/docs/nitpick-exceptions.ini +++ b/docs/nitpick-exceptions.ini @@ -39,7 +39,6 @@ py-class= callable Consumer confluent_kafka.Message - psycopg.Connection any= ; API 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 7668fa806e..ac0634e835 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py @@ -118,8 +118,10 @@ _logger = logging.getLogger(__name__) _OTEL_CURSOR_FACTORY_KEY = "_otel_orig_cursor_factory" -Connection = TypeVar("Connection", psycopg.Connection, psycopg.AsyncConnection) -Cursor = TypeVar("Cursor", psycopg.Cursor, psycopg.AsyncCursor) +ConnectionT = TypeVar( + "ConnectionT", psycopg.Connection, psycopg.AsyncConnection +) +CursorT = TypeVar("CursorT", psycopg.Cursor, psycopg.AsyncCursor) class PsycopgInstrumentor(BaseInstrumentor): @@ -195,8 +197,8 @@ def _uninstrument(self, **kwargs: Any): # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql @staticmethod def instrument_connection( - connection: Connection, tracer_provider: TracerProvider | None = None - ) -> Connection: + connection: ConnectionT, tracer_provider: TracerProvider | None = None + ) -> ConnectionT: if not hasattr(connection, "_is_instrumented_by_opentelemetry"): connection._is_instrumented_by_opentelemetry = False @@ -216,7 +218,7 @@ def instrument_connection( # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql @staticmethod - def uninstrument_connection(connection: Connection) -> Connection: + def uninstrument_connection(connection: ConnectionT) -> ConnectionT: connection.cursor_factory = getattr( connection, _OTEL_CURSOR_FACTORY_KEY, None ) @@ -264,7 +266,7 @@ async def wrapped_connection( class CursorTracer(dbapi.CursorTracer): - def get_operation_name(self, cursor: Cursor, args: list[Any]) -> str: + def get_operation_name(self, cursor: CursorT, args: list[Any]) -> str: if not args: return "" @@ -279,7 +281,7 @@ def get_operation_name(self, cursor: Cursor, args: list[Any]) -> str: return "" - def get_statement(self, cursor: Cursor, args: list[Any]) -> str: + def get_statement(self, cursor: CursorT, args: list[Any]) -> str: if not args: return "" diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/py.typed b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/py.typed new file mode 100644 index 0000000000..e69de29bb2 From 26815a8d2d683ec3dd732a82c94096806efab3ac Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 12 Dec 2024 10:19:43 +0100 Subject: [PATCH 6/8] add psycopg to nitpick again --- docs/nitpick-exceptions.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/nitpick-exceptions.ini b/docs/nitpick-exceptions.ini index 4b1b06f95b..479cf84d16 100644 --- a/docs/nitpick-exceptions.ini +++ b/docs/nitpick-exceptions.ini @@ -39,6 +39,7 @@ py-class= callable Consumer confluent_kafka.Message + psycopg.Connection any= ; API From edc0acc5777693445c6c085169a915c02923c6f8 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 12 Dec 2024 10:24:09 +0100 Subject: [PATCH 7/8] add psycopg to nitpick again --- docs/nitpick-exceptions.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/nitpick-exceptions.ini b/docs/nitpick-exceptions.ini index 479cf84d16..1a029e00a1 100644 --- a/docs/nitpick-exceptions.ini +++ b/docs/nitpick-exceptions.ini @@ -40,6 +40,7 @@ py-class= Consumer confluent_kafka.Message psycopg.Connection + psycopg.AsyncConnection any= ; API From 2862863b25a584fc0b11c49d6b85d839fbff4d0e Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Tue, 14 Jan 2025 09:28:18 +0000 Subject: [PATCH 8/8] move py.typed to the right folder --- .../src/opentelemetry/instrumentation/{ => psycopg}/py.typed | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/{ => psycopg}/py.typed (100%) diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/py.typed b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/py.typed similarity index 100% rename from instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/py.typed rename to instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/py.typed