opentelemetry-instrumentation-dbapi: fix libpq version detection (#3796)

* opentelemetry-instrumentation-dbapi: fix libpq version detection

This fixes a crash in the detection of libpq version in the psycopg module
when we wrap a module that is not the one containing the pq module that
has the version.

* Fix tests

* More fixes

* Please pylint

* Add changelog

* Fix typo

* Fix test name

* Add docker tests for psycopg

* Simplify psycopg handling

* Fix non root module tests

Author: xrmx

CHANGELOG.md

@@ -13,12 +13,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `opentelemetry-instrumentation-dbapi`: fix crash retrieving libpq version when enabling commenter with psycopg
+  ([#3796](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3796))
+
 ### Added
 - `opentelemetry-instrumentation`: botocore: Add support for AWS Secrets Manager semantic convention attribute
   ([#3765](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3765))
 - Add `rstcheck` to pre-commit to stop introducing invalid RST
   ([#3777](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3777))
-
 - `opentelemetry-exporter-credential-provider-gcp`: create this package which provides support for supplying your machine's Application Default Credentials (https://cloud.google.com/docs/authentication/application-default-credentials) to the OTLP Exporters created automatically by OpenTelemetry Python's auto instrumentation. These credentials authorize OTLP traces to be sent to `telemetry.googleapis.com`. 
 [#3766](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3766).
 

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

@@ -350,11 +350,24 @@ def calculate_commenter_data(self) -> dict[str, Any]:
         }
 
         if self.database_system == "postgresql":
-            if hasattr(self.connect_module, "__libpq_version__"):
-                libpq_version = self.connect_module.__libpq_version__
-            else:
-                libpq_version = self.connect_module.pq.__build_version__
-            commenter_data.update({"libpq_version": libpq_version})
+            libpq_version = None
+            # psycopg
+            try:
+                libpq_version = self.connect_module.pq.version()
+            except AttributeError:
+                pass
+
+            # psycopg2
+            if libpq_version is None:
+                # this the libpq version the client has been built against
+                libpq_version = getattr(
+                    self.connect_module, "__libpq_version__", None
+                )
+
+            # we instrument psycopg modules that are not the root one, in that case you
+            # won't get the libpq_version
+            if libpq_version is not None:
+                commenter_data.update({"libpq_version": libpq_version})
         elif self.database_system == "mysql":
             mysqlc_version = ""
             if db_driver == "MySQLdb":

instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py

@@ -279,7 +279,7 @@ def test_executemany_comment(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
-        connect_module.__libpq_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -312,7 +312,7 @@ def test_executemany_comment_stmt_enabled(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
-        connect_module.__libpq_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -420,7 +420,7 @@ def test_executemany_comment_stmt_enabled_matches_db_statement_attribute(
     ):
         connect_module = mock.MagicMock()
         connect_module.__version__ = mock.MagicMock()
-        connect_module.__libpq_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -456,12 +456,12 @@ def test_executemany_comment_stmt_enabled_matches_db_statement_attribute(
         ).group()
         self.assertEqual(cursor_span_id, db_statement_span_id)
 
-    def test_compatible_build_version_psycopg_psycopg2_libpq(self):
+    def test_compatible_build_version_psycopg2_libpq(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
-        connect_module.pq = mock.MagicMock()
-        connect_module.pq.__build_version__ = 123
+        connect_module.pq.version.side_effect = AttributeError
+        connect_module.__libpq_version__ = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -490,14 +490,47 @@ def test_compatible_build_version_psycopg_psycopg2_libpq(self):
             "Select 1;",
         )
 
-    def test_compatible_build_version_psycopg_psycopg2_libpq_stmt_enabled(
+    def test_compatible_build_version_psycopg_libpq(self):
+        connect_module = mock.MagicMock()
+        connect_module.__name__ = "test"
+        connect_module.__version__ = mock.MagicMock()
+        connect_module.pq.version.return_value = 123
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+
+        db_integration = dbapi.DatabaseApiIntegration(
+            "instrumenting_module_test_name",
+            "postgresql",
+            enable_commenter=True,
+            commenter_options={"db_driver": False, "dbapi_level": False},
+            connect_module=connect_module,
+        )
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+        self.assertEqual(
+            span.attributes[SpanAttributes.DB_STATEMENT],
+            "Select 1;",
+        )
+
+    def test_no_libpq_version_when_we_dont_instrument_the_root_module(
         self,
     ):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
-        connect_module.pq = mock.MagicMock()
-        connect_module.pq.__build_version__ = 123
+        connect_module.pq.version.side_effect = AttributeError
+        connect_module.__libpq_version__ = None
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -517,20 +550,21 @@ def test_compatible_build_version_psycopg_psycopg2_libpq_stmt_enabled(
         cursor.executemany("Select 1;")
         self.assertRegex(
             cursor.query,
-            r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+            r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
         span = spans_list[0]
         self.assertRegex(
             span.attributes[SpanAttributes.DB_STATEMENT],
-            r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+            r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
 
     def test_executemany_psycopg2_integration_comment(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "psycopg2"
         connect_module.__version__ = "1.2.3"
+        connect_module.pq.version.side_effect = AttributeError
         connect_module.__libpq_version__ = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
@@ -564,6 +598,7 @@ def test_executemany_psycopg2_integration_comment_stmt_enabled(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "psycopg2"
         connect_module.__version__ = "1.2.3"
+        connect_module.pq.version.side_effect = AttributeError
         connect_module.__libpq_version__ = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
@@ -598,8 +633,7 @@ def test_executemany_psycopg_integration_comment(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "psycopg"
         connect_module.__version__ = "1.2.3"
-        connect_module.pq = mock.MagicMock()
-        connect_module.pq.__build_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -632,8 +666,7 @@ def test_executemany_psycopg_integration_comment_stmt_enabled(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "psycopg"
         connect_module.__version__ = "1.2.3"
-        connect_module.pq = mock.MagicMock()
-        connect_module.pq.__build_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -887,8 +920,7 @@ def test_executemany_pymysql_integration_comment_stmt_enabled(self):
     def test_executemany_flask_integration_comment(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "test"
-        connect_module.__version__ = mock.MagicMock()
-        connect_module.__libpq_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -904,7 +936,7 @@ def test_executemany_flask_integration_comment(self):
         sqlcommenter_context = context.set_value(
             "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {"flask": 1}, current_context
         )
-        context.attach(sqlcommenter_context)
+        token = context.attach(sqlcommenter_context)
 
         mock_connection = db_integration.wrapped_connection(
             mock_connect, {}, {}
@@ -923,16 +955,13 @@ def test_executemany_flask_integration_comment(self):
             "Select 1;",
         )
 
-        clear_context = context.set_value(
-            "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {}, current_context
-        )
-        context.attach(clear_context)
+        context.detach(token)
 
     def test_executemany_flask_integration_comment_stmt_enabled(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
-        connect_module.__libpq_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = 123
         connect_module.threadsafety = 123
         connect_module.paramstyle = "test"
@@ -949,7 +978,7 @@ def test_executemany_flask_integration_comment_stmt_enabled(self):
         sqlcommenter_context = context.set_value(
             "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {"flask": 1}, current_context
         )
-        context.attach(sqlcommenter_context)
+        token = context.attach(sqlcommenter_context)
 
         mock_connection = db_integration.wrapped_connection(
             mock_connect, {}, {}
@@ -968,10 +997,7 @@ def test_executemany_flask_integration_comment_stmt_enabled(self):
             r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',flask=1,libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
 
-        clear_context = context.set_value(
-            "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {}, current_context
-        )
-        context.attach(clear_context)
+        context.detach(token)
 
     def test_callproc(self):
         db_integration = dbapi.DatabaseApiIntegration(
@@ -1096,7 +1122,7 @@ def test_non_string_sql_conversion(self):
         connect_module = mock.MagicMock()
         connect_module.__name__ = "test"
         connect_module.__version__ = "1.0"
-        connect_module.__libpq_version__ = 123
+        connect_module.pq.version.return_value = 123
         connect_module.apilevel = "2.0"
         connect_module.threadsafety = 1
         connect_module.paramstyle = "test"

tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import psycopg
 import psycopg2
 from test_psycopg_functional import (
     POSTGRES_DB_NAME,
@@ -21,11 +22,12 @@
     POSTGRES_USER,
 )
 
+from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor
 from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
 from opentelemetry.test.test_base import TestBase
 
 
-class TestFunctionalPsycopg(TestBase):
+class TestFunctionalPsycopg2(TestBase):
     def setUp(self):
         super().setUp()
         self._tracer = self.tracer_provider.get_tracer(__name__)
@@ -52,3 +54,49 @@ def test_commenter_enabled(self):
             self._cursor.query.decode("ascii"),
             r"SELECT  1 /\*db_driver='psycopg2(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
+
+
+class TestFunctionalPsycopg(TestBase):
+    def setUp(self):
+        super().setUp()
+        PsycopgInstrumentor().instrument(enable_commenter=True)
+
+    def tearDown(self):
+        PsycopgInstrumentor().uninstrument()
+        super().tearDown()
+
+    def test_commenter_enabled_root_module(self):
+        connection = psycopg.connect(
+            dbname=POSTGRES_DB_NAME,
+            user=POSTGRES_USER,
+            password=POSTGRES_PASSWORD,
+            host=POSTGRES_HOST,
+            port=POSTGRES_PORT,
+        )
+        cursor = connection.cursor()
+        cursor.execute("SELECT  1;")
+        self.assertRegex(
+            cursor._query.query.decode("ascii"),
+            r"SELECT  1 /\*db_driver='psycopg(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+
+        cursor.close()
+        connection.close()
+
+    def test_commenter_enabled_not_root_module(self):
+        connection = psycopg.Connection.connect(
+            dbname=POSTGRES_DB_NAME,
+            user=POSTGRES_USER,
+            password=POSTGRES_PASSWORD,
+            host=POSTGRES_HOST,
+            port=POSTGRES_PORT,
+        )
+        cursor = connection.cursor()
+        cursor.execute("SELECT  1;")
+        self.assertRegex(
+            cursor._query.query.decode("ascii"),
+            r"SELECT  1 /\*db_driver='Connection%%3Aunknown',dbapi_level='\d.\d',dbapi_threadsafety='unknown',driver_paramstyle='unknown',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+
+        cursor.close()
+        connection.close()