* refactor on package name psycopg insteead of psycopg3

open-telemetry · Mar 1, 2024 · 953a113 · 953a113
1 parent 057504a
commit 953a113
Show file tree

Hide file tree

Showing 16 changed files with 303 additions and 1,161 deletions.
diff --git a/.github/workflows/instrumentations_1.yml b/.github/workflows/instrumentations_1.yml
@@ -31,7 +31,6 @@ jobs:
           - "richconsole"
           - "psycopg"
           - "prometheus-remote-write"
-          - "psycopg3"
           - "sdkextension-aws"
           - "propagator-aws-xray"
           - "propagator-ot-trace"

diff --git a/docs-requirements.txt b/docs-requirements.txt
@@ -36,7 +36,7 @@ kafka-python>=2.0,<3.0
 mysql-connector-python~=8.0
 mysqlclient~=2.1.1
 psutil>=5
-psycopg>=3.1.17
+psycopg~=3.1.17
 pika>=0.12.0
 pymongo~=3.1
 PyMySQL~=0.9.3

diff --git a/docs/instrumentation/psycopg3/psycopg3.rst b/docs/instrumentation/psycopg3/psycopg3.rst
diff --git a/instrumentation/README.md b/instrumentation/README.md
@@ -31,7 +31,6 @@
 | [opentelemetry-instrumentation-pika](./opentelemetry-instrumentation-pika) | pika >= 0.12.0 | No
 | [opentelemetry-instrumentation-psycopg](./opentelemetry-instrumentation-psycopg) | psycopg >= 3.1.0 | No
 | [opentelemetry-instrumentation-psycopg2](./opentelemetry-instrumentation-psycopg2) | psycopg2 >= 2.7.3.1 | No
-| [opentelemetry-instrumentation-psycopg3](./opentelemetry-instrumentation-psycopg3) | psycopg >= 3.1.12 | No
 | [opentelemetry-instrumentation-pymemcache](./opentelemetry-instrumentation-pymemcache) | pymemcache >= 1.3.5, < 5 | No
 | [opentelemetry-instrumentation-pymongo](./opentelemetry-instrumentation-pymongo) | pymongo >= 3.1, < 5.0 | No
 | [opentelemetry-instrumentation-pymysql](./opentelemetry-instrumentation-pymysql) | PyMySQL < 2 | No

diff --git a/...entelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py b/...entelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py
@@ -106,6 +106,7 @@
 from typing import Collection
 
 import psycopg
+from psycopg import AsyncCursor as pg_async_cursor
 from psycopg import Cursor as pg_cursor  # pylint: disable=no-name-in-module
 from psycopg.sql import Composed  # pylint: disable=no-name-in-module
 
@@ -151,9 +152,36 @@ def _instrument(self, **kwargs):
             commenter_options=commenter_options,
         )
 
+        dbapi.wrap_connect(
+            __name__,
+            psycopg.Connection,
+            "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,
+            "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.Connection, "connect")
+        dbapi.unwrap_connect(psycopg.AsyncConnection, "connect")
 
     # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
     @staticmethod
@@ -204,6 +232,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 +307,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__,
+            Psycopg3Instrumentor._DATABASE_SYSTEM,
+            connection_attributes=Psycopg3Instrumentor._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