Repository: open-telemetry/opentelemetry-python-contrib

Issue:
Allow users to specify which distro and configurator to use in auto instrumentation
**Is your feature request related to a problem?**
Yes. Currently, if a user has multiple packages with distro or configurator entry points installed, they have no control over which are used for the auto instrumentation flow.

**Describe the solution you'd like**
Copying what we do for exporters, I am adding 2 env vars where users can specify the entry point names of the distros and configurators they want used. I added warnings for each distro and configurator not used.

**Describe alternatives you've considered**
I considered requiring thes environment variables to be used but instead chose to make them optional. If the env vars are not set, then existing behavior continues: the first distro and configurators found will be used.

**Additional context**
With the way the package is set up, this also adds the --distro and --configurator command line args. Also, we currently have not tests for the sitecustomize file. I found it tricky to implement because importing the module triggers instrumentation. Open to suggestions.


Discussion:

Resolving patch:
diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/_load.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/_load.py
new file mode 100644
--- /dev/null
+++ b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/_load.py
@@ -0,0 +1,124 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from logging import getLogger
+from os import environ
+
+from pkg_resources import iter_entry_points
+
+from opentelemetry.instrumentation.dependencies import (
+    get_dist_dependency_conflicts,
+)
+from opentelemetry.instrumentation.distro import BaseDistro, DefaultDistro
+from opentelemetry.instrumentation.environment_variables import (
+    OTEL_PYTHON_CONFIGURATOR,
+    OTEL_PYTHON_DISABLED_INSTRUMENTATIONS,
+    OTEL_PYTHON_DISTRO,
+)
+from opentelemetry.instrumentation.version import __version__
+
+_logger = getLogger(__name__)
+
+
+def _load_distro() -> BaseDistro:
+    distro_name = environ.get(OTEL_PYTHON_DISTRO, None)
+    for entry_point in iter_entry_points("opentelemetry_distro"):
+        try:
+            # If no distro is specified, use first to come up.
+            if distro_name is None or distro_name == entry_point.name:
+                distro = entry_point.load()()
+                if not isinstance(distro, BaseDistro):
+                    _logger.debug(
+                        "%s is not an OpenTelemetry Distro. Skipping",
+                        entry_point.name,
+                    )
+                    continue
+                _logger.debug(
+                    "Distribution %s will be configured", entry_point.name
+                )
+                return distro
+        except Exception as exc:  # pylint: disable=broad-except
+            _logger.exception(
+                "Distribution %s configuration failed", entry_point.name
+            )
+            raise exc
+    return DefaultDistro()
+
+
+def _load_instrumentors(distro):
+    package_to_exclude = environ.get(OTEL_PYTHON_DISABLED_INSTRUMENTATIONS, [])
+    if isinstance(package_to_exclude, str):
+        package_to_exclude = package_to_exclude.split(",")
+        # to handle users entering "requests , flask" or "requests, flask" with spaces
+        package_to_exclude = [x.strip() for x in package_to_exclude]
+
+    for entry_point in iter_entry_points("opentelemetry_pre_instrument"):
+        entry_point.load()()
+
+    for entry_point in iter_entry_points("opentelemetry_instrumentor"):
+        if entry_point.name in package_to_exclude:
+            _logger.debug(
+                "Instrumentation skipped for library %s", entry_point.name
+            )
+            continue
+
+        try:
+            conflict = get_dist_dependency_conflicts(entry_point.dist)
+            if conflict:
+                _logger.debug(
+                    "Skipping instrumentation %s: %s",
+                    entry_point.name,
+                    conflict,
+                )
+                continue
+
+            # tell instrumentation to not run dep checks again as we already did it above
+            distro.load_instrumentor(entry_point, skip_dep_check=True)
+            _logger.debug("Instrumented %s", entry_point.name)
+        except Exception as exc:  # pylint: disable=broad-except
+            _logger.exception("Instrumenting of %s failed", entry_point.name)
+            raise exc
+
+    for entry_point in iter_entry_points("opentelemetry_post_instrument"):
+        entry_point.load()()
+
+
+def _load_configurators():
+    configurator_name = environ.get(OTEL_PYTHON_CONFIGURATOR, None)
+    configured = None
+    for entry_point in iter_entry_points("opentelemetry_configurator"):
+        if configured is not None:
+            _logger.warning(
+                "Configuration of %s not loaded, %s already loaded",
+                entry_point.name,
+                configured,
+            )
+            continue
+        try:
+            if (
+                configurator_name is None
+                or configurator_name == entry_point.name
+            ):
+                entry_point.load()().configure(auto_instrumentation_version=__version__)  # type: ignore
+                configured = entry_point.name
+            else:
+                _logger.warning(
+                    "Configuration of %s not loaded because %s is set by %s",
+                    entry_point.name,
+                    configurator_name,
+                    OTEL_PYTHON_CONFIGURATOR,
+                )
+        except Exception as exc:  # pylint: disable=broad-except
+            _logger.exception("Configuration of %s failed", entry_point.name)
+            raise exc
diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/sitecustomize.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/sitecustomize.py
--- a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/sitecustomize.py
+++ b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/auto_instrumentation/sitecustomize.py
@@ -16,99 +16,16 @@
 from os import environ
 from os.path import abspath, dirname, pathsep
 
