From 2c60b8930cb0bcd20383834a077b077bb30b210f Mon Sep 17 00:00:00 2001
From: Nathaniel Ruiz Nowell <enowell@amazon.com>
Date: Sat, 21 Nov 2020 02:04:46 -0800
Subject: [PATCH] Generalize exports and IDs generator code

---
 .../auto_instrumentation/components.py        | 56 +++++++++++--------
 1 file changed, 34 insertions(+), 22 deletions(-)

diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/components.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/components.py
index 4fdb3c8ee13..8b4ef4f3e8b 100644
--- a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/components.py
+++ b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/components.py
@@ -90,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):
@@ -121,16 +137,12 @@ def import_exporters(
 
 
 def import_ids_generator(ids_generator_name: str) -> trace.IdsGenerator:
-    ids_generator_impl = None
-
-    for id_generator in iter_entry_points("opentelemetry_ids_generator"):
-        if id_generator.name == ids_generator_name.strip():
-            ids_generator_impl = id_generator.load()
-            break
-    else:
-        raise RuntimeError(
-            "Requested IDs Generator not found: {0}".format(ids_generator_name)
-        )
+    # 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