Skip to content

Commit

Permalink
Add docstrings for instrument_connection() for psycopg, psycopg2, and…
Browse files Browse the repository at this point in the history
… pymysql
  • Loading branch information
beijiez committed Dec 13, 2024
1 parent db31f28 commit 74b3636
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@
cursor.close()
cnx.close()
instrumented_connection = MySQLClientInstrumentor.instrument_connection(
instrumented_cnx = MySQLClientInstrumentor.instrument_connection(
cnx,
enable_commenter=True,
commenter_options={
Expand All @@ -65,11 +65,11 @@
"driver_paramstyle": False
}
)
cursor = instrumented_connection.cursor()
cursor = instrumented_cnx.cursor()
cursor.execute("INSERT INTO test (testField) VALUES (123)"
instrumented_connection.commit()
instrumented_cnx.commit()
cursor.close()
instrumented_connection.close()
instrumented_cnx.close()
For example,
::
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -92,11 +92,18 @@
PsycopgInstrumentor().instrument()
cnx = psycopg.connect(database='Database')
cursor = cnx.cursor()
cursor.execute("INSERT INTO test (testField) VALUES (123)")
cursor.close()
cnx.close()
instrumented_cnx = instrument_connection(cnx)
cursor = instrumented_cnx.cursor()
cursor.execute("INSERT INTO test (testField) VALUES (123)")
cursor.close()
instrumented_cnx.close()
API
---
"""
Expand Down Expand Up @@ -196,6 +203,18 @@ def _uninstrument(self, **kwargs):
# TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
@staticmethod
def instrument_connection(connection, tracer_provider=None):
"""Enable instrumentation in a psycopg connection.
Args:
connection: psycopg.Connection
The psycopg connection object to be instrumented.
tracer_provider: opentelemetry.trace.TracerProvider, optional
The TracerProvider to use for instrumentation. If not provided,
the global TracerProvider will be used.
Returns:
An instrumented psycopg connection object.
"""
if not hasattr(connection, "_is_instrumented_by_opentelemetry"):
connection._is_instrumented_by_opentelemetry = False

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -92,11 +92,17 @@
Psycopg2Instrumentor().instrument()
cnx = psycopg2.connect(database='Database')
cursor = cnx.cursor()
cursor.execute("INSERT INTO test (testField) VALUES (123)")
cursor.close()
cnx.close()
instrumented_cnx = Psycopg2Instrumentor.instrument_connection(cnx)
cursor = instrumented_cnx.cursor()
cursor.execute("INSERT INTO test (testField) VALUES (123)")
cursor.close()
instrumented_cnx.close()
API
---
"""
Expand Down Expand Up @@ -160,6 +166,19 @@ def _uninstrument(self, **kwargs):
# TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
@staticmethod
def instrument_connection(connection, tracer_provider=None):
"""Enable instrumentation in a psycopg2 connection.
Args:
connection: psycopg2.extensions.connection
The psycopg2 connection object to be instrumented.
tracer_provider: opentelemetry.trace.TracerProvider, optional
The TracerProvider to use for instrumentation. If not specified,
the global TracerProvider will be used.
Returns:
An instrumented psycopg2 connection object.
"""

if not hasattr(connection, "_is_instrumented_by_opentelemetry"):
connection._is_instrumented_by_opentelemetry = False

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,20 @@
cnx.commit()
cursor.close()
cnx.close()
instrumented_cnx = PyMySQLInstrumentor().instrument_connection(
cnx,
enable_commenter=True,
commenter_options={
"db_driver": True,
"mysql_client_version": True
}
)
cursor = instrumented_cnx.cursor()
cursor.execute("INSERT INTO test (testField) VALUES (123)"
instrumented_cnx.commit()
cursor.close()
instrumented_cnx.close()
For example,
Expand Down Expand Up @@ -165,10 +179,20 @@ def instrument_connection(
"""Enable instrumentation in a PyMySQL connection.
Args:
connection: The connection to instrument.
tracer_provider: The optional tracer provider to use. If omitted
the current globally configured one is used.
connection (pymysql.Connection):
The existing PyMySQL connection instance that needs to be instrumented.
This connection was typically created using `pymysql.connect()` and is wrapped with OpenTelemetry tracing.
tracer_provider (TracerProvider, optional):
An optional `TracerProvider` instance that specifies which tracer provider should be used.
If not provided, the globally configured OpenTelemetry tracer provider is automatically applied.
enable_commenter (bool, optional):
A flag to enable the SQL Commenter feature. If `True`, query logs will be enriched with additional
contextual metadata (e.g., database version, traceparent IDs, driver information).
commenter_options (dict, optional):
A dictionary containing configuration options for the SQL Commenter feature.
You can specify various options, such as enabling driver information, database version logging,
traceparent propagation, and other customizable metadata enhancements.
See *SQLCommenter Configurations* above for more information.
Returns:
An instrumented connection.
"""
Expand Down

0 comments on commit 74b3636

Please sign in to comment.