feat(dsm): implement protobuf schema tracking

.gitlab/tests/contrib.yml

@@ -179,6 +179,11 @@ molten:
   variables:
     SUITE_NAME: "molten"
 
+protobuf:
+  extends: .test_base_riot_snapshot
+  variables:
+    SUITE_NAME: "protobuf"
+
 opentracer:
   extends: .test_base_riot
   variables:

.riot/requirements/1721018.txt

@@ -0,0 +1,22 @@
+#
+# This file is autogenerated by pip-compile with Python 3.10
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate .riot/requirements/1721018.in
+#
+attrs==24.2.0
+coverage[toml]==7.6.1
+exceptiongroup==1.2.2
+hypothesis==6.45.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pluggy==1.5.0
+protobuf==5.28.0
+pytest==8.3.3
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1

.riot/requirements/1bf9721.txt

@@ -0,0 +1,20 @@
+#
+# This file is autogenerated by pip-compile with Python 3.12
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate .riot/requirements/1bf9721.in
+#
+attrs==24.2.0
+coverage[toml]==7.6.1
+hypothesis==6.45.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pluggy==1.5.0
+protobuf==5.28.0
+pytest==8.3.3
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0

.riot/requirements/26aada0.txt

@@ -0,0 +1,20 @@
+#
+# This file is autogenerated by pip-compile with Python 3.11
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate .riot/requirements/26aada0.in
+#
+attrs==24.2.0
+coverage[toml]==7.6.1
+hypothesis==6.45.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pluggy==1.5.0
+protobuf==5.28.0
+pytest==8.3.3
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0

.riot/requirements/3b28562.txt

@@ -0,0 +1,24 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate --resolver=backtracking .riot/requirements/3b28562.in
+#
+attrs==24.2.0
+coverage[toml]==7.6.1
+exceptiongroup==1.2.2
+hypothesis==6.45.0
+importlib-metadata==8.4.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pluggy==1.5.0
+protobuf==5.28.0
+pytest==8.3.3
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1
+zipp==3.20.1

.riot/requirements/e222783.txt

@@ -0,0 +1,24 @@
+#
+# This file is autogenerated by pip-compile with Python 3.8
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate --resolver=backtracking .riot/requirements/e222783.in
+#
+attrs==24.2.0
+coverage[toml]==7.6.1
+exceptiongroup==1.2.2
+hypothesis==6.45.0
+importlib-metadata==8.4.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pluggy==1.5.0
+protobuf==5.28.0
+pytest==8.3.3
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1
+zipp==3.20.1

ddtrace/contrib/internal/protobuf/patch.py

