Add IDs Generator as Configurable Property of Auto Instrumentation (open-telemetry#1404)

NathanielRN · web-flow · commit da7597c357b8 · 2020-11-24T08:46:51.000-08:00
diff --git a/opentelemetry-api/setup.cfg b/opentelemetry-api/setup.cfg
@@ -57,6 +57,8 @@ opentelemetry_tracer_provider =
 opentelemetry_propagator =
     tracecontext = opentelemetry.trace.propagation.tracecontext:TraceContextTextMapPropagator
     baggage = opentelemetry.baggage.propagation:BaggagePropagator
+opentelemetry_ids_generator =
+    random = opentelemetry.trace.ids_generator:RandomIdsGenerator
 
 [options.extras_require]
 test =
diff --git a/opentelemetry-instrumentation/CHANGELOG.md b/opentelemetry-instrumentation/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add IDs Generator as Configurable Property of Auto Instrumentation
+  ([#1404](https://github.com/open-telemetry/opentelemetry-python/pull/1404))
 - Added support for `OTEL_EXPORTER` to the `opentelemetry-instrument` command ([#1036](https://github.com/open-telemetry/opentelemetry-python/pull/1036))
 
 ## Version 0.14b0
diff --git a/opentelemetry-instrumentation/README.rst b/opentelemetry-instrumentation/README.rst
@@ -73,6 +73,11 @@ Well known trace exporter names:
 
 When present the value is passed on to the relevant exporter initializer as ``service_name`` argument.
 
+* ``--ids-generator`` or ``OTEL_IDS_GENERATOR``
+
+Used to specify which IDs Generator to use for the global Tracer Provider. By default, it
+will use the random IDs generator.
+
 The code in ``program.py`` needs to use one of the packages for which there is
 an OpenTelemetry integration. For a list of the available integrations please
 check `here <https://opentelemetry-python.readthedocs.io/en/stable/index.html#integrations>`_
@@ -93,6 +98,13 @@ The above command will pass ``-e otlp`` to the instrument command and ``--port=3
 The above command will configure global trace provider, attach zipkin and otlp exporters to it and then
 start celery with the rest of the arguments. 
 
+::
+
+    opentelemetry-instrument --ids-generator random flask run --port=3000
+
+The above command will configure the global trace provider to use the Random IDs Generator, and then
+pass ``--port=3000`` to ``flask run``.
+
 References
 ----------
 
diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/__init__.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/__init__.py
@@ -47,6 +47,18 @@ def parse_args():
         """,
     )
 
+    parser.add_argument(
+        "--ids-generator",
+        required=False,
+        help="""
+        The IDs Generator to be used with the Tracer Provider.
+
+        Examples:
+
+            --ids-generator=random
+        """,
+    )
+
     parser.add_argument(
         "-s",
         "--service-name",
@@ -70,6 +82,8 @@ def load_config_from_cli_args(args):
         environ["OTEL_EXPORTER"] = args.exporter
     if args.service_name:
         environ["OTEL_SERVICE_NAME"] = args.service_name
+    if args.ids_generator:
+        environ["OTEL_IDS_GENERATOR"] = args.ids_generator
 
 
 def run() -> None:
diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/components.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/components.py
@@ -34,6 +34,13 @@
 EXPORTER_OTLP_METRIC = "otlp_metric"
 _DEFAULT_EXPORTER = EXPORTER_OTLP
 
+RANDOM_IDS_GENERATOR = "random"
+_DEFAULT_IDS_GENERATOR = RANDOM_IDS_GENERATOR
+
+
+def get_ids_generator() -> str:
+    return Configuration().IDS_GENERATOR or _DEFAULT_IDS_GENERATOR
+
 
 def get_service_name() -> str:
     return Configuration().SERVICE_NAME or ""
@@ -55,10 +62,13 @@ def get_exporter_names() -> Sequence[str]:
     return names
 
 
-def init_tracing(exporters: Sequence[SpanExporter]):
+def init_tracing(
+    exporters: Sequence[SpanExporter], ids_generator: trace.IdsGenerator
+):
     service_name = get_service_name()
     provider = TracerProvider(
         resource=Resource.create({"service.name": service_name}),
+        ids_generator=ids_generator(),
     )
     trace.set_tracer_provider(provider)
 
@@ -80,23 +90,39 @@ def init_metrics(exporters: Sequence[MetricsExporter]):
         logger.warning("automatic metric initialization is not supported yet.")
 
 
-def import_exporters(
-    exporter_names: Sequence[str],
-) -> Tuple[Sequence[SpanExporter], Sequence[MetricsExporter]]:
-    trace_exporters, metric_exporters = {}, {}
-
-    exporters = {
-        ep.name: ep for ep in iter_entry_points("opentelemetry_exporter")
+def import_tracer_provider_config_components(
+    selected_components, entry_point_name
+) -> Sequence[Tuple[str, object]]:
+    component_entry_points = {
+        ep.name: ep for ep in iter_entry_points(entry_point_name)
     }
-
-    for exporter_name in exporter_names:
-        entry_point = exporters.get(exporter_name, None)
+    component_impls = []
+    for selected_component in selected_components:
+        entry_point = component_entry_points.get(selected_component, None)
         if not entry_point:
             raise RuntimeError(
-                "Requested exporter not found: {0}".format(exporter_name)
+                "Requested component '{}' not found in entry points for '{}'".format(
+                    selected_component, entry_point_name
+                )
             )
 
-        exporter_impl = entry_point.load()
+        component_impl = entry_point.load()
+        component_impls.append((selected_component, component_impl))
+
+    return component_impls
+
+
+def import_exporters(
+    exporter_names: Sequence[str],
+) -> Tuple[Sequence[SpanExporter], Sequence[MetricsExporter]]:
+    trace_exporters, metric_exporters = {}, {}
+
+    for (
+        exporter_name,
+        exporter_impl,
+    ) in import_tracer_provider_config_components(
+        exporter_names, "opentelemetry_exporter"
+    ):
         if issubclass(exporter_impl, SpanExporter):
             trace_exporters[exporter_name] = exporter_impl
         elif issubclass(exporter_impl, MetricsExporter):
@@ -110,10 +136,26 @@ def import_exporters(
     return trace_exporters, metric_exporters
 
 
+def import_ids_generator(ids_generator_name: str) -> trace.IdsGenerator:
+    # pylint: disable=unbalanced-tuple-unpacking
+    [
+        (ids_generator_name, ids_generator_impl)
+    ] = import_tracer_provider_config_components(
+        [ids_generator_name.strip()], "opentelemetry_ids_generator"
+    )
+
+    if issubclass(ids_generator_impl, trace.IdsGenerator):
+        return ids_generator_impl
+
+    raise RuntimeError("{0} is not an IdsGenerator".format(ids_generator_name))
+
+
 def initialize_components():
     exporter_names = get_exporter_names()
     trace_exporters, metric_exporters = import_exporters(exporter_names)
-    init_tracing(trace_exporters)
+    ids_generator_name = get_ids_generator()
+    ids_generator = import_ids_generator(ids_generator_name)
+    init_tracing(trace_exporters, ids_generator)
 
     # We don't support automatic initialization for metric yet but have added
     # some boilerplate in order to make sure current implementation does not
diff --git a/opentelemetry-instrumentation/tests/test_auto_tracing.py b/opentelemetry-instrumentation/tests/test_auto_tracing.py
@@ -20,10 +20,12 @@
 from opentelemetry.configuration import Configuration
 from opentelemetry.instrumentation.auto_instrumentation import components
 from opentelemetry.sdk.resources import Resource
+from opentelemetry.trace.ids_generator import RandomIdsGenerator
 
 
 class Provider:
-    def __init__(self, resource=None):
+    def __init__(self, resource=None, ids_generator=None):
+        self.ids_generator = ids_generator
         self.processor = None
         self.resource = resource
 
@@ -48,6 +50,23 @@ class OTLPExporter:
     pass
 
 
+class IdsGenerator:
+    pass
+
+
+class CustomIdsGenerator(IdsGenerator):
+    pass
+
+
+class IterEntryPoint:
+    def __init__(self, name, class_type):
+        self.name = name
+        self.class_type = class_type
+
+    def load(self):
+        return self.class_type
+
+
 class TestTraceInit(TestCase):
     def setUp(self):
         super()
@@ -77,11 +96,12 @@ def tearDown(self):
     def test_trace_init_default(self):
         environ["OTEL_SERVICE_NAME"] = "my-test-service"
         Configuration._reset()
-        components.init_tracing({"zipkin": Exporter})
+        components.init_tracing({"zipkin": Exporter}, RandomIdsGenerator)
 
         self.assertEqual(self.set_provider_mock.call_count, 1)
         provider = self.set_provider_mock.call_args[0][0]
         self.assertIsInstance(provider, Provider)
+        self.assertIsInstance(provider.ids_generator, RandomIdsGenerator)
         self.assertIsInstance(provider.processor, Processor)
         self.assertIsInstance(provider.processor.exporter, Exporter)
         self.assertEqual(
@@ -91,11 +111,12 @@ def test_trace_init_default(self):
     def test_trace_init_otlp(self):
         environ["OTEL_SERVICE_NAME"] = "my-otlp-test-service"
         Configuration._reset()
-        components.init_tracing({"otlp": OTLPExporter})
+        components.init_tracing({"otlp": OTLPExporter}, RandomIdsGenerator)
 
         self.assertEqual(self.set_provider_mock.call_count, 1)
         provider = self.set_provider_mock.call_args[0][0]
         self.assertIsInstance(provider, Provider)
+        self.assertIsInstance(provider.ids_generator, RandomIdsGenerator)
         self.assertIsInstance(provider.processor, Processor)
         self.assertIsInstance(provider.processor.exporter, OTLPExporter)
         self.assertIsInstance(provider.resource, Resource)
@@ -104,3 +125,24 @@ def test_trace_init_otlp(self):
             "my-otlp-test-service",
         )
         del environ["OTEL_SERVICE_NAME"]
+
+    @patch.dict(environ, {"OTEL_IDS_GENERATOR": "custom_ids_generator"})
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation.components.trace.IdsGenerator",
+        new=IdsGenerator,
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation.components.iter_entry_points"
+    )
+    def test_trace_init_custom_ids_generator(self, mock_iter_entry_points):
+        mock_iter_entry_points.configure_mock(
+            return_value=[
+                IterEntryPoint("custom_ids_generator", CustomIdsGenerator)
+            ]
+        )
+        Configuration._reset()
+        ids_generator_name = components.get_ids_generator()
+        ids_generator = components.import_ids_generator(ids_generator_name)
+        components.init_tracing({}, ids_generator)
+        provider = self.set_provider_mock.call_args[0][0]
+        self.assertIsInstance(provider.ids_generator, CustomIdsGenerator)
