Skip to content

Commit 4e992dd

Browse files
DB-API instrument_connection accepts optional connect_module (#3027)
* DbApi instrument_connection accepts optional connect_module * Changelog * Add test * Adjust tests
1 parent d0cbf8e commit 4e992dd

File tree

3 files changed

+42
-0
lines changed

3 files changed

+42
-0
lines changed

Diff for: CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
2121
([#2937](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2937))
2222
- `opentelemetry-instrumentation-dbapi` Add sqlcomment to `db.statement` attribute
2323
([#2935](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2935))
24+
- `opentelemetry-instrumentation-dbapi` instrument_connection accepts optional connect_module
25+
([#3027](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3027))
2426

2527
### Fixed
2628

Diff for: instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py

+3
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,7 @@ def instrument_connection(
190190
capture_parameters: bool = False,
191191
enable_commenter: bool = False,
192192
commenter_options: dict = None,
193+
connect_module: typing.Callable[..., typing.Any] = None,
193194
):
194195
"""Enable instrumentation in a database connection.
195196
@@ -204,6 +205,7 @@ def instrument_connection(
204205
capture_parameters: Configure if db.statement.parameters should be captured.
205206
enable_commenter: Flag to enable/disable sqlcommenter.
206207
commenter_options: Configurations for tags to be appended at the sql query.
208+
connect_module: Module name where connect method is available.
207209
208210
Returns:
209211
An instrumented connection.
@@ -221,6 +223,7 @@ def instrument_connection(
221223
capture_parameters=capture_parameters,
222224
enable_commenter=enable_commenter,
223225
commenter_options=commenter_options,
226+
connect_module=connect_module,
224227
)
225228
db_integration.get_connection_attributes(connection)
226229
return get_traced_connection_proxy(connection, db_integration)

Diff for: instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py

+37
Original file line numberDiff line numberDiff line change
@@ -592,6 +592,43 @@ def test_instrument_connection(self):
592592
connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
593593
self.assertIs(connection2.__wrapped__, connection)
594594

595+
@mock.patch("opentelemetry.instrumentation.dbapi.DatabaseApiIntegration")
596+
def test_instrument_connection_kwargs_defaults(self, mock_dbapiint):
597+
dbapi.instrument_connection(self.tracer, mock.Mock(), "foo")
598+
kwargs = mock_dbapiint.call_args[1]
599+
self.assertEqual(kwargs["connection_attributes"], None)
600+
self.assertEqual(kwargs["version"], "")
601+
self.assertEqual(kwargs["tracer_provider"], None)
602+
self.assertEqual(kwargs["capture_parameters"], False)
603+
self.assertEqual(kwargs["enable_commenter"], False)
604+
self.assertEqual(kwargs["commenter_options"], None)
605+
self.assertEqual(kwargs["connect_module"], None)
606+
607+
@mock.patch("opentelemetry.instrumentation.dbapi.DatabaseApiIntegration")
608+
def test_instrument_connection_kwargs_provided(self, mock_dbapiint):
609+
mock_tracer_provider = mock.MagicMock()
610+
mock_connect_module = mock.MagicMock()
611+
dbapi.instrument_connection(
612+
self.tracer,
613+
mock.Mock(),
614+
"foo",
615+
connection_attributes={"foo": "bar"},
616+
version="test",
617+
tracer_provider=mock_tracer_provider,
618+
capture_parameters=True,
619+
enable_commenter=True,
620+
commenter_options={"foo": "bar"},
621+
connect_module=mock_connect_module,
622+
)
623+
kwargs = mock_dbapiint.call_args[1]
624+
self.assertEqual(kwargs["connection_attributes"], {"foo": "bar"})
625+
self.assertEqual(kwargs["version"], "test")
626+
self.assertIs(kwargs["tracer_provider"], mock_tracer_provider)
627+
self.assertEqual(kwargs["capture_parameters"], True)
628+
self.assertEqual(kwargs["enable_commenter"], True)
629+
self.assertEqual(kwargs["commenter_options"], {"foo": "bar"})
630+
self.assertIs(kwargs["connect_module"], mock_connect_module)
631+
595632
def test_uninstrument_connection(self):
596633
connection = mock.Mock()
597634
# Set connection.database to avoid a failure because mock can't

0 commit comments

Comments
 (0)