Skip to content

Commit

Permalink
add attribute for db port
Browse files Browse the repository at this point in the history
  • Loading branch information
glee2429 committed Nov 16, 2024
1 parent 803bb32 commit 7096019
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 71 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -63,14 +63,14 @@


def trace_integration(
connect_module: typing.Callable[..., typing.Any],
connect_method_name: str,
database_system: str,
connection_attributes: typing.Dict = None,
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
db_api_integration_factory=None,
connect_module: typing.Callable[..., typing.Any],
connect_method_name: str,
database_system: str,
connection_attributes: typing.Dict = None,
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
db_api_integration_factory=None,
):
"""Integrate with DB API library.
https://www.python.org/dev/peps/pep-0249/
Expand Down Expand Up @@ -104,17 +104,17 @@ def trace_integration(


def wrap_connect(
name: str,
connect_module: typing.Callable[..., typing.Any],
connect_method_name: str,
database_system: str,
connection_attributes: typing.Dict = None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
db_api_integration_factory=None,
commenter_options: dict = None,
name: str,
connect_module: typing.Callable[..., typing.Any],
connect_method_name: str,
database_system: str,
connection_attributes: typing.Dict = None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
db_api_integration_factory=None,
commenter_options: dict = None,
):
"""Integrate with DB API library.
https://www.python.org/dev/peps/pep-0249/
Expand All @@ -136,15 +136,15 @@ def wrap_connect(
"""
db_api_integration_factory = (
db_api_integration_factory or DatabaseApiIntegration
db_api_integration_factory or DatabaseApiIntegration
)

# pylint: disable=unused-argument
def wrap_connect_(
wrapped: typing.Callable[..., typing.Any],
instance: typing.Any,
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
wrapped: typing.Callable[..., typing.Any],
instance: typing.Any,
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
):
db_integration = db_api_integration_factory(
name,
Expand All @@ -168,7 +168,7 @@ def wrap_connect_(


def unwrap_connect(
connect_module: typing.Callable[..., typing.Any], connect_method_name: str
connect_module: typing.Callable[..., typing.Any], connect_method_name: str
):
"""Disable integration with DB API library.
https://www.python.org/dev/peps/pep-0249/
Expand All @@ -181,15 +181,15 @@ def unwrap_connect(


def instrument_connection(
name: str,
connection,
database_system: str,
connection_attributes: typing.Dict = None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
commenter_options: dict = None,
name: str,
connection,
database_system: str,
connection_attributes: typing.Dict = None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
commenter_options: dict = None,
):
"""Enable instrumentation in a database connection.
Expand Down Expand Up @@ -244,16 +244,16 @@ def uninstrument_connection(connection):

class DatabaseApiIntegration:
def __init__(
self,
name: str,
database_system: str,
connection_attributes=None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
commenter_options: dict = None,
connect_module: typing.Callable[..., typing.Any] = None,
self,
name: str,
database_system: str,
connection_attributes=None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
capture_parameters: bool = False,
enable_commenter: bool = False,
commenter_options: dict = None,
connect_module: typing.Callable[..., typing.Any] = None,
):
self.connection_attributes = connection_attributes
if self.connection_attributes is None:
Expand Down Expand Up @@ -283,8 +283,8 @@ def __init__(
self.commenter_data = self.calculate_commenter_data()

def _get_db_version(
self,
db_driver,
self,
db_driver,
):
if db_driver in _DB_DRIVER_ALIASES:
return util_version(_DB_DRIVER_ALIASES[db_driver])
Expand All @@ -296,7 +296,7 @@ def _get_db_version(
return db_version

def calculate_commenter_data(
self,
self,
):
commenter_data = {}
if not self.enable_commenter:
Expand Down Expand Up @@ -346,10 +346,10 @@ def calculate_commenter_data(
return commenter_data

def wrapped_connection(
self,
connect_method: typing.Callable[..., typing.Any],
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
self,
connect_method: typing.Callable[..., typing.Any],
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
):
"""Add object proxy to connection object."""
connection = connect_method(*args, **kwargs)
Expand Down Expand Up @@ -391,7 +391,7 @@ def get_connection_attributes(self, connection):


def get_traced_connection_proxy(
connection, db_api_integration, *args, **kwargs
connection, db_api_integration, *args, **kwargs
):
# pylint: disable=abstract-method
class TracedConnectionProxy(wrapt.ObjectProxy):
Expand Down Expand Up @@ -435,10 +435,10 @@ def __init__(self, db_api_integration: DatabaseApiIntegration) -> None:
self._leading_comment_remover = re.compile(r"^/\*.*?\*/")

def _populate_span(
self,
span: trace_api.Span,
cursor,
*args: typing.Tuple[typing.Any, typing.Any],
self,
span: trace_api.Span,
cursor,
*args: typing.Tuple[typing.Any, typing.Any],
):
if not span.is_recording():
return
Expand All @@ -451,9 +451,11 @@ def _populate_span(
)
span.set_attribute(SpanAttributes.DB_STATEMENT, statement)

span.set_attribute(SpanAttributes.DB_PORT, self._db_api_integration.connection_props.get("port"))

for (
attribute_key,
attribute_value,
attribute_key,
attribute_value,
) in self._db_api_integration.span_attributes.items():
span.set_attribute(attribute_key, attribute_value)

Expand All @@ -475,11 +477,11 @@ def get_statement(self, cursor, args): # pylint: disable=no-self-use
return statement

def traced_execution(
self,
cursor,
query_method: typing.Callable[..., typing.Any],
*args: typing.Tuple[typing.Any, typing.Any],
**kwargs: typing.Dict[typing.Any, typing.Any],
self,
cursor,
query_method: typing.Callable[..., typing.Any],
*args: typing.Tuple[typing.Any, typing.Any],
**kwargs: typing.Dict[typing.Any, typing.Any],
):
name = self.get_operation_name(cursor, args)
if not name:
Expand All @@ -490,7 +492,7 @@ def traced_execution(
)

with self._db_api_integration._tracer.start_as_current_span(
name, kind=SpanKind.CLIENT
name, kind=SpanKind.CLIENT
) as span:
self._populate_span(span, cursor, *args)
if args and self._commenter_enabled:
Expand All @@ -499,12 +501,12 @@ def traced_execution(

# lazy capture of mysql-connector client version using cursor
if (
self._db_api_integration.database_system == "mysql"
and self._db_api_integration.connect_module.__name__
== "mysql.connector"
and not self._db_api_integration.commenter_data[
"mysql_client_version"
]
self._db_api_integration.database_system == "mysql"
and self._db_api_integration.connect_module.__name__
== "mysql.connector"
and not self._db_api_integration.commenter_data[
"mysql_client_version"
]
):
self._db_api_integration.commenter_data[
"mysql_client_version"
Expand All @@ -514,7 +516,7 @@ def traced_execution(
self._db_api_integration.commenter_data
)
if self._commenter_options.get(
"opentelemetry_values", True
"opentelemetry_values", True
):
commenter_data.update(**_get_opentelemetry_values())

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,9 @@ def test_span_succeeded(self):
self.assertEqual(
span.attributes[SpanAttributes.DB_STATEMENT], "Test query"
)
self.assertEqual(
span.attributes[SpanAttributes.DB_PORT], 123
)
self.assertFalse("db.statement.parameters" in span.attributes)
self.assertEqual(span.attributes[SpanAttributes.DB_USER], "testuser")
self.assertEqual(
Expand Down

0 comments on commit 7096019

Please sign in to comment.