@@ -0,0 +1,120 @@
+from google import protobuf
+from google.protobuf.internal import builder
+import wrapt
+
+from ddtrace import config
+from ddtrace.internal.utils.wrappers import unwrap
+from ddtrace.pin import Pin
+
+from .schema_iterator import SchemaExtractor
+
+
+config._add(
+    "protobuf",
+    dict(),
+)
+
+
+_WRAPPED_MESSAGE_CLASSES = []
+
+
+def get_version():
+    # type: () -> str
+    return getattr(protobuf, "__version__", "")
+
+
+def patch():
+    """Patch the instrumented methods"""
+    if getattr(protobuf, "_datadog_patch", False):
+        return
+    protobuf._datadog_patch = True
+
+    _w = wrapt.wrap_function_wrapper
+
+    _w("google.protobuf.internal", "builder.BuildTopDescriptorsAndMessages", _traced_build)
+    Pin().onto(builder)
+
+
+def unpatch():
+    if getattr(protobuf, "_datadog_patch", False):
+        protobuf._datadog_patch = False
+
+        unwrap(protobuf.internal.builder, "BuildTopDescriptorsAndMessages")
+
+        global _WRAPPED_MESSAGE_CLASSES
+        for wrapped_message_class in _WRAPPED_MESSAGE_CLASSES:
+            _unwrap_message(wrapped_message_class)
+
+        _WRAPPED_MESSAGE_CLASSES = []
+
+
+def _unwrap_message(message_class):
+    unwrap(message_class, "SerializeToString")
+    unwrap(message_class, "ParseFromString")
+
+
+def _wrap_message(message_descriptor, message_class):
+    def serialize_wrapper(wrapped, instance, args, kwargs):
+        return _traced_serialize_message(wrapped, instance, args, kwargs, msg_descriptor=message_descriptor)
+
+    def deserialize_wrapper(wrapped, instance, args, kwargs):
+        return _traced_deserialize_message(wrapped, instance, args, kwargs, msg_descriptor=message_descriptor)
+
+    _w = wrapt.wrap_function_wrapper
+    _w(message_class, "SerializeToString", serialize_wrapper)
+    _w(message_class, "ParseFromString", deserialize_wrapper)
+
+    global _WRAPPED_MESSAGE_CLASSES
+    _WRAPPED_MESSAGE_CLASSES.append(message_class)
+    Pin().onto(message_class)
+
+
+#
+# tracing functions
+#
+def _traced_build(func, instance, args, kwargs):
+    file_des = args[0]
+
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        func(*args, **kwargs)
+
+    try:
+        func(*args, **kwargs)
+    finally:
+        if config._data_streams_enabled:
+            generated_message_classes = args[2]
+            message_descriptors = file_des.message_types_by_name.items()
+            for message_idx in range(len(message_descriptors)):
+                message_class_name = message_descriptors[message_idx][0]
+                message_descriptor = message_descriptors[message_idx][1]
+                message_class = generated_message_classes[message_class_name]
+                _wrap_message(message_descriptor=message_descriptor, message_class=message_class)
+
+
+def _traced_deserialize_message(func, instance, args, kwargs, msg_descriptor):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        func(*args, **kwargs)
+
+    active = pin.tracer.current_span()
+
+    try:
+        func(*args, **kwargs)
+    finally:
+        if config._data_streams_enabled and active:
+            SchemaExtractor.attach_schema_on_span(msg_descriptor, active, SchemaExtractor.DESERIALIZATION)
+
+
+def _traced_serialize_message(func, instance, args, kwargs, msg_descriptor):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled() or not msg_descriptor:
+        return func(*args, **kwargs)
+
+    active = pin.tracer.current_span()
+
+    try:
+        return func(*args, **kwargs)
+    finally:
+        if config._data_streams_enabled and active:
+            SchemaExtractor.attach_schema_on_span(msg_descriptor, active, SchemaExtractor.SERIALIZATION)

ddtrace/contrib/internal/protobuf/schema_iterator.py

