Skip to content

Commit

Permalink
DB-API instrument_connection accepts optional connect_module (#3027)
Browse files Browse the repository at this point in the history
* DbApi instrument_connection accepts optional connect_module

* Changelog

* Add test

* Adjust tests
  • Loading branch information
tammy-baylis-swi authored Nov 20, 2024
1 parent d0cbf8e commit 4e992dd
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 0 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 4e992dd

Please sign in to comment.