Skip to content

Commit

Permalink
dbapi.instrument_connection optional kwargs db_api_integration_factor…
Browse files Browse the repository at this point in the history
…y, get_cnx_proxy
  • Loading branch information
tammy-baylis-swi committed Nov 26, 2024
1 parent 94bf04a commit 19bfd7c
Showing 1 changed file with 14 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,8 @@ def instrument_connection(
enable_commenter: bool = False,
commenter_options: dict = None,
connect_module: typing.Callable[..., typing.Any] = None,
db_api_integration_factory: typing.ClassVar = None,
get_cnx_proxy: typing.Callable[..., typing.Any] = None,
):
"""Enable instrumentation in a database connection.
Expand All @@ -210,6 +212,10 @@ def instrument_connection(
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.
db_api_integration_factory: The `DatabaseApiIntegration` to use. If none is passed the
default one is used.
get_cnx_proxy: Method to get traced connextion proxy. If none is passed the
default one is used.
Returns:
An instrumented connection.
Expand All @@ -218,7 +224,11 @@ def instrument_connection(
_logger.warning("Connection already instrumented")
return connection

db_integration = DatabaseApiIntegration(
db_api_integration_factory = (
db_api_integration_factory or DatabaseApiIntegration
)

db_integration = db_api_integration_factory(
name,
database_system,
connection_attributes=connection_attributes,
Expand All @@ -230,7 +240,9 @@ def instrument_connection(
connect_module=connect_module,
)
db_integration.get_connection_attributes(connection)
return get_traced_connection_proxy(connection, db_integration)

get_cnx_proxy = get_cnx_proxy or get_traced_connection_proxy
return get_cnx_proxy(connection, db_integration)


def uninstrument_connection(connection):
Expand Down

0 comments on commit 19bfd7c

Please sign in to comment.