@@ -0,0 +1,114 @@
+from google.protobuf.descriptor import FieldDescriptor
+
+# from google._upb._message import Descriptor
+from ddtrace._trace.span import Span
+from ddtrace.ext import schema as SCHEMA_TAGS
+from ddtrace.internal.datastreams import data_streams_processor
+from ddtrace.internal.datastreams.schemas.schema_builder import SchemaBuilder
+from ddtrace.internal.datastreams.schemas.schema_iterator import SchemaIterator
+
+
+class SchemaExtractor(SchemaIterator):
+    SERIALIZATION = "serialization"
+    DESERIALIZATION = "deserialization"
+    PROTOBUF = "protobuf"
+
+    def __init__(self, schema):
+        self.schema = schema
+
+    @staticmethod
+    def extract_property(
+        field: FieldDescriptor, schema_name: str, field_name: str, builder: SchemaBuilder, depth: int
+    ) -> int:
+        array = False
+        type_ = None
+        format_ = None
+        description = None
+        ref = None
+        enum_values = None
+
+        type_format = SchemaExtractor.get_type_and_format(field.type)
+        type_ = type_format[0]
+        format_ = type_format[1]
+
+        if type_ is None and format_ == "Message type":
+            format_ = None
+            ref = "#/components/schemas/" + field.message_type.name
+            if not SchemaExtractor.extract_schema(field.message_type, builder, depth):
+                return False
+        elif format_ == "enum":
+            enum_values = [value.name for value in field.enum_type.values]
+        return builder.add_property(schema_name, field_name, array, type_, description, ref, format_, enum_values)
+
+    @staticmethod
+    def extract_schema(schema, builder: SchemaBuilder, depth: int) -> bool:
+        depth += 1
+        schema_name = schema.name
+        if not builder.should_extract_schema(schema_name, depth):
+            return False
+        try:
+            for field_name, field_desc in schema.fields_by_name.items():
+                if not SchemaExtractor.extract_property(field_desc, schema_name, field_name, builder, depth):
+                    return False
+        except Exception:
+            return False
+        return True
+
+    @staticmethod
+    def extract_schemas(descriptor) -> bool:
+        return data_streams_processor().get_schema(descriptor.name, SchemaExtractor(descriptor))
+
+    def iterate_over_schema(self, builder: SchemaBuilder):
+        self.extract_schema(self.schema, builder, 0)
+
+    @staticmethod
+    def attach_schema_on_span(descriptor, span: Span, operation: str):
+        if descriptor is None or span is None:
+            return
+
+        span.set_tag(SCHEMA_TAGS.SCHEMA_TYPE, SchemaExtractor.PROTOBUF)
+        span.set_tag(SCHEMA_TAGS.SCHEMA_NAME, descriptor.name)
+        span.set_tag(SCHEMA_TAGS.SCHEMA_OPERATION, operation)
+
+        if not data_streams_processor().can_sample_schema(operation):
+            return
+
+        prio = span.context.sampling_priority
+        if prio is None or prio <= 0:
+            return
+
+        weight = data_streams_processor().try_sample_schema(operation)
+        if weight == 0:
+            return
+
+        schema_data = SchemaExtractor.extract_schemas(descriptor)
+
+        span.set_tag(SCHEMA_TAGS.SCHEMA_DEFINITION, schema_data.definition)
+        span.set_metric(SCHEMA_TAGS.SCHEMA_WEIGHT, weight)
+        span.set_tag(SCHEMA_TAGS.SCHEMA_ID, schema_data.id)
+
+    @staticmethod
+    def get_type_and_format(type_: int) -> tuple:
+        type_format_mapping = {
+            1: ("number", "double"),
+            2: ("number", "float"),
+            3: ("integer", "int64"),
+            4: ("integer", "uint64"),
+            5: ("integer", "int32"),
+            6: ("integer", "fixed64"),
+            7: ("integer", "fixed32"),
+            8: ("boolean", None),
+            9: ("string", None),
+            10: ("object", "Group type"),  # Group types are deprecated
+            11: (None, "Message type"),  # Placeholder for handling messages
+            12: ("string", "byte"),
+            13: ("integer", "uint32"),
+            14: ("string", "enum"),
+            15: ("integer", "sfixed32"),
+            16: ("integer", "sfixed64"),
+            17: ("integer", "int32"),
+            18: ("integer", "int64"),
+        }
+
+        # Default values for unknown types
+        return type_format_mapping.get(type_, ("string", None))

ddtrace/contrib/protobuf/__init__.py

@@ -0,0 +1,30 @@
+"""
+The Protobuf integration will trace all Protobuf read / write calls made with the ``google.protobuf``
+library. This integration is enabled by default.
+
+Enabling
+~~~~~~~~
+
+The protobuf integration is enabled by default. Use
+:func:`patch()<ddtrace.patch>` to enable the integration::
+
+    from ddtrace import patch
+    patch(protobuf=True)
+
+Configuration
+~~~~~~~~~~~~~
+
+"""
+from ...internal.utils.importlib import require_modules
+
+
+required_modules = ["protobuf"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        # Expose public methods
+        from ..internal.protobuf.patch import get_version
+        from ..internal.protobuf.patch import patch
+        from ..internal.protobuf.patch import unpatch
+
+        __all__ = ["patch", "unpatch", "get_version"]