Repository: open-telemetry/opentelemetry-python

Issue:
Add instrumentation for botocore
Add a new instrumentation to support [botocore](https://pypi.org/project/botocore/).


Discussion:

Resolving patch:
diff --git a/ext/opentelemetry-ext-botocore/setup.py b/ext/opentelemetry-ext-botocore/setup.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-botocore/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", "botocore", "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-botocore/src/opentelemetry/ext/botocore/__init__.py b/ext/opentelemetry-ext-botocore/src/opentelemetry/ext/botocore/__init__.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-botocore/src/opentelemetry/ext/botocore/__init__.py
@@ -0,0 +1,209 @@
+# 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 `Botocore`_ to trace service requests.
+
+There are two options for instrumenting code. The first option is to use the
+``opentelemetry-instrument`` executable which will automatically
+instrument your Botocore client. The second is to programmatically enable
+instrumentation via the following code:
+
+.. _Botocore: https://pypi.org/project/botocore/
+
+Usage
+-----
+
+.. code:: python
+
+    from opentelemetry import trace
+    from opentelemetry.ext.botocore import BotocoreInstrumentor
+    from opentelemetry.sdk.trace import TracerProvider
+    import botocore
+
+    trace.set_tracer_provider(TracerProvider())
+
+    # Instrument Botocore
+    BotocoreInstrumentor().instrument(
+        tracer_provider=trace.get_tracer_provider()
+    )
+
+    # This will create a span with Botocore-specific attributes
+    session = botocore.session.get_session()
+    session.set_credentials(
+        access_key="access-key", secret_key="secret-key"
+    )
+    ec2 = self.session.create_client("ec2", region_name="us-west-2")
+    ec2.describe_instances()
+
+API
+---
+"""
+
+import logging
+
+from botocore.client import BaseClient
+from wrapt import ObjectProxy, wrap_function_wrapper
+
+from opentelemetry.ext.botocore.version import __version__
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import SpanKind, get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+class BotocoreInstrumentor(BaseInstrumentor):
+    """A instrumentor for Botocore
+
+    See `BaseInstrumentor`
+    """
+
+    def _instrument(self, **kwargs):
+
+        # 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(
+            "botocore.client",
+            "BaseClient._make_api_call",
+            self._patched_api_call,
+        )
+
+    def _uninstrument(self, **kwargs):
+        unwrap(BaseClient, "_make_api_call")
+
+    def _patched_api_call(self, original_func, instance, args, kwargs):
+
+        endpoint_name = deep_getattr(instance, "_endpoint._endpoint_prefix")
+
+        with self._tracer.start_as_current_span(
+            "{}.command".format(endpoint_name), kind=SpanKind.CONSUMER,
+        ) as span:
+
+            operation = None
+            if args:
+                operation = args[0]
+                span.resource = "%s.%s" % (endpoint_name, operation.lower())
+
+            else:
+                span.resource = endpoint_name
+
+            add_span_arg_tags(
+                span,
+                endpoint_name,
+                args,
+                ("action", "params", "path", "verb"),
+                {"params", "path", "verb"},
+            )
+
+            region_name = deep_getattr(instance, "meta.region_name")
+
+            meta = {
+                "aws.agent": "botocore",
+                "aws.operation": operation,
+                "aws.region": region_name,
+            }
+            for key, value in meta.items():
+                span.set_attribute(key, value)
+
+            result = original_func(*args, **kwargs)
+
+            span.set_attribute(
+                "http.status_code",
+                result["ResponseMetadata"]["HTTPStatusCode"],
+            )
+            span.set_attribute(
+                "retry_attempts", result["ResponseMetadata"]["RetryAttempts"],
+            )
+
+            return result
+
+
+def unwrap(obj, attr):
+    function = getattr(obj, attr, None)
+    if (
+        function
+        and isinstance(function, ObjectProxy)
+        and hasattr(function, "__wrapped__")
+    ):
+        setattr(obj, attr, function.__wrapped__)
+
+
+def add_span_arg_tags(span, endpoint_name, args, args_names, args_traced):
+    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
+
+    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 deep_getattr(obj, attr_string, default=None):
+    """
+    Returns the attribute of ``obj`` at the dotted path given by
+    ``attr_string``, if no such attribute is reachable, returns ``default``.
+
+    >>> deep_getattr(cass, "cluster")
+    <cassandra.cluster.Cluster object at 0xa20c350
+
+    >>> deep_getattr(cass, "cluster.metadata.partitioner")
+    u"org.apache.cassandra.dht.Murmur3Partitioner"
+
+    >>> deep_getattr(cass, "i.dont.exist", default="default")
+    "default"
+    """
+    attrs = attr_string.split(".")
+    for attr in attrs:
+        try:
+            obj = getattr(obj, attr)
+        except AttributeError:
+            return default
+
+    return obj
diff --git a/ext/opentelemetry-ext-botocore/src/opentelemetry/ext/botocore/version.py b/ext/opentelemetry-ext-botocore/src/opentelemetry/ext/botocore/version.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-botocore/src/opentelemetry/ext/botocore/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.9.dev0"


Test patch:
diff --git a/ext/opentelemetry-ext-botocore/tests/__init__.py b/ext/opentelemetry-ext-botocore/tests/__init__.py
new file mode 100644
diff --git a/ext/opentelemetry-ext-botocore/tests/test_botocore_instrumentation.py b/ext/opentelemetry-ext-botocore/tests/test_botocore_instrumentation.py
new file mode 100644
--- /dev/null
+++ b/ext/opentelemetry-ext-botocore/tests/test_botocore_instrumentation.py
@@ -0,0 +1,197 @@
+import botocore.session
+from botocore.exceptions import ParamValidationError
+
+from moto import (  # pylint: disable=import-error
+    mock_ec2,
+    mock_kinesis,
+    mock_kms,
+    mock_lambda,
+    mock_s3,
+    mock_sqs,
+)
+from opentelemetry.ext.botocore import BotocoreInstrumentor
+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 TestBotocoreInstrumentor(TestBase):
+    """Botocore integration testsuite"""
+
+    def setUp(self):
+        super().setUp()
+        BotocoreInstrumentor().instrument()
+
+        self.session = botocore.session.get_session()
+        self.session.set_credentials(
+            access_key="access-key", secret_key="secret-key"
+        )
+
+    def tearDown(self):
+        super().tearDown()
+        BotocoreInstrumentor().uninstrument()
+
+    @mock_ec2
+    def test_traced_client(self):
+        ec2 = self.session.create_client("ec2", region_name="us-west-2")
+
+        ec2.describe_instances()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(len(spans), 1)
+        self.assertEqual(span.attributes["aws.agent"], "botocore")
+        self.assertEqual(span.attributes["aws.region"], "us-west-2")
+        self.assertEqual(span.attributes["aws.operation"], "DescribeInstances")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "ec2.describeinstances")
+        self.assertEqual(span.name, "ec2.command")
+
+    @mock_ec2
+    def test_traced_client_analytics(self):
+        ec2 = self.session.create_client("ec2", region_name="us-west-2")
+        ec2.describe_instances()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+
+    @mock_s3
+    def test_s3_client(self):
+        s3 = self.session.create_client("s3", region_name="us-west-2")
+
+        s3.list_buckets()
+        s3.list_buckets()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(len(spans), 2)
+        self.assertEqual(span.attributes["aws.operation"], "ListBuckets")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "s3.listbuckets")
+
+        # testing for span error
+        self.memory_exporter.get_finished_spans()
+        with self.assertRaises(ParamValidationError):
+            s3.list_objects(bucket="mybucket")
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[2]
+        self.assertEqual(span.resource, "s3.listobjects")
+
+    @mock_s3
+    def test_s3_put(self):
+        params = dict(Key="foo", Bucket="mybucket", Body=b"bar")
+        s3 = self.session.create_client("s3", region_name="us-west-2")
+        s3.create_bucket(Bucket="mybucket")
+        s3.put_object(**params)
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(len(spans), 2)
+        self.assertEqual(span.attributes["aws.operation"], "CreateBucket")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "s3.createbucket")
+        self.assertEqual(spans[1].attributes["aws.operation"], "PutObject")
+        self.assertEqual(spans[1].resource, "s3.putobject")
+        self.assertEqual(spans[1].attributes["params.Key"], str(params["Key"]))
+        self.assertEqual(
+            spans[1].attributes["params.Bucket"], str(params["Bucket"])
+        )
+        self.assertTrue("params.Body" not in spans[1].attributes.keys())
+
+    @mock_sqs
+    def test_sqs_client(self):
+        sqs = self.session.create_client("sqs", region_name="us-east-1")
+
+        sqs.list_queues()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(len(spans), 1)
+        self.assertEqual(span.attributes["aws.region"], "us-east-1")
+        self.assertEqual(span.attributes["aws.operation"], "ListQueues")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "sqs.listqueues")
+
+    @mock_kinesis
+    def test_kinesis_client(self):
+        kinesis = self.session.create_client(
+            "kinesis", region_name="us-east-1"
+        )
+
+        kinesis.list_streams()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(len(spans), 1)
+        self.assertEqual(span.attributes["aws.region"], "us-east-1")
+        self.assertEqual(span.attributes["aws.operation"], "ListStreams")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "kinesis.liststreams")
+
+    @mock_kinesis
+    def test_unpatch(self):
+        kinesis = self.session.create_client(
+            "kinesis", region_name="us-east-1"
+        )
+
+        BotocoreInstrumentor().uninstrument()
+
+        kinesis.list_streams()
+        spans = self.memory_exporter.get_finished_spans()
+        assert not spans, spans
+
+    @mock_sqs
+    def test_double_patch(self):
+        sqs = self.session.create_client("sqs", region_name="us-east-1")
+
+        BotocoreInstrumentor().instrument()
+        BotocoreInstrumentor().instrument()
+
+        sqs.list_queues()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 1)
+
+    @mock_lambda
+    def test_lambda_client(self):
+        lamb = self.session.create_client("lambda", region_name="us-east-1")
+
+        lamb.list_functions()
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(len(spans), 1)
+        self.assertEqual(span.attributes["aws.region"], "us-east-1")
+        self.assertEqual(span.attributes["aws.operation"], "ListFunctions")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "lambda.listfunctions")
+
+    @mock_kms
+    def test_kms_client(self):
+        kms = self.session.create_client("kms", region_name="us-east-1")
+
+        kms.list_keys(Limit=21)
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        span = spans[0]
+        self.assertEqual(len(spans), 1)
+        self.assertEqual(span.attributes["aws.region"], "us-east-1")
+        self.assertEqual(span.attributes["aws.operation"], "ListKeys")
+        assert_span_http_status_code(span, 200)
+        self.assertEqual(span.resource, "kms.listkeys")
+
+        # checking for protection on sts against security leak
+        self.assertTrue("params" not in span.attributes.keys())