From ea17a2a3621718c6a531c5fcbead612f24f64127 Mon Sep 17 00:00:00 2001
From: tammy-baylis-swi <tammy.baylis@solarwinds.com>
Date: Wed, 16 Oct 2024 16:55:49 -0700
Subject: [PATCH] Add pymysql sqlcomment support

---
 .../instrumentation/pymysql/__init__.py       | 77 ++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/instrumentation/opentelemetry-instrumentation-pymysql/src/opentelemetry/instrumentation/pymysql/__init__.py b/instrumentation/opentelemetry-instrumentation-pymysql/src/opentelemetry/instrumentation/pymysql/__init__.py
index c48a5b6b3d..459e0b1e87 100644
--- a/instrumentation/opentelemetry-instrumentation-pymysql/src/opentelemetry/instrumentation/pymysql/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-pymysql/src/opentelemetry/instrumentation/pymysql/__init__.py
@@ -36,6 +36,70 @@
     cursor.close()
     cnx.close()
 
+SQLCOMMENTER
+*****************************************
+You can optionally configure PyMySQL instrumentation to enable sqlcommenter which enriches
+the query with contextual information.
+
+.. code:: python
+
+    import MySQLdb
+    from opentelemetry.instrumentation.pymysql import PyMySQLInstrumentor
+
+    PyMySQLInstrumentor().instrument(enable_commenter=True, commenter_options={})
+
+    cnx = MySQLdb.connect(database="MySQL_Database")
+    cursor = cnx.cursor()
+    cursor.execute("INSERT INTO test (testField) VALUES (123)"
+    cnx.commit()
+    cursor.close()
+    cnx.close()
+
+For example,
+::
+   Invoking cursor.execute("INSERT INTO test (testField) VALUES (123)") will lead to sql query "INSERT INTO test (testField) VALUES (123)" but when SQLCommenter is enabled
+   the query will get appended with some configurable tags like "INSERT INTO test (testField) VALUES (123) /*tag=value*/;"
+
+SQLCommenter Configurations
+***************************
+We can configure the tags to be appended to the sqlquery log by adding configuration inside commenter_options(default:{}) keyword
+
+db_driver = True(Default) or False
+
+For example,
+::
+Enabling this flag will add MySQLdb and its version, e.g. /*MySQLdb%%3A1.2.3*/
+
+dbapi_threadsafety = True(Default) or False
+
+For example,
+::
+Enabling this flag will add threadsafety /*dbapi_threadsafety=2*/
+
+dbapi_level = True(Default) or False
+
+For example,
+::
+Enabling this flag will add dbapi_level /*dbapi_level='2.0'*/
+
+mysql_client_version = True(Default) or False
+
+For example,
+::
+Enabling this flag will add mysql_client_version /*mysql_client_version='123'*/
+
+driver_paramstyle = True(Default) or False
+
+For example,
+::
+Enabling this flag will add driver_paramstyle /*driver_paramstyle='pyformat'*/
+
+opentelemetry_values = True(Default) or False
+
+For example,
+::
+Enabling this flag will add traceparent values /*traceparent='00-03afa25236b8cd948fa853d67038ac79-405ff022e8247c46-01'*/
+
 API
 ---
 """
@@ -67,6 +131,8 @@ def _instrument(self, **kwargs):
         https://github.com/PyMySQL/PyMySQL/
         """
         tracer_provider = kwargs.get("tracer_provider")
+        enable_sqlcommenter = kwargs.get("enable_commenter", False)
+        commenter_options = kwargs.get("commenter_options", {})
 
         dbapi.wrap_connect(
             __name__,
@@ -76,6 +142,8 @@ def _instrument(self, **kwargs):
             _CONNECTION_ATTRIBUTES,
             version=__version__,
             tracer_provider=tracer_provider,
+            enable_commenter=enable_sqlcommenter,
+            commenter_options=commenter_options,
         )
 
     def _uninstrument(self, **kwargs):
@@ -83,7 +151,12 @@ def _uninstrument(self, **kwargs):
         dbapi.unwrap_connect(pymysql, "connect")
 
     @staticmethod
-    def instrument_connection(connection, tracer_provider=None):
+    def instrument_connection(
+        connection,
+        tracer_provider=None,
+        enable_commenter=None,
+        commenter_options=None,
+    ):
         """Enable instrumentation in a PyMySQL connection.
 
         Args:
@@ -102,6 +175,8 @@ def instrument_connection(connection, tracer_provider=None):
             _CONNECTION_ATTRIBUTES,
             version=__version__,
             tracer_provider=tracer_provider,
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
         )
 
     @staticmethod