Introduce a bootstrap command to auto-install packages

owais · owais · commit ca460091f347 · 2020-05-18T17:43:18.000+05:30
This commit introduces a new boostrap command that is shipped as part of
the opentelemetry-auto-instrumentation package. The command detects
installed libraries and installs the relevant auto-instrumentation
packages.
diff --git a/opentelemetry-auto-instrumentation/setup.cfg b/opentelemetry-auto-instrumentation/setup.cfg
@@ -49,3 +49,4 @@ where = src
 [options.entry_points]
 console_scripts =
     opentelemetry-auto-instrumentation = opentelemetry.auto_instrumentation.auto_instrumentation:run
+    opentelemetry-bootstrap = opentelemetry.auto_instrumentation.bootstrap:run
diff --git a/opentelemetry-auto-instrumentation/src/opentelemetry/auto_instrumentation/bootstrap.py b/opentelemetry-auto-instrumentation/src/opentelemetry/auto_instrumentation/bootstrap.py
@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+
+# 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.
+
+import argparse
+import pkgutil
+import subprocess
+import sys
+from logging import getLogger
+
+logger = getLogger(__file__)
+
+_ACTION_INSTALL = "install"
+_ACTION_REQUIREMENTS = "requirements"
+
+
+# target library to desired instrumentor path/versioned package name
+instrumentations = {
+    "dbapi": "opentelemetry-ext-dbapi>=0.6b0",
+    "flask": "opentelemetry-ext-flask>=0.6b0",
+    "grpc": "opentelemetry-ext-grpc>=0.6b0",
+    "requests": "opentelemetry-ext-requests>=0.6b0",
+    "mysql": "opentelemetry-ext-mysql>=0.6b0",
+    "psycopg2": "opentelemetry-ext-psycopg2>=0.6b0",
+    "pymongo": "opentelemetry-ext-pymongo>=0.6b0",
+    "pymysql": "opentelemetry-ext-pymysql",
+    "redis": "opentelemetry-ext-redis",
+    "sqlalchemy": "opentelemetry-ext-sqlalchemy",
+    "wsgi": "opentelemetry-ext-wsgi>=0.6b0",
+}
+
+# relevant instrumentors and tracers to uninstall and check for conflicts for target libraries
+libraries = {
+    "dbapi": ("opentelemetry-ext-dbapi",),
+    "flask": ("opentelemetry-ext-flask",),
+    "grpc": ("opentelemetry-ext-grpc",),
+    "requests": ("opentelemetry-ext-requests",),
+    "mysql": ("opentelemetry-ext-mysql",),
+    "psycopg2": ("opentelemetry-ext-psycopg2",),
+    "pymongo": ("opentelemetry-ext-pymongo",),
+    "pymysql": ("opentelemetry-ext-pymysql",),
+    "redis": ("opentelemetry-ext-redis",),
+    "sqlalchemy": ("opentelemetry-ext-sqlalchemy",),
+    "wsgi": ("opentelemetry-ext-wsgi",),
+}
+
+
+def _install_package(library, instrumentation):
+    """
+    Ensures that desired version is installed w/o upgrading its dependencies
+    by uninstalling where necessary (if `target` is not provided).
+
+
+    OpenTelemetry auto-instrumentation packages often have traced libraries
+    as instrumentation dependency (e.g. flask for opentelemetry-ext-flask),
+    so using -I on library could cause likely undesired Flask upgrade.
+    Using --no-dependencies alone would leave potential for nonfunctional
+    installations.
+    """
+    pip_list = _pip_freeze()
+    for package in libraries[library]:
+        if "{}==".format(package).lower() in pip_list:
+            logger.info(
+                "Existing %s installation detected.  Uninstalling.", package
+            )
+            _pip_uninstall(package)
+    _pip_install(instrumentation)
+
+
+def _pip_freeze():
+    return (
+        subprocess.check_output([sys.executable, "-m", "pip", "freeze"])
+        .decode()
+        .lower()
+    )
+
+
+def _pip_install(package):
+    # explicit upgrade strategy to override potential pip config
+    subprocess.check_call(
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "-U",
+            "--upgrade-strategy",
+            "only-if-needed",
+            package,
+        ]
+    )
+
+
+def _pip_uninstall(package):
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "uninstall", "-y", package]
+    )
+
+
+def _pip_check():
+    """Ensures none of the instrumentations have dependency conflicts.
+    Clean check reported as:
+    'No broken requirements found.'
+    Dependency conflicts are reported as:
+    'opentelemetry-ext-flask 1.0.1 has requirement opentelemetry-sdk<2.0,>=1.0, but you have opentelemetry-sdk 0.5.'
+    To not be too restrictive, we'll only check for relevant packages.
+    """
+    check_pipe = subprocess.Popen(
+        [sys.executable, "-m", "pip", "check"], stdout=subprocess.PIPE
+    )
+    pip_check = check_pipe.communicate()[0].decode()
+    pip_check_lower = pip_check.lower()
+    for package_tup in libraries.values():
+        for package in package_tup:
+            if package.lower() in pip_check_lower:
+                raise RuntimeError(
+                    "Dependency conflict found: {}".format(pip_check)
+                )
+
+
+def _is_installed(library):
+    return library in sys.modules or pkgutil.find_loader(library) is not None
+
+
+def _find_installed_libraries():
+    return {k: v for k, v in instrumentations.items() if _is_installed(k)}
+
+
+def _run_requirements(packages):
+    print("\n".join(packages.values()), end="")
+
+
+def _run_install(packages):
+    for pkg, inst in packages.items():
+        _install_package(pkg, inst)
+
+    _pip_check()
+
+
+def run() -> None:
+    parser = argparse.ArgumentParser(
+        description="""
+        opentelemetry-bootstrap detects installed libraries and automatically
+        installs the relevant instrumentation packages for them.
+        """
+    )
+    parser.add_argument(
+        "-a",
+        "--action",
+        choices=[_ACTION_INSTALL, _ACTION_REQUIREMENTS],
+        default=_ACTION_REQUIREMENTS,
+        help="""
+        install - uses pip to install the new requirements using to the
+                  currently active site-package.
+        requirements - prints out the new requirements to stdout. Action can
+                       be piped and appended to a requirements.txt file.
+        """,
+    )
+    args = parser.parse_args()
+
+    cmd = {
+        _ACTION_INSTALL: _run_install,
+        _ACTION_REQUIREMENTS: _run_requirements,
+    }[args.action]
+    cmd(_find_installed_libraries())
diff --git a/opentelemetry-auto-instrumentation/tests/test_bootstrap.py b/opentelemetry-auto-instrumentation/tests/test_bootstrap.py
@@ -0,0 +1,125 @@
+# 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 functools import reduce
+from io import StringIO
+from random import sample
+from unittest import TestCase
+from unittest.mock import call, patch
+
+from opentelemetry.auto_instrumentation import bootstrap
+
+
+def sample_packages(packages, rate):
+    sampled = sample(list(packages), int(len(packages) * rate),)
+    return {k: v for k, v in packages.items() if k in sampled}
+
+
+class TestBootstrap(TestCase):
+
+    installed_libraries = {}
+    installed_instrumentations = {}
+
+    @classmethod
+    def setUpClass(cls):
+        # select random 60% of instrumentations
+        cls.installed_libraries = sample_packages(
+            bootstrap.instrumentations, 0.6
+        )
+
+        # treat 50% of sampled packages as pre-installed
+        cls.installed_instrumentations = sample_packages(
+            cls.installed_libraries, 0.5
+        )
+
+        cls.pkg_patcher = patch(
+            "opentelemetry.auto_instrumentation.bootstrap._find_installed_libraries",
+            return_value=cls.installed_libraries,
+        )
+
+        pip_freeze_output = []
+        for inst in cls.installed_instrumentations.values():
+            inst = inst.replace(">=", "==")
+            if "==" not in inst:
+                inst = "{}==x.y".format(inst)
+            pip_freeze_output.append(inst)
+
+        cls.pip_freeze_patcher = patch(
+            "opentelemetry.auto_instrumentation.bootstrap._pip_freeze",
+            return_value="\n".join(pip_freeze_output),
+        )
+        cls.pip_install_patcher = patch(
+            "opentelemetry.auto_instrumentation.bootstrap._pip_install",
+        )
+        cls.pip_uninstall_patcher = patch(
+            "opentelemetry.auto_instrumentation.bootstrap._pip_uninstall",
+        )
+        cls.pip_check_patcher = patch(
+            "opentelemetry.auto_instrumentation.bootstrap._pip_check",
+        )
+
+        cls.pkg_patcher.start()
+        cls.mock_pip_freeze = cls.pip_freeze_patcher.start()
+        cls.mock_pip_install = cls.pip_install_patcher.start()
+        cls.mock_pip_uninstall = cls.pip_uninstall_patcher.start()
+        cls.mock_pip_check = cls.pip_check_patcher.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.pip_check_patcher.start()
+        cls.pip_uninstall_patcher.start()
+        cls.pip_install_patcher.start()
+        cls.pip_freeze_patcher.start()
+        cls.pkg_patcher.stop()
+
+    @patch("sys.argv", ["bootstrap", "-a", "pipenv"])
+    def test_run_unknown_cmd(self):
+        with self.assertRaises(SystemExit):
+            bootstrap.run()
+
+    @patch("sys.argv", ["bootstrap", "-a", "requirements"])
+    def test_run_cmd_print(self):
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            bootstrap.run()
+            self.assertEqual(
+                fake_out.getvalue(),
+                "\n".join(self.installed_libraries.values()),
+            )
+
+    @patch("sys.argv", ["bootstrap", "-a", "install"])
+    def test_run_cmd_install(self):
+        bootstrap.run()
+
+        self.assertEqual(
+            self.mock_pip_freeze.call_count, len(self.installed_libraries)
+        )
+
+        to_uninstall = reduce(
+            lambda x, y: x + y,
+            [
+                pkgs
+                for lib, pkgs in bootstrap.libraries.items()
+                if lib in self.installed_instrumentations
+            ],
+        )
+        self.mock_pip_uninstall.assert_has_calls(
+            [call(i) for i in to_uninstall], any_order=True
+        )
+
+        self.mock_pip_install.assert_has_calls(
+            [call(i) for i in self.installed_libraries.values()],
+            any_order=True,
+        )
+        self.assertEqual(self.mock_pip_check.call_count, 1)
