Merge branch 'main' into enable-commenter-instrument-connection

Author: tammy-baylis-swi

docs/nitpick-exceptions.ini

@@ -41,6 +41,8 @@ py-class=
     callable
     Consumer
     confluent_kafka.Message
+    psycopg.Connection
+    psycopg.AsyncConnection
     ObjectProxy
 
 any=

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

@@ -137,28 +137,29 @@
 ---
 """
 
+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 import trace as trace_api
 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"
 
+ConnectionT = TypeVar(
+    "ConnectionT", psycopg.Connection, psycopg.AsyncConnection
+)
+CursorT = TypeVar("CursorT", psycopg.Cursor, psycopg.AsyncCursor)
+
 
 class PsycopgInstrumentor(BaseInstrumentor):
     _CONNECTION_ATTRIBUTES = {
@@ -173,7 +174,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: https://www.psycopg.org/psycopg3/docs/
         """
@@ -224,7 +225,7 @@ def _instrument(self, **kwargs):
             enable_attribute_commenter=enable_attribute_commenter,
         )
 
-    def _uninstrument(self, **kwargs):
+    def _uninstrument(self, **kwargs: Any):
         """ "Disable Psycopg instrumentation"""
         dbapi.unwrap_connect(psycopg, "connect")  # pylint: disable=no-member
         dbapi.unwrap_connect(
@@ -243,7 +244,7 @@ def instrument_connection(
         tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
         enable_commenter: bool = False,
         commenter_options: dict = None,
-        enable_attribute_commenter=None,
+        enable_attribute_commenter: bool = False,
     ):
         """Enable instrumentation of a Psycopg connection.
 
@@ -285,7 +286,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):
+    def uninstrument_connection(connection: ConnectionT) -> ConnectionT:
         connection.cursor_factory = getattr(
             connection, _OTEL_CURSOR_FACTORY_KEY, None
         )
@@ -297,9 +298,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)
@@ -315,9 +316,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)
@@ -333,7 +334,7 @@ async def wrapped_connection(
 
 
 class CursorTracer(dbapi.CursorTracer):
-    def get_operation_name(self, cursor, args):
+    def get_operation_name(self, cursor: CursorT, args: list[Any]) -> str:
         if not args:
             return ""
 
@@ -348,7 +349,7 @@ def get_operation_name(self, cursor, args):
 
         return ""
 
-    def get_statement(self, cursor, args):
+    def get_statement(self, cursor: CursorT, args: list[Any]) -> str:
         if not args:
             return ""
 
@@ -379,21 +380,21 @@ def _new_cursor_factory(
             enable_attribute_commenter=enable_attribute_commenter,
         )
 
-    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
             )
@@ -402,7 +403,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(
@@ -412,21 +415,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
             )

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",)

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

@@ -133,10 +133,10 @@ class PostgresqlIntegrationTestMixin:
     def setUp(self):
         super().setUp()
         self.cursor_mock = mock.patch(
-            "opentelemetry.instrumentation.psycopg.pg_cursor", MockCursor
+            "opentelemetry.instrumentation.psycopg.psycopg.Cursor", MockCursor
         )
         self.cursor_async_mock = mock.patch(
-            "opentelemetry.instrumentation.psycopg.pg_async_cursor",
+            "opentelemetry.instrumentation.psycopg.psycopg.AsyncCursor",
             MockAsyncCursor,
         )
         self.connection_mock = mock.patch("psycopg.connect", MockConnection)