DB-API instrument_connection accepts optional connect_module (open-te…

…lemetry#3027)

* DbApi instrument_connection accepts optional connect_module

* Changelog

* Add test

* Adjust tests

CHANGELOG.md

@@ -21,6 +21,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#2937](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2937))
 - `opentelemetry-instrumentation-dbapi` Add sqlcomment to `db.statement` attribute
   ([#2935](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2935))
+- `opentelemetry-instrumentation-dbapi` instrument_connection accepts optional connect_module
+  ([#3027](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3027))
 
 ### Fixed
 

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

@@ -190,6 +190,7 @@ def instrument_connection(
     capture_parameters: bool = False,
     enable_commenter: bool = False,
     commenter_options: dict = None,
+    connect_module: typing.Callable[..., typing.Any] = None,
 ):
     """Enable instrumentation in a database connection.
 
@@ -204,6 +205,7 @@ def instrument_connection(
         capture_parameters: Configure if db.statement.parameters should be captured.
         enable_commenter: Flag to enable/disable sqlcommenter.
         commenter_options: Configurations for tags to be appended at the sql query.
+        connect_module: Module name where connect method is available.
 
     Returns:
         An instrumented connection.
@@ -221,6 +223,7 @@ def instrument_connection(
         capture_parameters=capture_parameters,
         enable_commenter=enable_commenter,
         commenter_options=commenter_options,
+        connect_module=connect_module,
     )
     db_integration.get_connection_attributes(connection)
     return get_traced_connection_proxy(connection, db_integration)

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

@@ -592,6 +592,43 @@ def test_instrument_connection(self):
         connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
         self.assertIs(connection2.__wrapped__, connection)
 
+    @mock.patch("opentelemetry.instrumentation.dbapi.DatabaseApiIntegration")
+    def test_instrument_connection_kwargs_defaults(self, mock_dbapiint):
+        dbapi.instrument_connection(self.tracer, mock.Mock(), "foo")
+        kwargs = mock_dbapiint.call_args[1]
+        self.assertEqual(kwargs["connection_attributes"], None)
+        self.assertEqual(kwargs["version"], "")
+        self.assertEqual(kwargs["tracer_provider"], None)
+        self.assertEqual(kwargs["capture_parameters"], False)
+        self.assertEqual(kwargs["enable_commenter"], False)
+        self.assertEqual(kwargs["commenter_options"], None)
+        self.assertEqual(kwargs["connect_module"], None)
+
+    @mock.patch("opentelemetry.instrumentation.dbapi.DatabaseApiIntegration")
+    def test_instrument_connection_kwargs_provided(self, mock_dbapiint):
+        mock_tracer_provider = mock.MagicMock()
+        mock_connect_module = mock.MagicMock()
+        dbapi.instrument_connection(
+            self.tracer,
+            mock.Mock(),
+            "foo",
+            connection_attributes={"foo": "bar"},
+            version="test",
+            tracer_provider=mock_tracer_provider,
+            capture_parameters=True,
+            enable_commenter=True,
+            commenter_options={"foo": "bar"},
+            connect_module=mock_connect_module,
+        )
+        kwargs = mock_dbapiint.call_args[1]
+        self.assertEqual(kwargs["connection_attributes"], {"foo": "bar"})
+        self.assertEqual(kwargs["version"], "test")
+        self.assertIs(kwargs["tracer_provider"], mock_tracer_provider)
+        self.assertEqual(kwargs["capture_parameters"], True)
+        self.assertEqual(kwargs["enable_commenter"], True)
+        self.assertEqual(kwargs["commenter_options"], {"foo": "bar"})
+        self.assertIs(kwargs["connect_module"], mock_connect_module)
+
     def test_uninstrument_connection(self):
         connection = mock.Mock()
         # Set connection.database to avoid a failure because mock can't