-from pkg_resources import iter_entry_points
-
-from opentelemetry.instrumentation.dependencies import (
-    get_dist_dependency_conflicts,
-)
-from opentelemetry.instrumentation.distro import BaseDistro, DefaultDistro
-from opentelemetry.instrumentation.environment_variables import (
-    OTEL_PYTHON_DISABLED_INSTRUMENTATIONS,
+from opentelemetry.instrumentation.auto_instrumentation._load import (
+    _load_configurators,
+    _load_distro,
+    _load_instrumentors,
 )
 from opentelemetry.instrumentation.utils import _python_path_without_directory
-from opentelemetry.instrumentation.version import __version__
 
 logger = getLogger(__name__)
 
 
-def _load_distros() -> BaseDistro:
-    for entry_point in iter_entry_points("opentelemetry_distro"):
-        try:
-            distro = entry_point.load()()
-            if not isinstance(distro, BaseDistro):
-                logger.debug(
-                    "%s is not an OpenTelemetry Distro. Skipping",
-                    entry_point.name,
-                )
-                continue
-            logger.debug(
-                "Distribution %s will be configured", entry_point.name
-            )
-            return distro
-        except Exception as exc:  # pylint: disable=broad-except
-            logger.exception(
-                "Distribution %s configuration failed", entry_point.name
-            )
-            raise exc
-    return DefaultDistro()
-
-
-def _load_instrumentors(distro):
-    package_to_exclude = environ.get(OTEL_PYTHON_DISABLED_INSTRUMENTATIONS, [])
-    if isinstance(package_to_exclude, str):
-        package_to_exclude = package_to_exclude.split(",")
-        # to handle users entering "requests , flask" or "requests, flask" with spaces
-        package_to_exclude = [x.strip() for x in package_to_exclude]
-
-    for entry_point in iter_entry_points("opentelemetry_pre_instrument"):
-        entry_point.load()()
-
-    for entry_point in iter_entry_points("opentelemetry_instrumentor"):
-        if entry_point.name in package_to_exclude:
-            logger.debug(
-                "Instrumentation skipped for library %s", entry_point.name
-            )
-            continue
-
-        try:
-            conflict = get_dist_dependency_conflicts(entry_point.dist)
-            if conflict:
-                logger.debug(
-                    "Skipping instrumentation %s: %s",
-                    entry_point.name,
-                    conflict,
-                )
-                continue
-
-            # tell instrumentation to not run dep checks again as we already did it above
-            distro.load_instrumentor(entry_point, skip_dep_check=True)
-            logger.debug("Instrumented %s", entry_point.name)
-        except Exception as exc:  # pylint: disable=broad-except
-            logger.exception("Instrumenting of %s failed", entry_point.name)
-            raise exc
-
-    for entry_point in iter_entry_points("opentelemetry_post_instrument"):
-        entry_point.load()()
-
-
-def _load_configurators():
-    configured = None
-    for entry_point in iter_entry_points("opentelemetry_configurator"):
-        if configured is not None:
-            logger.warning(
-                "Configuration of %s not loaded, %s already loaded",
-                entry_point.name,
-                configured,
-            )
-            continue
-        try:
-            entry_point.load()().configure(auto_instrumentation_version=__version__)  # type: ignore
-            configured = entry_point.name
-        except Exception as exc:  # pylint: disable=broad-except
-            logger.exception("Configuration of %s failed", entry_point.name)
-            raise exc
-
-
 def initialize():
     # prevents auto-instrumentation of subprocesses if code execs another python process
     environ["PYTHONPATH"] = _python_path_without_directory(
@@ -116,7 +33,7 @@ def initialize():
     )
 
     try:
-        distro = _load_distros()
+        distro = _load_distro()
         distro.configure()
         _load_configurators()
         _load_instrumentors(distro)
diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/environment_variables.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/environment_variables.py
--- a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/environment_variables.py
+++ b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/environment_variables.py
@@ -16,3 +16,13 @@
 """
 .. envvar:: OTEL_PYTHON_DISABLED_INSTRUMENTATIONS
 """
+
+OTEL_PYTHON_DISTRO = "OTEL_PYTHON_DISTRO"
+"""
+.. envvar:: OTEL_PYTHON_DISTRO
+"""
+
+OTEL_PYTHON_CONFIGURATOR = "OTEL_PYTHON_CONFIGURATOR"
+"""
+.. envvar:: OTEL_PYTHON_CONFIGURATOR
+"""


Test patch:
diff --git a/opentelemetry-instrumentation/tests/auto_instrumentation/test_load.py b/opentelemetry-instrumentation/tests/auto_instrumentation/test_load.py
new file mode 100644
--- /dev/null
+++ b/opentelemetry-instrumentation/tests/auto_instrumentation/test_load.py
@@ -0,0 +1,312 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# type: ignore
+
+from unittest import TestCase
+from unittest.mock import Mock, call, patch
+
+from opentelemetry.instrumentation.auto_instrumentation import _load
+from opentelemetry.instrumentation.environment_variables import (
+    OTEL_PYTHON_CONFIGURATOR,
+    OTEL_PYTHON_DISABLED_INSTRUMENTATIONS,
+    OTEL_PYTHON_DISTRO,
+)
+from opentelemetry.instrumentation.version import __version__
+
+
+class TestLoad(TestCase):
+    @patch.dict(
+        "os.environ", {OTEL_PYTHON_CONFIGURATOR: "custom_configurator2"}
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_configurators(self, iter_mock):
+        # Add multiple entry points but only specify the 2nd in the environment variable.
+        ep_mock1 = Mock()
+        ep_mock1.name = "custom_configurator1"
+        configurator_mock1 = Mock()
+        ep_mock1.load.return_value = configurator_mock1
+        ep_mock2 = Mock()
+        ep_mock2.name = "custom_configurator2"
+        configurator_mock2 = Mock()
+        ep_mock2.load.return_value = configurator_mock2
+        ep_mock3 = Mock()
+        ep_mock3.name = "custom_configurator3"
+        configurator_mock3 = Mock()
+        ep_mock3.load.return_value = configurator_mock3
+
+        iter_mock.return_value = (ep_mock1, ep_mock2, ep_mock3)
+        _load._load_configurators()
+        configurator_mock1.assert_not_called()
+        configurator_mock2().configure.assert_called_once_with(
+            auto_instrumentation_version=__version__
+        )
+        configurator_mock3.assert_not_called()
+
+    @patch.dict(
+        "os.environ", {OTEL_PYTHON_CONFIGURATOR: "custom_configurator2"}
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_configurators_no_ep(
+        self,
+        iter_mock,
+    ):
+        iter_mock.return_value = ()
+        # Confirm method does not crash if not entry points exist.
+        _load._load_configurators()
+
+    @patch.dict(
+        "os.environ", {OTEL_PYTHON_CONFIGURATOR: "custom_configurator2"}
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_configurators_error(self, iter_mock):
+        # Add multiple entry points but only specify the 2nd in the environment variable.
+        ep_mock1 = Mock()
+        ep_mock1.name = "custom_configurator1"
+        configurator_mock1 = Mock()
+        ep_mock1.load.return_value = configurator_mock1
+        ep_mock2 = Mock()
+        ep_mock2.name = "custom_configurator2"
+        configurator_mock2 = Mock()
+        configurator_mock2().configure.side_effect = Exception()
+        ep_mock2.load.return_value = configurator_mock2
+        ep_mock3 = Mock()
+        ep_mock3.name = "custom_configurator3"
+        configurator_mock3 = Mock()
+        ep_mock3.load.return_value = configurator_mock3
+
+        iter_mock.return_value = (ep_mock1, ep_mock2, ep_mock3)
+        # Confirm failed configuration raises exception.
+        self.assertRaises(Exception, _load._load_configurators)
+
+    @patch.dict("os.environ", {OTEL_PYTHON_DISTRO: "custom_distro2"})
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.isinstance"
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_distro(self, iter_mock, isinstance_mock):
+        # Add multiple entry points but only specify the 2nd in the environment variable.
+        ep_mock1 = Mock()
+        ep_mock1.name = "custom_distro1"
+        distro_mock1 = Mock()
+        ep_mock1.load.return_value = distro_mock1
+        ep_mock2 = Mock()
+        ep_mock2.name = "custom_distro2"
+        distro_mock2 = Mock()
+        ep_mock2.load.return_value = distro_mock2
+        ep_mock3 = Mock()
+        ep_mock3.name = "custom_distro3"
+        distro_mock3 = Mock()
+        ep_mock3.load.return_value = distro_mock3
+
+        iter_mock.return_value = (ep_mock1, ep_mock2, ep_mock3)
+        # Mock entry points to be instances of BaseDistro.
+        isinstance_mock.return_value = True
+        self.assertEqual(
+            _load._load_distro(),
+            distro_mock2(),
+        )
+
+    @patch.dict("os.environ", {OTEL_PYTHON_DISTRO: "custom_distro2"})
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.isinstance"
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.DefaultDistro"
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_distro_not_distro(
+        self, iter_mock, default_distro_mock, isinstance_mock
+    ):
+        # Add multiple entry points but only specify the 2nd in the environment variable.
+        ep_mock1 = Mock()
+        ep_mock1.name = "custom_distro1"
+        distro_mock1 = Mock()
+        ep_mock1.load.return_value = distro_mock1
+        ep_mock2 = Mock()
+        ep_mock2.name = "custom_distro2"
+        distro_mock2 = Mock()
+        ep_mock2.load.return_value = distro_mock2
+        ep_mock3 = Mock()
+        ep_mock3.name = "custom_distro3"
+        distro_mock3 = Mock()
+        ep_mock3.load.return_value = distro_mock3
+
+        iter_mock.return_value = (ep_mock1, ep_mock2, ep_mock3)
+        # Confirm default distro is used if specified entry point is not a BaseDistro
+        isinstance_mock.return_value = False
+        self.assertEqual(
+            _load._load_distro(),
+            default_distro_mock(),
+        )
+
+    @patch.dict("os.environ", {OTEL_PYTHON_DISTRO: "custom_distro2"})
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.DefaultDistro"
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_distro_no_ep(self, iter_mock, default_distro_mock):
+        iter_mock.return_value = ()
+        # Confirm default distro is used if there are no entry points.
+        self.assertEqual(
+            _load._load_distro(),
+            default_distro_mock(),
+        )
+
+    @patch.dict("os.environ", {OTEL_PYTHON_DISTRO: "custom_distro2"})
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.isinstance"
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_distro_error(self, iter_mock, isinstance_mock):
+        ep_mock1 = Mock()
+        ep_mock1.name = "custom_distro1"
+        distro_mock1 = Mock()
+        ep_mock1.load.return_value = distro_mock1
+        ep_mock2 = Mock()
+        ep_mock2.name = "custom_distro2"
+        distro_mock2 = Mock()
+        distro_mock2.side_effect = Exception()
+        ep_mock2.load.return_value = distro_mock2
+        ep_mock3 = Mock()
+        ep_mock3.name = "custom_distro3"
+        distro_mock3 = Mock()
+        ep_mock3.load.return_value = distro_mock3
+
+        iter_mock.return_value = (ep_mock1, ep_mock2, ep_mock3)
+        isinstance_mock.return_value = True
+        # Confirm method raises exception if it fails to load a distro.
+        self.assertRaises(Exception, _load._load_distro)
+
+    @patch.dict(
+        "os.environ",
+        {OTEL_PYTHON_DISABLED_INSTRUMENTATIONS: " instr1 , instr3 "},
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.get_dist_dependency_conflicts"
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_instrumentors(self, iter_mock, dep_mock):
+        # Mock opentelemetry_pre_instrument entry points
+        pre_ep_mock1 = Mock()
+        pre_ep_mock1.name = "pre1"
+        pre_mock1 = Mock()
+        pre_ep_mock1.load.return_value = pre_mock1
+
+        pre_ep_mock2 = Mock()
+        pre_ep_mock2.name = "pre2"
+        pre_mock2 = Mock()
+        pre_ep_mock2.load.return_value = pre_mock2
+
+        # Mock opentelemetry_instrumentor entry points
+        ep_mock1 = Mock()
+        ep_mock1.name = "instr1"
+
+        ep_mock2 = Mock()
+        ep_mock2.name = "instr2"
+
+        ep_mock3 = Mock()
+        ep_mock3.name = "instr3"
+
+        ep_mock4 = Mock()
+        ep_mock4.name = "instr4"
+
+        # Mock opentelemetry_instrumentor entry points
+        post_ep_mock1 = Mock()
+        post_ep_mock1.name = "post1"
+        post_mock1 = Mock()
+        post_ep_mock1.load.return_value = post_mock1
+
+        post_ep_mock2 = Mock()
+        post_ep_mock2.name = "post2"
+        post_mock2 = Mock()
+        post_ep_mock2.load.return_value = post_mock2
+
+        distro_mock = Mock()
+
+        # Mock entry points in order
+        iter_mock.side_effect = [
+            (pre_ep_mock1, pre_ep_mock2),
+            (ep_mock1, ep_mock2, ep_mock3, ep_mock4),
+            (post_ep_mock1, post_ep_mock2),
+        ]
+        # No dependency conflict
+        dep_mock.return_value = None
+        _load._load_instrumentors(distro_mock)
+        # All opentelemetry_pre_instrument entry points should be loaded
+        pre_mock1.assert_called_once()
+        pre_mock2.assert_called_once()
+        self.assertEqual(iter_mock.call_count, 3)
+        # Only non-disabled instrumentations should be loaded
+        distro_mock.load_instrumentor.assert_has_calls(
+            [
+                call(ep_mock2, skip_dep_check=True),
+                call(ep_mock4, skip_dep_check=True),
+            ]
+        )
+        self.assertEqual(distro_mock.load_instrumentor.call_count, 2)
+        # All opentelemetry_post_instrument entry points should be loaded
+        post_mock1.assert_called_once()
+        post_mock2.assert_called_once()
+
+    @patch.dict(
+        "os.environ",
+        {OTEL_PYTHON_DISABLED_INSTRUMENTATIONS: " instr1 , instr3 "},
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.get_dist_dependency_conflicts"
+    )
+    @patch(
+        "opentelemetry.instrumentation.auto_instrumentation._load.iter_entry_points"
+    )
+    def test_load_instrumentors_dep_conflict(self, iter_mock, dep_mock):
+        ep_mock1 = Mock()
+        ep_mock1.name = "instr1"
+
+        ep_mock2 = Mock()
+        ep_mock2.name = "instr2"
+
+        ep_mock3 = Mock()
+        ep_mock3.name = "instr3"
+
+        ep_mock4 = Mock()
+        ep_mock4.name = "instr4"
+
+        distro_mock = Mock()
+
+        iter_mock.return_value = (ep_mock1, ep_mock2, ep_mock3, ep_mock4)
+        # If a dependency conflict is raised, that instrumentation should not be loaded, but others still should.
+        dep_mock.side_effect = [None, "DependencyConflict"]
+        _load._load_instrumentors(distro_mock)
+        distro_mock.load_instrumentor.assert_has_calls(
+            [
+                call(ep_mock2, skip_dep_check=True),
+            ]
+        )
+        distro_mock.load_instrumentor.assert_called_once()
diff --git a/opentelemetry-instrumentation/tests/test_run.py b/opentelemetry-instrumentation/tests/auto_instrumentation/test_run.py
similarity index 100%
rename from opentelemetry-instrumentation/tests/test_run.py
rename to opentelemetry-instrumentation/tests/auto_instrumentation/test_run.py