Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
…y-python-contrib into fix/1477
  • Loading branch information
aryabharat committed Jan 20, 2025
2 parents ba3b51a + e90eefe commit cc7332d
Show file tree
Hide file tree
Showing 5 changed files with 10 additions and 17 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -488,6 +488,9 @@ def _update_args_with_added_sql_comment(self, args, cursor) -> tuple:
args_list = list(args)
self._capture_mysql_version(cursor)
commenter_data = self._get_commenter_data()
# Convert sql statement to string, handling psycopg2.sql.Composable object
if hasattr(args_list[0], "as_string"):
args_list[0] = args_list[0].as_string(cursor.connection)
statement = _add_sql_comment(args_list[0], **commenter_data)
args_list[0] = statement
args = tuple(args_list)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,10 @@ def __call__(self, execute: Type[T], sql, params, many, context) -> T:
db_driver = context["connection"].settings_dict.get("ENGINE", "")
resolver_match = self.request.resolver_match

# Convert sql statement to string, handling psycopg2.sql.Composable object
if hasattr(sql, "as_string"):
sql = sql.as_string(context["connection"])

sql = _add_sql_comment(
sql,
# Information about the controller.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,9 @@ def _before_cur_exec(
commenter_data = self._get_commenter_data(conn)

if self.enable_attribute_commenter:
# Convert sql statement to string, handling psycopg2.sql.Composable object
if hasattr(statement, "as_string"):
statement = statement.as_string(conn)
# sqlcomment is added to executed query and db.statement span attribute
statement = _add_sql_comment(
statement, **commenter_data
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ def _add_sql_comment(sql, **meta) -> str:
"""
meta.update(**_add_framework_tags())
comment = _generate_sql_comment(**meta)
# converting to str to handle any type errors
sql = str(sql)
sql = sql.rstrip()
if sql[-1] == ";":
sql = sql[:-1] + comment + ";"
Expand Down
15 changes: 0 additions & 15 deletions opentelemetry-instrumentation/tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,21 +208,6 @@ def test_add_sql_comments_without_comments(self):

self.assertEqual(commented_sql_without_semicolon, "Select 1")

def test_psycopg2_sql_composable(self):
"""Test handling of psycopg2.sql.Composable input"""
# Mock psycopg2.sql.Composable object
class MockComposable:
def __str__(self):
return "SELECT * FROM table_name"

sql_query = MockComposable()
comments = {"trace_id": "abc123"}

result = _add_sql_comment(sql_query, **comments)
expected = "SELECT * FROM table_name /*trace_id='abc123'*/"

self.assertEqual(result, expected)

def test_is_instrumentation_enabled_by_default(self):
self.assertTrue(is_instrumentation_enabled())
self.assertTrue(is_http_instrumentation_enabled())
Expand Down

0 comments on commit cc7332d

Please sign in to comment.