diff --git a/ddtrace/propagation/_database_monitoring.py b/ddtrace/propagation/_database_monitoring.py index 4210e4cbec6..4002f864dd8 100644 --- a/ddtrace/propagation/_database_monitoring.py +++ b/ddtrace/propagation/_database_monitoring.py @@ -23,6 +23,7 @@ DBM_DATABASE_SERVICE_NAME_KEY = "dddbs" DBM_PEER_HOSTNAME_KEY = "ddh" DBM_PEER_DB_NAME_KEY = "dddb" +DBM_PEER_SERVICE_KEY = "ddprs" DBM_ENVIRONMENT_KEY = "dde" DBM_VERSION_KEY = "ddpv" DBM_TRACE_PARENT_KEY = "traceparent" @@ -56,12 +57,14 @@ def __init__( sql_injector=default_sql_injector, peer_hostname_tag="out.host", peer_db_name_tag="db.name", + peer_service_tag="peer.service", ): self.sql_pos = sql_pos self.sql_kw = sql_kw self.sql_injector = sql_injector self.peer_hostname_tag = peer_hostname_tag self.peer_db_name_tag = peer_db_name_tag + self.peer_service_tag = peer_service_tag def inject(self, dbspan, args, kwargs): # run sampling before injection to propagate correct sampling priority @@ -114,6 +117,10 @@ def _get_dbm_comment(self, db_span): if peer_hostname: dbm_tags[DBM_PEER_HOSTNAME_KEY] = peer_hostname + peer_service = db_span.get_tag(self.peer_service_tag) + if peer_service: + dbm_tags[DBM_PEER_SERVICE_KEY] = peer_service + if dbm_config.propagation_mode == "full": db_span.set_tag_str(DBM_TRACE_INJECTED_TAG, "true") dbm_tags[DBM_TRACE_PARENT_KEY] = db_span.context._traceparent diff --git a/tests/contrib/aiomysql/test_aiomysql.py b/tests/contrib/aiomysql/test_aiomysql.py index 35e0a7e09c6..2247b2dba6f 100644 --- a/tests/contrib/aiomysql/test_aiomysql.py +++ b/tests/contrib/aiomysql/test_aiomysql.py @@ -230,7 +230,9 @@ class AioMySQLTestCase(AsyncioTestCase): TEST_SERVICE = "mysql" conn = None - async def _get_conn_tracer(self): + async def _get_conn_tracer(self, tags=None): + tags = tags if tags is not None else {} + if not self.conn: self.conn = await aiomysql.connect(**AIOMYSQL_CONFIG) assert not self.conn.closed @@ -239,7 +241,7 @@ async def _get_conn_tracer(self): assert pin # Customize the service # we have to apply it on the existing one since new one won't inherit `app` - pin.clone(tracer=self.tracer).onto(self.conn) + pin.clone(tracer=self.tracer, tags={**tags, **pin.tags}).onto(self.conn) return self.conn, self.tracer @@ -429,3 +431,26 @@ async def test_aiomysql_dbm_propagation_comment_peer_service_enabled(self): await shared_tests._test_dbm_propagation_comment_peer_service_enabled( config=AIOMYSQL_CONFIG, cursor=cursor, wrapped_instance=cursor.__wrapped__ ) + + @mark_asyncio + @AsyncioTestCase.run_in_subprocess( + env_overrides=dict( + DD_DBM_PROPAGATION_MODE="service", + DD_SERVICE="orders-app", + DD_ENV="staging", + DD_VERSION="v7343437-d7ac743", + DD_TRACE_SPAN_ATTRIBUTE_SCHEMA="v1", + ) + ) + async def test_aiomysql_dbm_propagation_comment_with_peer_service_tag(self): + """tests if dbm comment is set in mysql""" + conn, tracer = await self._get_conn_tracer({"peer.service": "peer_service_name"}) + cursor = await conn.cursor() + cursor.__wrapped__ = mock.AsyncMock() + + await shared_tests._test_dbm_propagation_comment_with_peer_service_tag( + config=AIOMYSQL_CONFIG, + cursor=cursor, + wrapped_instance=cursor.__wrapped__, + peer_service_name="peer_service_name", + ) diff --git a/tests/contrib/shared_tests.py b/tests/contrib/shared_tests.py index 2ccb319551f..97d1df32cfa 100644 --- a/tests/contrib/shared_tests.py +++ b/tests/contrib/shared_tests.py @@ -94,3 +94,18 @@ async def _test_dbm_propagation_comment_peer_service_enabled(config, cursor, wra await _test_execute(dbm_comment, cursor, wrapped_instance) if execute_many: await _test_execute_many(dbm_comment, cursor, wrapped_instance) + + +async def _test_dbm_propagation_comment_with_peer_service_tag( + config, cursor, wrapped_instance, peer_service_name, execute_many=True +): + """tests if dbm comment is set in mysql""" + db_name = config["db"] + + dbm_comment = ( + f"/*dddb='{db_name}',dddbs='test',dde='staging',ddh='127.0.0.1',ddprs='{peer_service_name}',ddps='orders-app'," + "ddpv='v7343437-d7ac743'*/ " + ) + await _test_execute(dbm_comment, cursor, wrapped_instance) + if execute_many: + await _test_execute_many(dbm_comment, cursor, wrapped_instance)