Repository: open-telemetry/opentelemetry-python

Issue:
Add instrumentation for Boto
Add an instrumentation for [Boto](https://boto3.amazonaws.com/v1/documentation/api/latest/index.html)


Discussion:

Resolving patch:
diff --git a/ext/opentelemetry-ext-boto/setup.py b/ext/opentelemetry-ext-boto/setup.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-boto/setup.py
@@ -0,0 +1,26 @@
+# 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 os
+
+import setuptools
+
+BASE_DIR = os.path.dirname(__file__)
+VERSION_FILENAME = os.path.join(
+    BASE_DIR, "src", "opentelemetry", "ext", "boto", "version.py"
+)
+PACKAGE_INFO = {}
+with open(VERSION_FILENAME) as f:
+    exec(f.read(), PACKAGE_INFO)
+
+setuptools.setup(version=PACKAGE_INFO["__version__"])
diff --git a/ext/opentelemetry-ext-boto/src/opentelemetry/ext/boto/__init__.py b/ext/opentelemetry-ext-boto/src/opentelemetry/ext/boto/__init__.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-boto/src/opentelemetry/ext/boto/__init__.py
@@ -0,0 +1,245 @@
+# 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.
+"""
+Instrument `Boto`_ to trace service requests.
+
+There are two options for instrumenting code. The first option is to use the
+``opentelemetry-auto-instrumentation`` executable which will automatically
+instrument your Boto client. The second is to programmatically enable
+instrumentation via the following code:
+
+.. _boto: https://pypi.org/project/boto/
+
+Usage
+-----
+
+.. code:: python
+
+    from opentelemetry import trace
+    from opentelemetry.ext.boto import BotoInstrumentor
+    from opentelemetry.sdk.trace import TracerProvider
+    import boto
+
+    trace.set_tracer_provider(TracerProvider())
+
+    # Instrument Boto
+    BotoInstrumentor().instrument(tracer_provider=trace.get_tracer_provider())
+
+    # This will create a span with Boto-specific attributes
+    ec2 = boto.ec2.connect_to_region("us-west-2")
+    ec2.get_all_instances()
+
+API
+---
+"""
+
+import logging
+from inspect import currentframe
+
+from boto.connection import AWSAuthConnection, AWSQueryConnection
+from wrapt import ObjectProxy, wrap_function_wrapper
+
+from opentelemetry.auto_instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.ext.boto.version import __version__
+from opentelemetry.trace import SpanKind, get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def _get_instance_region_name(instance):
+    region = getattr(instance, "region", None)
+
+    if not region:
+        return None
+    if isinstance(region, str):
+        return region.split(":")[1]
+    return region.name
+
+
+class BotoInstrumentor(BaseInstrumentor):
+    """A instrumentor for Boto
+
+    See `BaseInstrumentor`
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._original_boto = None
+
+    def _instrument(self, **kwargs):
+        # AWSQueryConnection and AWSAuthConnection are two different classes
+        # called by different services for connection.
+        # For exemple EC2 uses AWSQueryConnection and S3 uses
+        # AWSAuthConnection
+
+        # FIXME should the tracer provider be accessed via Configuration
+        # instead?
+        # pylint: disable=attribute-defined-outside-init
+        self._tracer = get_tracer(
+            __name__, __version__, kwargs.get("tracer_provider")
+        )
+
+        wrap_function_wrapper(
+            "boto.connection",
+            "AWSQueryConnection.make_request",
+            self._patched_query_request,
+        )
+        wrap_function_wrapper(
+            "boto.connection",
+            "AWSAuthConnection.make_request",
+            self._patched_auth_request,
+        )
+
+    def _uninstrument(self, **kwargs):
+        unwrap(AWSQueryConnection, "make_request")
+        unwrap(AWSAuthConnection, "make_request")
+
+    def _common_request(  # pylint: disable=too-many-locals
+        self,
+        args_name,
+        traced_args,
+        operation_name,
+        original_func,
+        instance,
+        args,
+        kwargs,
+    ):
+
+        endpoint_name = getattr(instance, "host").split(".")[0]
+
+        with self._tracer.start_as_current_span(
+            "{}.command".format(endpoint_name), kind=SpanKind.CONSUMER,
+        ) as span:
+            if args:
+                http_method = args[0]
+                span.resource = "%s.%s" % (endpoint_name, http_method.lower())
+            else:
+                span.resource = endpoint_name
+
+            add_span_arg_tags(
+                span, endpoint_name, args, args_name, traced_args,
+            )
+
+            # Obtaining region name
+            region_name = _get_instance_region_name(instance)
+
+            meta = {
+                "aws.agent": "boto",
+                "aws.operation": operation_name,
+            }
+            if region_name:
+                meta["aws.region"] = region_name
+
+            for key, value in meta.items():
+                span.set_attribute(key, value)
+
+            # Original func returns a boto.connection.HTTPResponse object
+            result = original_func(*args, **kwargs)
+            span.set_attribute("http.status_code", getattr(result, "status"))
+            span.set_attribute("http.method", getattr(result, "_method"))
+
+            return result
+
+    def _patched_query_request(self, original_func, instance, args, kwargs):
+
+        return self._common_request(
+            ("operation_name", "params", "path", "verb"),
+            ["operation_name", "params", "path"],
+            args[0] if args else None,
+            original_func,
+            instance,
+            args,
+            kwargs,
+        )
+
+    def _patched_auth_request(self, original_func, instance, args, kwargs):
+        operation_name = None
+
+        frame = currentframe().f_back
+        operation_name = None
+        while frame:
+            if frame.f_code.co_name == "make_request":
+                operation_name = frame.f_back.f_code.co_name
+                break
+            frame = frame.f_back
+
+        return self._common_request(
+            (
+                "method",
+                "path",
+                "headers",
+                "data",
+                "host",
+                "auth_path",
+                "sender",
+            ),
+            ["path", "data", "host"],
+            operation_name,
+            original_func,
+            instance,
+            args,
+            kwargs,
+        )
+
+
+def truncate_arg_value(value, max_len=1024):
+    """Truncate values which are bytes and greater than ``max_len``.
+    Useful for parameters like "Body" in ``put_object`` operations.
+    """
+    if isinstance(value, bytes) and len(value) > max_len:
+        return b"..."
+
+    return value
+
+
+def add_span_arg_tags(span, endpoint_name, args, args_names, args_traced):
+    if endpoint_name not in ["kms", "sts"]:
+        tags = dict(
+            (name, value)
+            for (name, value) in zip(args_names, args)
+            if name in args_traced
+        )
+        tags = flatten_dict(tags)
+        for key, value in {
+            k: truncate_arg_value(v)
+            for k, v in tags.items()
+            if k not in {"s3": ["params.Body"]}.get(endpoint_name, [])
+        }.items():
+            span.set_attribute(key, value)
+
+
+def flatten_dict(dict_, sep=".", prefix=""):
+    """
+    Returns a normalized dict of depth 1 with keys in order of embedding
+    """
+    # adapted from https://stackoverflow.com/a/19647596
+    return (
+        {
+            prefix + sep + k if prefix else k: v
+            for kk, vv in dict_.items()
+            for k, v in flatten_dict(vv, sep, kk).items()
+        }
+        if isinstance(dict_, dict)
+        else {prefix: dict_}
+    )
+
+
+def unwrap(obj, attr):
+    function = getattr(obj, attr, None)
+    if (
+        function
+        and isinstance(function, ObjectProxy)
+        and hasattr(function, "__wrapped__")
+    ):
+        setattr(obj, attr, function.__wrapped__)
diff --git a/ext/opentelemetry-ext-boto/src/opentelemetry/ext/boto/version.py b/ext/opentelemetry-ext-boto/src/opentelemetry/ext/boto/version.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-boto/src/opentelemetry/ext/boto/version.py
@@ -0,0 +1,15 @@
+# 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.
+
+__version__ = "0.8.dev0"


Test patch:
diff --git a/ext/opentelemetry-ext-boto/tests/__init__.py b/ext/opentelemetry-ext-boto/tests/__init__.py
new file mode 100644
diff --git a/ext/opentelemetry-ext-boto/tests/conftest.py b/ext/opentelemetry-ext-boto/tests/conftest.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-boto/tests/conftest.py
@@ -0,0 +1,31 @@
+# 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 os import environ
+
+
+def pytest_sessionstart(session):
+    # pylint: disable=unused-argument
+    environ["AWS_ACCESS_KEY_ID"] = "testing"
+    environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    environ["AWS_SECURITY_TOKEN"] = "testing"
+    environ["AWS_SESSION_TOKEN"] = "testing"
+
+
+def pytest_sessionfinish(session):
+    # pylint: disable=unused-argument
+    environ.pop("AWS_ACCESS_KEY_ID")
+    environ.pop("AWS_SECRET_ACCESS_KEY")
+    environ.pop("AWS_SECURITY_TOKEN")
+    environ.pop("AWS_SESSION_TOKEN")
diff --git a/ext/opentelemetry-ext-boto/tests/test_boto_instrumentation.py b/ext/opentelemetry-ext-boto/tests/test_boto_instrumentation.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-boto/tests/test_boto_instrumentation.py
@@ -0,0 +1,242 @@
+# 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 unittest import skipUnless
+
+import boto.awslambda
+import boto.ec2
+import boto.elasticache
+import boto.s3
+import boto.sts
+
+from moto import (  # pylint: disable=import-error
+    mock_ec2_deprecated,
+    mock_lambda_deprecated,
+    mock_s3_deprecated,
+    mock_sts_deprecated,
+)
+from opentelemetry.ext.boto import BotoInstrumentor
+from opentelemetry.test.test_base import TestBase
+
+
+def assert_span_http_status_code(span, code):
+    """Assert on the span's 'http.status_code' tag"""
+    tag = span.attributes["http.status_code"]
+    assert tag == code, "%r != %r" % (tag, code)
+
+
+class TestBotoInstrumentor(TestBase):
+    """Botocore integration testsuite"""
+
+    def setUp(self):
+        super().setUp()
+        BotoInstrumentor().instrument()
+
+    def tearDown(self):
+        BotoInstrumentor().uninstrument()
+
+    @mock_ec2_deprecated
+    def test_ec2_client(self):
+        ec2 = boto.ec2.connect_to_region("us-west-2")
+
+        ec2.get_all_instances()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.attributes["aws.operation"], "DescribeInstances")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.attributes["http.method"], "POST")
+        self.assertEqual(span.attributes["aws.region"], "us-west-2")
+
+        # Create an instance
+        ec2.run_instances(21)
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 2)
+        span = spans[1]
+        self.assertEqual(span.attributes["aws.operation"], "RunInstances")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "ec2.runinstances")
+        self.assertEqual(span.attributes["http.method"], "POST")
+        self.assertEqual(span.attributes["aws.region"], "us-west-2")
+        self.assertEqual(span.name, "ec2.command")
+
+    @mock_ec2_deprecated
+    def test_analytics_enabled_with_rate(self):
+        ec2 = boto.ec2.connect_to_region("us-west-2")
+
+        ec2.get_all_instances()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+
+    @mock_ec2_deprecated
+    def test_analytics_enabled_without_rate(self):
+        ec2 = boto.ec2.connect_to_region("us-west-2")
+
+        ec2.get_all_instances()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+
+    @mock_s3_deprecated
+    def test_s3_client(self):
+        s3 = boto.s3.connect_to_region("us-east-1")
+
+        s3.get_all_buckets()
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.attributes["http.method"], "GET")
+        self.assertEqual(span.attributes["aws.operation"], "get_all_buckets")
+
+        # Create a bucket command
+        s3.create_bucket("cheese")
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 2)
+        span = spans[1]
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.attributes["http.method"], "PUT")
+        self.assertEqual(span.attributes["path"], "/")
+        self.assertEqual(span.attributes["aws.operation"], "create_bucket")
+
+        # Get the created bucket
+        s3.get_bucket("cheese")
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 3)
+        span = spans[2]
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "s3.head")
+        self.assertEqual(span.attributes["http.method"], "HEAD")
+        self.assertEqual(span.attributes["aws.operation"], "head_bucket")
+        self.assertEqual(span.name, "s3.command")
+
+        # Checking for resource incase of error
+        try:
+            s3.get_bucket("big_bucket")
+        except Exception:  # pylint: disable=broad-except
+            spans = self.memory_exporter.get_finished_spans()
+            assert spans
+            span = spans[2]
+            self.assertEqual(span.resource, "s3.head")
+
+    @mock_s3_deprecated
+    def test_s3_put(self):
+        s3 = boto.s3.connect_to_region("us-east-1")
+        s3.create_bucket("mybucket")
+        bucket = s3.get_bucket("mybucket")
+        key = boto.s3.key.Key(bucket)
+        key.key = "foo"
+        key.set_contents_from_string("bar")
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        # create bucket
+        self.assertEqual(len(spans), 3)
+        self.assertEqual(spans[0].attributes["aws.operation"], "create_bucket")
+        assert_span_http_status_code(spans[0], 200)
+        self.assertEqual(spans[0].resource, "s3.put")
+        # get bucket
+        self.assertEqual(spans[1].attributes["aws.operation"], "head_bucket")
+        self.assertEqual(spans[1].resource, "s3.head")
+        # put object
+        self.assertEqual(
+            spans[2].attributes["aws.operation"], "_send_file_internal"
+        )
+        self.assertEqual(spans[2].resource, "s3.put")
+
+    @mock_lambda_deprecated
+    def test_unpatch(self):
+
+        lamb = boto.awslambda.connect_to_region("us-east-2")
+
+        BotoInstrumentor().uninstrument()
+
+        # multiple calls
+        lamb.list_functions()
+        spans = self.memory_exporter.get_finished_spans()
+        assert not spans, spans
+
+    @mock_s3_deprecated
+    def test_double_patch(self):
+        s3 = boto.s3.connect_to_region("us-east-1")
+
+        BotoInstrumentor().instrument()
+        BotoInstrumentor().instrument()
+
+        # Get the created bucket
+        s3.create_bucket("cheese")
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 1)
+
+    @mock_lambda_deprecated
+    def test_lambda_client(self):
+        lamb = boto.awslambda.connect_to_region("us-east-2")
+
+        # multiple calls
+        lamb.list_functions()
+        lamb.list_functions()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 2)
+        span = spans[0]
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "lambda.get")
+        self.assertEqual(span.attributes["http.method"], "GET")
+        self.assertEqual(span.attributes["aws.region"], "us-east-2")
+        self.assertEqual(span.attributes["aws.operation"], "list_functions")
+
+    @mock_sts_deprecated
+    def test_sts_client(self):
+        sts = boto.sts.connect_to_region("us-west-2")
+
+        sts.get_federation_token(12, duration=10)
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(span.resource, "sts.getfederationtoken")
+        self.assertEqual(span.attributes["aws.region"], "us-west-2")
+        self.assertEqual(
+            span.attributes["aws.operation"], "GetFederationToken"
+        )
+
+        # checking for protection on sts against security leak
+        self.assertTrue("args.path" not in span.attributes.keys())
+
+    @skipUnless(
+        False,
+        (
+            "Test to reproduce the case where args sent to patched function "
+            "are None, can't be mocked: needs AWS credentials"
+        ),
+    )
+    def test_elasticache_client(self):
+        elasticache = boto.elasticache.connect_to_region("us-west-2")
+
+        elasticache.describe_cache_clusters()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(span.resource, "elasticache")
+        self.assertEqual(span.attributes["aws.region"], "us-west-2")