From 19bfd7cff91980d257ff691fb9f54819a78905ea Mon Sep 17 00:00:00 2001
From: tammy-baylis-swi <tammy.baylis@solarwinds.com>
Date: Tue, 26 Nov 2024 10:48:51 -0800
Subject: [PATCH] dbapi.instrument_connection optional kwargs
 db_api_integration_factory, get_cnx_proxy

---
 .../instrumentation/dbapi/__init__.py            | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py
index b80112c204..c2b63e59fe 100644
--- a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py
@@ -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.
 
@@ -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.
@@ -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,
@@ -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):