Skip to content

Commit dfd5bbc

Browse files
tammy-baylis-switammy-baylis-swi
committed
Refactor db-api traced_execution
1 parent b8b5a67 commit dfd5bbc

File tree

2 files changed

+50
-98
lines changed

2 files changed

+50
-98
lines changed

‎instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py‎

Lines changed: 49 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,46 @@ def __init__(self, db_api_integration: DatabaseApiIntegration) -> None:
451451
self._connect_module = self._db_api_integration.connect_module
452452
self._leading_comment_remover = re.compile(r"^/\*.*?\*/")
453453

454+
def _capture_mysql_version(self, cursor) -> None:
455+
"""Lazy capture of mysql-connector client version using cursor, if applicable"""
456+
if (
457+
self._db_api_integration.database_system == "mysql"
458+
and self._db_api_integration.connect_module.__name__
459+
== "mysql.connector"
460+
and not self._db_api_integration.commenter_data[
461+
"mysql_client_version"
462+
]
463+
):
464+
self._db_api_integration.commenter_data["mysql_client_version"] = (
465+
cursor._cnx._cmysql.get_client_info()
466+
)
467+
468+
def _get_commenter_data(self) -> dict:
469+
"""Uses DB-API integration to return commenter data for sqlcomment"""
470+
commenter_data = dict(self._db_api_integration.commenter_data)
471+
if self._commenter_options.get("opentelemetry_values", True):
472+
commenter_data.update(**_get_opentelemetry_values())
473+
return {
474+
k: v
475+
for k, v in commenter_data.items()
476+
if self._commenter_options.get(k, True)
477+
}
478+
479+
def _update_args_with_added_sql_comment(self, args, cursor) -> tuple:
480+
"""Updates args with cursor info and adds sqlcomment to query statement"""
481+
try:
482+
args_list = list(args)
483+
self._capture_mysql_version(cursor)
484+
commenter_data = self._get_commenter_data()
485+
statement = _add_sql_comment(args_list[0], **commenter_data)
486+
args_list[0] = statement
487+
args = tuple(args_list)
488+
except Exception as exc: # pylint: disable=broad-except
489+
_logger.exception(
490+
"Exception while generating sql comment: %s", exc
491+
)
492+
return args
493+
454494
def _populate_span(
455495
self,
456496
span: trace_api.Span,
@@ -511,110 +551,21 @@ def traced_execution(
511551
) as span:
512552
if span.is_recording():
513553
if args and self._commenter_enabled:
514-
# sqlcomment in query and span attribute
515554
if self._enable_attribute_commenter:
516-
try:
517-
args_list = list(args)
518-
519-
# lazy capture of mysql-connector client version using cursor
520-
if (
521-
self._db_api_integration.database_system
522-
== "mysql"
523-
and self._db_api_integration.connect_module.__name__
524-
== "mysql.connector"
525-
and not self._db_api_integration.commenter_data[
526-
"mysql_client_version"
527-
]
528-
):
529-
self._db_api_integration.commenter_data[
530-
"mysql_client_version"
531-
] = cursor._cnx._cmysql.get_client_info()
532-
533-
commenter_data = dict(
534-
self._db_api_integration.commenter_data
535-
)
536-
if self._commenter_options.get(
537-
"opentelemetry_values", True
538-
):
539-
commenter_data.update(
540-
**_get_opentelemetry_values()
541-
)
542-
543-
# Filter down to just the requested attributes.
544-
commenter_data = {
545-
k: v
546-
for k, v in commenter_data.items()
547-
if self._commenter_options.get(k, True)
548-
}
549-
statement = _add_sql_comment(
550-
args_list[0], **commenter_data
551-
)
552-
553-
args_list[0] = statement
554-
args = tuple(args_list)
555-
556-
except Exception as exc: # pylint: disable=broad-except
557-
_logger.exception(
558-
"Exception while generating sql comment: %s",
559-
exc,
560-
)
561-
555+
# sqlcomment in query and span attribute
556+
args = self._update_args_with_added_sql_comment(
557+
args, cursor
558+
)
562559
self._populate_span(span, cursor, *args)
563-
564-
# sqlcomment in query only
565560
else:
561+
# sqlcomment in query only
566562
self._populate_span(span, cursor, *args)
567-
568-
try:
569-
args_list = list(args)
570-
571-
# lazy capture of mysql-connector client version using cursor
572-
if (
573-
self._db_api_integration.database_system
574-
== "mysql"
575-
and self._db_api_integration.connect_module.__name__
576-
== "mysql.connector"
577-
and not self._db_api_integration.commenter_data[
578-
"mysql_client_version"
579-
]
580-
):
581-
self._db_api_integration.commenter_data[
582-
"mysql_client_version"
583-
] = cursor._cnx._cmysql.get_client_info()
584-
585-
commenter_data = dict(
586-
self._db_api_integration.commenter_data
587-
)
588-
if self._commenter_options.get(
589-
"opentelemetry_values", True
590-
):
591-
commenter_data.update(
592-
**_get_opentelemetry_values()
593-
)
594-
595-
# Filter down to just the requested attributes.
596-
commenter_data = {
597-
k: v
598-
for k, v in commenter_data.items()
599-
if self._commenter_options.get(k, True)
600-
}
601-
statement = _add_sql_comment(
602-
args_list[0], **commenter_data
603-
)
604-
605-
args_list[0] = statement
606-
args = tuple(args_list)
607-
608-
except Exception as exc: # pylint: disable=broad-except
609-
_logger.exception(
610-
"Exception while generating sql comment: %s",
611-
exc,
612-
)
613-
614-
# No sqlcommenting
563+
args = self._update_args_with_added_sql_comment(
564+
args, cursor
565+
)
615566
else:
567+
# no sqlcomment
616568
self._populate_span(span, cursor, *args)
617-
618569
return query_method(*args, **kwargs)
619570

620571

‎instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
# pylint: disable=too-many-lines
1516

1617
import logging
1718
import re

0 commit comments

Comments
 (0)