ft: add instrumentation to s3

Signed-off-by: Cagri Yonca <[email protected]>

src/instana/instrumentation/aws/boto3.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Tuple, Type
 
 from instana.instrumentation.aws.dynamodb import collect_dynamodb_attributes
+from instana.instrumentation.aws.s3 import collect_s3_attributes
 from opentelemetry.semconv.trace import SpanAttributes
 
 if TYPE_CHECKING:
@@ -76,6 +77,10 @@ def make_api_call_with_instana(
                     collect_dynamodb_attributes(
                         wrapped, instance, args, kwargs, parent_context
                     )
+                elif instance.meta.service_model.service_name == "s3":
+                    collect_s3_attributes(
+                        wrapped, instance, args, kwargs, parent_context
+                    )
                 else:
                     with tracer.start_as_current_span(
                         "boto3", span_context=parent_context

src/instana/instrumentation/aws/s3.py

@@ -0,0 +1,86 @@
+# (c) Copyright IBM Corp. 2021
+# (c) Copyright Instana Inc. 2020
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Type
+
+from instana.span_context import SpanContext
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+try:
+    import wrapt
+
+    from instana.log import logger
+    from instana.singletons import tracer
+    from instana.util.traceutils import (
+        get_tracer_tuple,
+        tracing_is_off,
+    )
+
+    operations = {
+        "upload_file": "UploadFile",
+        "upload_fileobj": "UploadFileObj",
+        "download_file": "DownloadFile",
+        "download_fileobj": "DownloadFileObj",
+    }
+
+    def collect_s3_attributes(
+        wrapped: Callable[..., Dict[str, Any]],
+        instance: Type["BaseClient"],
+        args: Sequence[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        parent_context: SpanContext,
+    ) -> None:
+        with tracer.start_as_current_span("s3", span_context=parent_context) as span:
+            try:
+                span.set_attribute("s3.op", args[0])
+                if "Bucket" in args[1].keys():
+                    span.set_attribute("s3.bucket", args[1]["Bucket"])
+            except Exception as exc:
+                span.record_exception(exc)
+                logger.debug(
+                    "collect_dynamodb_attributes: collect error", exc_info=True
+                )
+
+    def collect_s3_injected_attributes(
+        wrapped: Callable[..., object],
+        instance: Type["BaseClient"],
+        args: Sequence[object],
+        kwargs: Dict[str, Any],
+    ) -> Callable[..., object]:
+        # If we're not tracing, just return
+        if tracing_is_off():
+            return wrapped(*args, **kwargs)
+
+        tracer, parent_span, _ = get_tracer_tuple()
+
+        parent_context = parent_span.get_span_context() if parent_span else None
+
+        with tracer.start_as_current_span("s3", span_context=parent_context) as span:
+            try:
+                span.set_attribute("s3.op", operations[wrapped.__name__])
+                if wrapped.__name__ in ["download_file", "download_fileobj"]:
+                    span.set_attribute("s3.bucket", args[0])
+                else:
+                    span.set_attribute("s3.bucket", args[1])
+                return wrapped(*args, **kwargs)
+            except Exception as exc:
+                span.record_exception(exc)
+                logger.debug(
+                    "s3_inject_method_with_instana: collect error", exc_info=True
+                )
+
+    for method in [
+        "upload_file",
+        "upload_fileobj",
+        "download_file",
+        "download_fileobj",
+    ]:
+        wrapt.wrap_function_wrapper(
+            "boto3.s3.inject", method, collect_s3_injected_attributes
+        )
+
+    logger.debug("Instrumenting s3")
+except ImportError:
+    pass

src/instana/span/kind.py

@@ -50,6 +50,7 @@
     "redis",
     "rpc-client",
     "sqlalchemy",
+    "s3",
     "tornado-client",
     "urllib3",
     "pymongo",

src/instana/span/registered_span.py

@@ -260,6 +260,10 @@ def _populate_exit_span_data(self, span: "InstanaSpan") -> None:
             # self.data["rpc"]["baggage"] = span.attributes.pop("rpc.baggage", None)
             self.data["rpc"]["error"] = span.attributes.pop("rpc.error", None)
 
+        elif span.name == "s3":
+            self.data["s3"]["op"] = span.attributes.pop("s3.op", None)
+            self.data["s3"]["bucket"] = span.attributes.pop("s3.bucket", None)
+
         elif span.name == "sqlalchemy":
             self.data["sqlalchemy"]["sql"] = span.attributes.pop("sqlalchemy.sql", None)
             self.data["sqlalchemy"]["eng"] = span.attributes.pop("sqlalchemy.eng", None)