From b3c0781a250c959f6cf69650ba51c94973bc2052 Mon Sep 17 00:00:00 2001 From: Cagri Yonca Date: Thu, 27 Feb 2025 11:33:36 +0100 Subject: [PATCH 1/3] enhancement: refactoring project structure and pep8 format Signed-off-by: Cagri Yonca --- docker-compose.yml | 1 - src/instana/__init__.py | 22 ++- src/instana/instrumentation/aws/boto3.py | 126 +++++++++++++ .../instrumentation/aws/lambda_inst.py | 146 ++++++++------- src/instana/instrumentation/boto3_inst.py | 174 ------------------ .../{cassandra_inst.py => cassandra.py} | 0 .../{couchbase_inst.py => couchbase.py} | 0 .../{fastapi_inst.py => fastapi.py} | 4 +- .../{gevent_inst.py => gevent.py} | 24 ++- .../{sanic_inst.py => sanic.py} | 1 - .../{starlette_inst.py => starlette.py} | 0 11 files changed, 231 insertions(+), 267 deletions(-) create mode 100644 src/instana/instrumentation/aws/boto3.py delete mode 100644 src/instana/instrumentation/boto3_inst.py rename src/instana/instrumentation/{cassandra_inst.py => cassandra.py} (100%) rename src/instana/instrumentation/{couchbase_inst.py => couchbase.py} (100%) rename src/instana/instrumentation/{fastapi_inst.py => fastapi.py} (98%) rename src/instana/instrumentation/{gevent_inst.py => gevent.py} (65%) rename src/instana/instrumentation/{sanic_inst.py => sanic.py} (99%) rename src/instana/instrumentation/{starlette_inst.py => starlette.py} (100%) diff --git a/docker-compose.yml b/docker-compose.yml index a60d89df..09a4b4f1 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -13,7 +13,6 @@ services: ports: - 9042:9042 - couchbase: image: public.ecr.aws/docker/library/couchbase:community ports: diff --git a/src/instana/__init__.py b/src/instana/__init__.py index 046ddccc..8c111576 100644 --- a/src/instana/__init__.py +++ b/src/instana/__init__.py @@ -168,11 +168,10 @@ def boot_agent() -> None: from instana.instrumentation import ( aioamqp, # noqa: F401 asyncio, # noqa: F401 - boto3_inst, # noqa: F401 - cassandra_inst, # noqa: F401 + cassandra, # noqa: F401 celery, # noqa: F401 - couchbase_inst, # noqa: F401 - fastapi_inst, # noqa: F401 + couchbase, # noqa: F401 + fastapi, # noqa: F401 flask, # noqa: F401 # gevent_inst, # noqa: F401 grpcio, # noqa: F401 @@ -185,9 +184,9 @@ def boot_agent() -> None: pymysql, # noqa: F401 pyramid, # noqa: F401 redis, # noqa: F401 - sanic_inst, # noqa: F401 + sanic, # noqa: F401 sqlalchemy, # noqa: F401 - starlette_inst, # noqa: F401 + starlette, # noqa: F401 urllib3, # noqa: F401 ) from instana.instrumentation.aiohttp import ( @@ -196,7 +195,10 @@ def boot_agent() -> None: from instana.instrumentation.aiohttp import ( server as aiohttp_server, # noqa: F401 ) - from instana.instrumentation.aws import lambda_inst # noqa: F401 + from instana.instrumentation.aws import ( + boto3, # noqa: F401 + lambda_inst, # noqa: F401 + ) from instana.instrumentation.django import middleware # noqa: F401 from instana.instrumentation.google.cloud import ( pubsub, # noqa: F401 @@ -209,12 +211,14 @@ def boot_agent() -> None: client as tornado_client, # noqa: F401 ) from instana.instrumentation.tornado import ( - client as tornado_client, # noqa: F401 server as tornado_server, # noqa: F401 ) # Hooks - from instana.hooks import hook_gunicorn, hook_uwsgi # noqa: F401 + from instana.hooks import ( + hook_gunicorn, # noqa: F401 + hook_uwsgi, # noqa: F401 + ) if "INSTANA_DISABLE" not in os.environ: diff --git a/src/instana/instrumentation/aws/boto3.py b/src/instana/instrumentation/aws/boto3.py new file mode 100644 index 00000000..7f7c6006 --- /dev/null +++ b/src/instana/instrumentation/aws/boto3.py @@ -0,0 +1,126 @@ +# (c) Copyright IBM Corp. 2025 + +from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Tuple, Type + +from opentelemetry.semconv.trace import SpanAttributes + +if TYPE_CHECKING: + from botocore.auth import SigV4Auth + from botocore.client import BaseClient + + from instana.span.span import InstanaSpan + +try: + import json + + import wrapt + + from instana.log import logger + from instana.propagators.format import Format + from instana.singletons import tracer + from instana.span.span import get_current_span + from instana.util.traceutils import ( + extract_custom_headers, + get_tracer_tuple, + tracing_is_off, + ) + + def lambda_inject_context(payload: Dict[str, Any], span: "InstanaSpan") -> None: + """ + When boto3 lambda client 'Invoke' is called, we want to inject the tracing context. + boto3/botocore has specific requirements: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/lambda.html#Lambda.Client.invoke + """ + try: + invoke_payload = payload.get("Payload", {}) + + if not isinstance(invoke_payload, dict): + invoke_payload = json.loads(invoke_payload) + + tracer.inject(span.context, Format.HTTP_HEADERS, invoke_payload) + payload["Payload"] = json.dumps(invoke_payload) + except Exception: + logger.debug("non-fatal lambda_inject_context: ", exc_info=True) + + @wrapt.patch_function_wrapper("botocore.auth", "SigV4Auth.add_auth") + def emit_add_auth_with_instana( + wrapped: Callable[..., None], + instance: "SigV4Auth", + args: Tuple[object], + kwargs: Dict[str, Any], + ) -> Callable[..., None]: + current_span = get_current_span() + if not tracing_is_off() and current_span and current_span.is_recording(): + extract_custom_headers(current_span, args[0].headers) + return wrapped(*args, **kwargs) + + @wrapt.patch_function_wrapper("botocore.client", "BaseClient._make_api_call") + def make_api_call_with_instana( + wrapped: Callable[..., Dict[str, Any]], + instance: Type["BaseClient"], + args: Sequence[Dict[str, Any]], + kwargs: Dict[str, Any], + ) -> Dict[str, Any]: + # 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 + + try: + with tracer.start_as_current_span( + "boto3", span_context=parent_context + ) as span: + try: + operation = args[0] + payload = args[1] + + span.set_attribute("op", operation) + span.set_attribute("ep", instance._endpoint.host) + span.set_attribute("reg", instance._client_config.region_name) + + span.set_attribute( + SpanAttributes.HTTP_URL, + instance._endpoint.host + ":443/" + args[0], + ) + span.set_attribute(SpanAttributes.HTTP_METHOD, "POST") + + # Don't collect payload for SecretsManager + if not hasattr(instance, "get_secret_value"): + span.set_attribute("payload", payload) + + # Inject context when invoking lambdas + if "lambda" in instance._endpoint.host and operation == "Invoke": + lambda_inject_context(payload, span) + + except Exception: + logger.debug( + "make_api_call_with_instana: collect error", + exc_info=True, + ) + + try: + result = wrapped(*args, **kwargs) + + if isinstance(result, dict): + http_dict = result.get("ResponseMetadata") + if isinstance(http_dict, dict): + status = http_dict.get("HTTPStatusCode") + if status is not None: + span.set_attribute("http.status_code", status) + headers = http_dict.get("HTTPHeaders") + extract_custom_headers(span, headers) + + return result + except Exception as exc: + span.mark_as_errored({"error": exc}) + raise + except Exception: + logger.debug("make_api_call_with_instana: collect error", exc_info=True) + else: + return wrapped(*args, **kwargs) + +except ImportError: + pass diff --git a/src/instana/instrumentation/aws/lambda_inst.py b/src/instana/instrumentation/aws/lambda_inst.py index 1dc5c959..62cabcc4 100644 --- a/src/instana/instrumentation/aws/lambda_inst.py +++ b/src/instana/instrumentation/aws/lambda_inst.py @@ -5,89 +5,93 @@ Instrumentation for AWS Lambda functions """ -import sys -import traceback from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple -import wrapt -from opentelemetry.semconv.trace import SpanAttributes - -from instana import get_aws_lambda_handler -from instana.instrumentation.aws.triggers import enrich_lambda_span, get_context -from instana.log import logger -from instana.singletons import env_is_aws_lambda, get_agent, get_tracer -from instana.util.ids import define_server_timing - if TYPE_CHECKING: from instana.agent.aws_lambda import AWSLambdaAgent +try: + import sys + import traceback -def lambda_handler_with_instana( - wrapped: Callable[..., object], - instance: object, - args: Tuple[object, ...], - kwargs: Dict[str, Any], -) -> object: - event = args[0] - agent: "AWSLambdaAgent" = get_agent() - tracer = get_tracer() + import wrapt + from opentelemetry.semconv.trace import SpanAttributes - agent.collector.collect_snapshot(*args) - incoming_ctx = get_context(tracer, event) + from instana import get_aws_lambda_handler + from instana.instrumentation.aws.triggers import enrich_lambda_span, get_context + from instana.log import logger + from instana.singletons import env_is_aws_lambda, get_agent, get_tracer + from instana.util.ids import define_server_timing - result = None - with tracer.start_as_current_span( - "aws.lambda.entry", span_context=incoming_ctx - ) as span: - enrich_lambda_span(agent, span, *args) - try: - result = wrapped(*args, **kwargs) + def lambda_handler_with_instana( + wrapped: Callable[..., object], + instance: object, + args: Tuple[object, ...], + kwargs: Dict[str, Any], + ) -> object: + event = args[0] + agent: "AWSLambdaAgent" = get_agent() + tracer = get_tracer() - if isinstance(result, dict): - server_timing_value = define_server_timing(span.context.trace_id) - if "headers" in result: - result["headers"]["Server-Timing"] = server_timing_value - elif "multiValueHeaders" in result: - result["multiValueHeaders"]["Server-Timing"] = [server_timing_value] - if "statusCode" in result and result.get("statusCode"): - status_code = int(result["statusCode"]) - span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, status_code) - if 500 <= status_code: - span.record_exception(f"HTTP status {status_code}") - except Exception as exc: - logger.debug(f"AWS Lambda lambda_handler_with_instana error: {exc}") - if span: - exc = traceback.format_exc() - span.record_exception(exc) - raise - finally: - agent.collector.shutdown() + agent.collector.collect_snapshot(*args) + incoming_ctx = get_context(tracer, event) + + result = None + with tracer.start_as_current_span( + "aws.lambda.entry", span_context=incoming_ctx + ) as span: + enrich_lambda_span(agent, span, *args) + try: + result = wrapped(*args, **kwargs) - if agent.collector.started: - agent.collector.shutdown() - - return result + if isinstance(result, dict): + server_timing_value = define_server_timing(span.context.trace_id) + if "headers" in result: + result["headers"]["Server-Timing"] = server_timing_value + elif "multiValueHeaders" in result: + result["multiValueHeaders"]["Server-Timing"] = [ + server_timing_value + ] + if "statusCode" in result and result.get("statusCode"): + status_code = int(result["statusCode"]) + span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, status_code) + if 500 <= status_code: + span.record_exception(f"HTTP status {status_code}") + except Exception as exc: + logger.debug(f"AWS Lambda lambda_handler_with_instana error: {exc}") + if span: + exc = traceback.format_exc() + span.record_exception(exc) + raise + finally: + agent.collector.shutdown() + if agent.collector.started: + agent.collector.shutdown() -if env_is_aws_lambda: - handler_module, handler_function = get_aws_lambda_handler() + return result - if handler_module and handler_function: - try: - logger.debug( - f"Instrumenting AWS Lambda handler ({handler_module}.{handler_function})" - ) - sys.path.insert(0, "/var/runtime") - sys.path.insert(0, "/var/task") - wrapt.wrap_function_wrapper( - handler_module, handler_function, lambda_handler_with_instana - ) - except (ModuleNotFoundError, ImportError) as exc: - logger.debug(f"AWS Lambda error: {exc}") + if env_is_aws_lambda: + handler_module, handler_function = get_aws_lambda_handler() + + if handler_module and handler_function: + try: + logger.debug( + f"Instrumenting AWS Lambda handler ({handler_module}.{handler_function})" + ) + sys.path.insert(0, "/var/runtime") + sys.path.insert(0, "/var/task") + wrapt.wrap_function_wrapper( + handler_module, handler_function, lambda_handler_with_instana + ) + except (ModuleNotFoundError, ImportError) as exc: + logger.debug(f"AWS Lambda error: {exc}") + logger.warning( + "Instana: Couldn't instrument AWS Lambda handler. Not monitoring." + ) + else: logger.warning( - "Instana: Couldn't instrument AWS Lambda handler. Not monitoring." + "Instana: Couldn't determine AWS Lambda Handler. Not monitoring." ) - else: - logger.warning( - "Instana: Couldn't determine AWS Lambda Handler. Not monitoring." - ) +except ImportError: + pass diff --git a/src/instana/instrumentation/boto3_inst.py b/src/instana/instrumentation/boto3_inst.py deleted file mode 100644 index 88e1c33f..00000000 --- a/src/instana/instrumentation/boto3_inst.py +++ /dev/null @@ -1,174 +0,0 @@ -# (c) Copyright IBM Corp. 2021 -# (c) Copyright Instana Inc. 2020 - - -import json -import wrapt -import inspect -from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple, Sequence, Type, Optional -from opentelemetry.semconv.trace import SpanAttributes - -from instana.log import logger -from instana.singletons import tracer, agent -from instana.util.traceutils import get_tracer_tuple, tracing_is_off, extract_custom_headers -from instana.propagators.format import Format -from instana.span.span import get_current_span - -if TYPE_CHECKING: - from instana.span.span import InstanaSpan - from botocore.auth import SigV4Auth - from botocore.client import BaseClient - -try: - import boto3 - from boto3.s3 import inject - - def lambda_inject_context(payload: Dict[str, Any], span: "InstanaSpan") -> None: - """ - When boto3 lambda client 'Invoke' is called, we want to inject the tracing context. - boto3/botocore has specific requirements: - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/lambda.html#Lambda.Client.invoke - """ - try: - invoke_payload = payload.get("Payload", {}) - - if not isinstance(invoke_payload, dict): - invoke_payload = json.loads(invoke_payload) - - tracer.inject(span.context, Format.HTTP_HEADERS, invoke_payload) - payload["Payload"] = json.dumps(invoke_payload) - except Exception: - logger.debug("non-fatal lambda_inject_context: ", exc_info=True) - - @wrapt.patch_function_wrapper("botocore.auth", "SigV4Auth.add_auth") - def emit_add_auth_with_instana( - wrapped: Callable[..., None], - instance: "SigV4Auth", - args: Tuple[object], - kwargs: Dict[str, Any], - ) -> Callable[..., None]: - current_span = get_current_span() - if not tracing_is_off() and current_span and current_span.is_recording(): - extract_custom_headers(current_span, args[0].headers) - return wrapped(*args, **kwargs) - - @wrapt.patch_function_wrapper("botocore.client", "BaseClient._make_api_call") - def make_api_call_with_instana( - wrapped: Callable[..., Dict[str, Any]], - instance: Type["BaseClient"], - arg_list: Sequence[Dict[str, Any]], - kwargs: Dict[str, Any], - ) -> Dict[str, Any]: - # If we're not tracing, just return - if tracing_is_off(): - return wrapped(*arg_list, **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("boto3", span_context=parent_context) as span: - try: - operation = arg_list[0] - payload = arg_list[1] - - span.set_attribute("op", operation) - span.set_attribute("ep", instance._endpoint.host) - span.set_attribute("reg", instance._client_config.region_name) - - span.set_attribute( - SpanAttributes.HTTP_URL, - instance._endpoint.host + ":443/" + arg_list[0], - ) - span.set_attribute(SpanAttributes.HTTP_METHOD, "POST") - - # Don't collect payload for SecretsManager - if not hasattr(instance, "get_secret_value"): - span.set_attribute("payload", payload) - - # Inject context when invoking lambdas - if "lambda" in instance._endpoint.host and operation == "Invoke": - lambda_inject_context(payload, span) - - except Exception: - logger.debug("make_api_call_with_instana: collect error", exc_info=True) - - try: - result = wrapped(*arg_list, **kwargs) - - if isinstance(result, dict): - http_dict = result.get("ResponseMetadata") - if isinstance(http_dict, dict): - status = http_dict.get("HTTPStatusCode") - if status is not None: - span.set_attribute("http.status_code", status) - headers = http_dict.get("HTTPHeaders") - extract_custom_headers(span, headers) - - return result - except Exception as exc: - span.mark_as_errored({"error": exc}) - raise - - def s3_inject_method_with_instana( - wrapped: Callable[..., object], - instance: Type["BaseClient"], - arg_list: Sequence[object], - kwargs: Dict[str, Any], - ) -> Callable[..., object]: - # If we're not tracing, just return - if tracing_is_off(): - return wrapped(*arg_list, **kwargs) - - fas = inspect.getfullargspec(wrapped) - fas_args = fas.args - fas_args.remove("self") - - tracer, parent_span, _ = get_tracer_tuple() - - parent_context = parent_span.get_span_context() if parent_span else None - - with tracer.start_as_current_span("boto3", span_context=parent_context) as span: - try: - operation = wrapped.__name__ - span.set_attribute("op", operation) - span.set_attribute("ep", instance._endpoint.host) - span.set_attribute("reg", instance._client_config.region_name) - - span.set_attribute( - SpanAttributes.HTTP_URL, - instance._endpoint.host + ":443/" + operation, - ) - span.set_attribute(SpanAttributes.HTTP_METHOD, "POST") - - arg_length = len(arg_list) - if arg_length > 0: - payload = {} - for index in range(arg_length): - if fas_args[index] in ["Filename", "Bucket", "Key"]: - payload[fas_args[index]] = arg_list[index] - span.set_attribute("payload", payload) - except Exception: - logger.debug( - "s3_inject_method_with_instana: collect error", exc_info=True - ) - - try: - return wrapped(*arg_list, **kwargs) - except Exception as exc: - span.mark_as_errored({"error": exc}) - raise - - for method in [ - "upload_file", - "upload_fileobj", - "download_file", - "download_fileobj", - ]: - wrapt.wrap_function_wrapper( - "boto3.s3.inject", method, s3_inject_method_with_instana - ) - - logger.debug("Instrumenting boto3") -except ImportError: - pass diff --git a/src/instana/instrumentation/cassandra_inst.py b/src/instana/instrumentation/cassandra.py similarity index 100% rename from src/instana/instrumentation/cassandra_inst.py rename to src/instana/instrumentation/cassandra.py diff --git a/src/instana/instrumentation/couchbase_inst.py b/src/instana/instrumentation/couchbase.py similarity index 100% rename from src/instana/instrumentation/couchbase_inst.py rename to src/instana/instrumentation/couchbase.py diff --git a/src/instana/instrumentation/fastapi_inst.py b/src/instana/instrumentation/fastapi.py similarity index 98% rename from src/instana/instrumentation/fastapi_inst.py rename to src/instana/instrumentation/fastapi.py index 5edee85c..b2e9b018 100644 --- a/src/instana/instrumentation/fastapi_inst.py +++ b/src/instana/instrumentation/fastapi.py @@ -29,7 +29,7 @@ from starlette.requests import Request from starlette.responses import Response - if not ( # pragma: no cover + if not ( # pragma: no cover hasattr(fastapi, "__version__") and ( fastapi.__version__[0] > "0" or int(fastapi.__version__.split(".")[1]) >= 51 @@ -84,7 +84,7 @@ def init_with_instana( logger.debug("Instrumenting FastAPI") # Reload GUnicorn when we are instrumenting an already running application - if "INSTANA_MAGIC" in os.environ and running_in_gunicorn(): # pragma: no cover + if "INSTANA_MAGIC" in os.environ and running_in_gunicorn(): # pragma: no cover os.kill(os.getpid(), signal.SIGHUP) except ImportError: diff --git a/src/instana/instrumentation/gevent_inst.py b/src/instana/instrumentation/gevent.py similarity index 65% rename from src/instana/instrumentation/gevent_inst.py rename to src/instana/instrumentation/gevent.py index 93ed3800..c083fb84 100644 --- a/src/instana/instrumentation/gevent_inst.py +++ b/src/instana/instrumentation/gevent.py @@ -11,7 +11,7 @@ def instrument_gevent(): - """ Adds context propagation to gevent greenlet spawning """ + """Adds context propagation to gevent greenlet spawning""" try: logger.debug("Instrumenting gevent") @@ -20,28 +20,34 @@ def instrument_gevent(): from opentracing.scope_managers.gevent import _GeventScope def spawn_callback(new_greenlet): - """ Handles context propagation for newly spawning greenlets """ + """Handles context propagation for newly spawning greenlets""" parent_scope = tracer.scope_manager.active if parent_scope is not None: # New greenlet, new clean slate. Clone and make active in this new greenlet # the currently active scope (but don't finish() the span on close - it's a # clone/not the original and we don't want to close it prematurely) # TODO: Change to our own ScopeManagers - parent_scope_clone = _GeventScope(parent_scope.manager, parent_scope.span, finish_on_close=False) - tracer._scope_manager._set_greenlet_scope(parent_scope_clone, new_greenlet) + parent_scope_clone = _GeventScope( + parent_scope.manager, parent_scope.span, finish_on_close=False + ) + tracer._scope_manager._set_greenlet_scope( + parent_scope_clone, new_greenlet + ) logger.debug(" -> Updating tracer to use gevent based context management") tracer._scope_manager = GeventScopeManager() gevent.Greenlet.add_spawn_callback(spawn_callback) - except: + except Exception: logger.debug("instrument_gevent: ", exc_info=True) -if not 'gevent' in sys.modules: +if "gevent" not in sys.modules: logger.debug("Instrumenting gevent: gevent not detected or loaded. Nothing done.") -elif not hasattr(sys.modules['gevent'], 'version_info'): +elif not hasattr(sys.modules["gevent"], "version_info"): logger.debug("gevent module has no 'version_info'. Skipping instrumentation.") -elif sys.modules['gevent'].version_info < (1, 4): - logger.debug("gevent < 1.4 detected. The Instana package supports gevent versions 1.4 and greater.") +elif sys.modules["gevent"].version_info < (1, 4): + logger.debug( + "gevent < 1.4 detected. The Instana package supports gevent versions 1.4 and greater." + ) else: instrument_gevent() diff --git a/src/instana/instrumentation/sanic_inst.py b/src/instana/instrumentation/sanic.py similarity index 99% rename from src/instana/instrumentation/sanic_inst.py rename to src/instana/instrumentation/sanic.py index 57758a6d..c3c1cac5 100644 --- a/src/instana/instrumentation/sanic_inst.py +++ b/src/instana/instrumentation/sanic.py @@ -21,7 +21,6 @@ from sanic.exceptions import SanicException from opentelemetry import context, trace - from opentelemetry.trace import SpanKind from opentelemetry.semconv.trace import SpanAttributes from instana.singletons import tracer, agent diff --git a/src/instana/instrumentation/starlette_inst.py b/src/instana/instrumentation/starlette.py similarity index 100% rename from src/instana/instrumentation/starlette_inst.py rename to src/instana/instrumentation/starlette.py From fa6e606694d8362f521e067bc05f9e45047c3d6d Mon Sep 17 00:00:00 2001 From: Cagri Yonca Date: Wed, 5 Mar 2025 12:54:26 +0100 Subject: [PATCH 2/3] ft: add instrumentation to dynamodb Signed-off-by: Cagri Yonca --- src/instana/instrumentation/aws/boto3.py | 66 ++-- src/instana/instrumentation/aws/dynamodb.py | 31 ++ src/instana/span/kind.py | 2 + src/instana/span/registered_span.py | 11 + tests/clients/boto3/test_boto3_dynamodb.py | 358 ++++++++++++++++++++ 5 files changed, 433 insertions(+), 35 deletions(-) create mode 100644 src/instana/instrumentation/aws/dynamodb.py create mode 100644 tests/clients/boto3/test_boto3_dynamodb.py diff --git a/src/instana/instrumentation/aws/boto3.py b/src/instana/instrumentation/aws/boto3.py index 7f7c6006..3350a1ac 100644 --- a/src/instana/instrumentation/aws/boto3.py +++ b/src/instana/instrumentation/aws/boto3.py @@ -1,16 +1,18 @@ # (c) Copyright IBM Corp. 2025 +try: + from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Tuple, Type -from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Tuple, Type + from opentelemetry.semconv.trace import SpanAttributes -from opentelemetry.semconv.trace import SpanAttributes + from instana.instrumentation.aws.dynamodb import create_dynamodb_span + from instana.instrumentation.aws.s3 import create_s3_span -if TYPE_CHECKING: - from botocore.auth import SigV4Auth - from botocore.client import BaseClient + if TYPE_CHECKING: + from botocore.auth import SigV4Auth + from botocore.client import BaseClient - from instana.span.span import InstanaSpan + from instana.span.span import InstanaSpan -try: import json import wrapt @@ -69,37 +71,34 @@ def make_api_call_with_instana( parent_context = parent_span.get_span_context() if parent_span else None - try: + if instance.meta.service_model.service_name == "dynamodb": + create_dynamodb_span(wrapped, instance, args, kwargs, parent_context) + elif instance.meta.service_model.service_name == "s3": + create_s3_span(wrapped, instance, args, kwargs, parent_context) + else: with tracer.start_as_current_span( "boto3", span_context=parent_context ) as span: - try: - operation = args[0] - payload = args[1] + operation = args[0] + payload = args[1] - span.set_attribute("op", operation) - span.set_attribute("ep", instance._endpoint.host) - span.set_attribute("reg", instance._client_config.region_name) + span.set_attribute("op", operation) + span.set_attribute("ep", instance._endpoint.host) + span.set_attribute("reg", instance._client_config.region_name) - span.set_attribute( - SpanAttributes.HTTP_URL, - instance._endpoint.host + ":443/" + args[0], - ) - span.set_attribute(SpanAttributes.HTTP_METHOD, "POST") + span.set_attribute( + SpanAttributes.HTTP_URL, + instance._endpoint.host + ":443/" + args[0], + ) + span.set_attribute(SpanAttributes.HTTP_METHOD, "POST") - # Don't collect payload for SecretsManager - if not hasattr(instance, "get_secret_value"): - span.set_attribute("payload", payload) + # Don't collect payload for SecretsManager + if not hasattr(instance, "get_secret_value"): + span.set_attribute("payload", payload) - # Inject context when invoking lambdas - if "lambda" in instance._endpoint.host and operation == "Invoke": - lambda_inject_context(payload, span) - - except Exception: - logger.debug( - "make_api_call_with_instana: collect error", - exc_info=True, - ) + # Inject context when invoking lambdas + if "lambda" in instance._endpoint.host and operation == "Invoke": + lambda_inject_context(payload, span) try: result = wrapped(*args, **kwargs) @@ -117,10 +116,7 @@ def make_api_call_with_instana( except Exception as exc: span.mark_as_errored({"error": exc}) raise - except Exception: - logger.debug("make_api_call_with_instana: collect error", exc_info=True) - else: - return wrapped(*args, **kwargs) + return wrapped(*args, **kwargs) except ImportError: pass diff --git a/src/instana/instrumentation/aws/dynamodb.py b/src/instana/instrumentation/aws/dynamodb.py new file mode 100644 index 00000000..ef9fe251 --- /dev/null +++ b/src/instana/instrumentation/aws/dynamodb.py @@ -0,0 +1,31 @@ +# (c) Copyright IBM Corp. 2025 + +from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Type + +if TYPE_CHECKING: + from botocore.client import BaseClient + +from instana.log import logger +from instana.singletons import tracer +from instana.span_context import SpanContext + + +def create_dynamodb_span( + 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("dynamodb", span_context=parent_context) as span: + try: + span.set_attribute("dynamodb.op", args[0]) + span.set_attribute("dynamodb.region", instance._client_config.region_name) + if "TableName" in args[1].keys(): + span.set_attribute("dynamodb.table", args[1]["TableName"]) + except Exception as exc: + span.record_exception(exc) + logger.debug("create_dynamodb_span: collect error", exc_info=True) + + +logger.debug("Instrumenting DynamoDB") diff --git a/src/instana/span/kind.py b/src/instana/span/kind.py index f3487c39..b93fa207 100644 --- a/src/instana/span/kind.py +++ b/src/instana/span/kind.py @@ -40,6 +40,7 @@ "cassandra", "celery-client", "couchbase", + "dynamodb", "log", "memcache", "mongo", @@ -49,6 +50,7 @@ "redis", "rpc-client", "sqlalchemy", + "s3", "tornado-client", "urllib3", "pymongo", diff --git a/src/instana/span/registered_span.py b/src/instana/span/registered_span.py index 66769ebd..a658f0b8 100644 --- a/src/instana/span/registered_span.py +++ b/src/instana/span/registered_span.py @@ -229,6 +229,13 @@ def _populate_exit_span_data(self, span: "InstanaSpan") -> None: ) self.data["couchbase"]["sql"] = span.attributes.pop("couchbase.sql", None) + elif span.name == "dynamodb": + self.data["dynamodb"]["op"] = span.attributes.pop("dynamodb.op", None) + self.data["dynamodb"]["region"] = span.attributes.pop( + "dynamodb.region", None + ) + self.data["dynamodb"]["table"] = span.attributes.pop("dynamodb.table", None) + elif span.name == "rabbitmq": self.data["rabbitmq"]["exchange"] = span.attributes.pop("exchange", None) self.data["rabbitmq"]["queue"] = span.attributes.pop("queue", None) @@ -253,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) diff --git a/tests/clients/boto3/test_boto3_dynamodb.py b/tests/clients/boto3/test_boto3_dynamodb.py new file mode 100644 index 00000000..649f07df --- /dev/null +++ b/tests/clients/boto3/test_boto3_dynamodb.py @@ -0,0 +1,358 @@ +# (c) Copyright IBM Corp. 2021 +# (c) Copyright Instana Inc. 2020 + +from typing import Generator + +import boto3 +import pytest +from moto import mock_aws + +from instana.singletons import agent, tracer +from tests.helpers import get_first_span_by_filter + + +class TestDynamoDB: + @pytest.fixture(autouse=True) + def _resource(self) -> Generator[None, None, None]: + self.recorder = tracer.span_processor + self.recorder.clear_spans() + self.mock = mock_aws() + self.mock.start() + self.dynamodb = boto3.client("dynamodb", region_name="us-west-2") + yield + self.mock.stop() + agent.options.allow_exit_as_root = False + + def test_vanilla_create_table(self) -> None: + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + result = self.dynamodb.list_tables() + assert len(result["TableNames"]) == 1 + assert result["TableNames"][0] == "dynamodb-table" + + def test_dynamodb_create_table(self) -> None: + with tracer.start_as_current_span("test"): + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + result = self.dynamodb.list_tables() + assert len(result["TableNames"]) == 1 + assert result["TableNames"][0] == "dynamodb-table" + + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "CreateTable" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" + + def test_dynamodb_create_table_as_root_exit_span(self) -> None: + agent.options.allow_exit_as_root = True + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + agent.options.allow_exit_as_root = False + result = self.dynamodb.list_tables() + assert len(result["TableNames"]) == 1 + assert result["TableNames"][0] == "dynamodb-table" + + spans = self.recorder.queued_spans() + assert len(spans) == 1 + dynamodb_span = spans[0] + assert dynamodb_span + assert dynamodb_span.n == "dynamodb" + assert not dynamodb_span.p + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "CreateTable" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" + + def test_dynamodb_list_tables(self) -> None: + with tracer.start_as_current_span("test"): + result = self.dynamodb.list_tables() + + assert len(result["TableNames"]) == 0 + assert result["ResponseMetadata"]["HTTPStatusCode"] == 200 + + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "ListTables" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + + def test_dynamodb_put_item(self) -> None: + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + with tracer.start_as_current_span("test"): + self.dynamodb.put_item( + TableName="dynamodb-table", + Item={"id": {"S": "1"}, "name": {"S": "John"}}, + ) + + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "PutItem" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" + + def test_dynamodb_scan(self) -> None: + test_item = {"id": {"S": "1"}, "name": {"S": "John"}} + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + self.dynamodb.put_item( + TableName="dynamodb-table", + Item=test_item, + ) + with tracer.start_as_current_span("test"): + result = self.dynamodb.scan(TableName="dynamodb-table") + + assert result["Items"] == [test_item] + assert result["Count"] == 1 + assert result["ScannedCount"] == 1 + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "Scan" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" + + def test_dynamodb_get_item(self) -> None: + test_item = {"id": {"S": "1"}, "name": {"S": "John"}} + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + self.dynamodb.put_item( + TableName="dynamodb-table", + Item=test_item, + ) + with tracer.start_as_current_span("test"): + result = self.dynamodb.get_item( + TableName="dynamodb-table", Key={"id": {"S": "1"}} + ) + + assert result["Item"] == test_item + assert result["ResponseMetadata"] + + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "GetItem" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" + + def test_dynamodb_update_item(self) -> None: + test_item = {"id": {"S": "1"}, "name": {"S": "John"}} + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + self.dynamodb.put_item( + TableName="dynamodb-table", + Item=test_item, + ) + with tracer.start_as_current_span("test"): + self.dynamodb.update_item( + TableName="dynamodb-table", + Key={"id": {"S": "1"}}, # Specify the key + UpdateExpression="SET #attr_name = :new_name", + ExpressionAttributeNames={"#attr_name": "name"}, # Use alias for "name" + ExpressionAttributeValues={":new_name": {"S": "Updated John"}}, + ReturnValues="UPDATED_NEW", + ) + + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "UpdateItem" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" + + def test_dynamodb_delete_item(self) -> None: + test_item = {"id": {"S": "1"}, "name": {"S": "John"}} + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + self.dynamodb.put_item( + TableName="dynamodb-table", + Item=test_item, + ) + with tracer.start_as_current_span("test"): + self.dynamodb.delete_item( + TableName="dynamodb-table", Key={"id": {"S": "1"}} + ) + + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "DeleteItem" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" + + def test_dynamodb_query_item(self) -> None: + test_item = {"id": {"S": "1"}, "name": {"S": "John"}} + self.dynamodb.create_table( + TableName="dynamodb-table", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}, + ) + self.dynamodb.put_item( + TableName="dynamodb-table", + Item=test_item, + ) + self.dynamodb.put_item( + TableName="dynamodb-table", Item={"id": {"S": "2"}, "name": {"S": "Jack"}} + ) + with tracer.start_as_current_span("test"): + self.dynamodb.query( + TableName="dynamodb-table", + KeyConditionExpression="id = :pk_val", + ExpressionAttributeValues={":pk_val": {"S": "1"}}, + ) + + spans = self.recorder.queued_spans() + assert len(spans) == 2 + + filter = lambda span: span.n == "sdk" # noqa: E731 + test_span = get_first_span_by_filter(spans, filter) + assert test_span + + filter = lambda span: span.n == "dynamodb" # noqa: E731 + dynamodb_span = get_first_span_by_filter(spans, filter) + assert dynamodb_span + + assert dynamodb_span.t == test_span.t + assert dynamodb_span.p == test_span.s + + assert not test_span.ec + assert not dynamodb_span.ec + + assert dynamodb_span.data["dynamodb"]["op"] == "Query" + assert dynamodb_span.data["dynamodb"]["region"] == "us-west-2" + assert dynamodb_span.data["dynamodb"]["table"] == "dynamodb-table" From 7e10124c7b23e478f7047da07432125c213377fd Mon Sep 17 00:00:00 2001 From: Cagri Yonca Date: Wed, 5 Mar 2025 12:54:34 +0100 Subject: [PATCH 3/3] ft: add instrumentation to s3 Signed-off-by: Cagri Yonca --- src/instana/instrumentation/aws/s3.py | 83 ++++++ tests/clients/boto3/test_boto3_s3.py | 373 ++++++-------------------- 2 files changed, 164 insertions(+), 292 deletions(-) create mode 100644 src/instana/instrumentation/aws/s3.py diff --git a/src/instana/instrumentation/aws/s3.py b/src/instana/instrumentation/aws/s3.py new file mode 100644 index 00000000..932d902a --- /dev/null +++ b/src/instana/instrumentation/aws/s3.py @@ -0,0 +1,83 @@ +# (c) Copyright IBM Corp. 2021 +# (c) Copyright Instana Inc. 2020 + +try: + from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Type + + from instana.span_context import SpanContext + + if TYPE_CHECKING: + from botocore.client import BaseClient + 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 create_s3_span( + 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("create_s3_span: 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( + "collect_s3_injected_attributes: 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 diff --git a/tests/clients/boto3/test_boto3_s3.py b/tests/clients/boto3/test_boto3_s3.py index 6410a6ea..b772ab42 100644 --- a/tests/clients/boto3/test_boto3_s3.py +++ b/tests/clients/boto3/test_boto3_s3.py @@ -50,56 +50,40 @@ def test_s3_create_bucket(self) -> None: spans = self.recorder.queued_spans() assert len(spans) == 2 - filter = lambda span: span.n == "sdk" + filter = lambda span: span.n == "sdk" # noqa: E731 test_span = get_first_span_by_filter(spans, filter) assert test_span - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span + filter = lambda span: span.n == "s3" # noqa: E731 + s3_span = get_first_span_by_filter(spans, filter) + assert s3_span - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s + assert s3_span.t == test_span.t + assert s3_span.p == test_span.s assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "CreateBucket" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - assert boto_span.data["boto3"]["payload"] == {"Bucket": "aws_bucket_name"} - assert boto_span.data["http"]["status"] == 200 - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] == "https://s3.amazonaws.com:443/CreateBucket" - ) + assert not s3_span.ec + + assert s3_span.data["s3"]["op"] == "CreateBucket" + assert s3_span.data["s3"]["bucket"] == "aws_bucket_name" def test_s3_create_bucket_as_root_exit_span(self) -> None: agent.options.allow_exit_as_root = True self.s3.create_bucket(Bucket="aws_bucket_name") agent.options.allow_exit_as_root = False - result = self.s3.list_buckets() - assert len(result["Buckets"]) == 1 - assert result["Buckets"][0]["Name"] == "aws_bucket_name" + self.s3.list_buckets() spans = self.recorder.queued_spans() assert len(spans) == 1 - boto_span = spans[0] - assert boto_span - assert boto_span.n == "boto3" - assert not boto_span.p - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "CreateBucket" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - assert boto_span.data["boto3"]["payload"] == {"Bucket": "aws_bucket_name"} - assert boto_span.data["http"]["status"] == 200 - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] == "https://s3.amazonaws.com:443/CreateBucket" - ) + + s3_span = spans[0] + assert s3_span + + assert not s3_span.ec + + assert s3_span.data["s3"]["op"] == "CreateBucket" + assert s3_span.data["s3"]["bucket"] == "aws_bucket_name" def test_s3_list_buckets(self) -> None: with tracer.start_as_current_span("test"): @@ -111,29 +95,22 @@ def test_s3_list_buckets(self) -> None: spans = self.recorder.queued_spans() assert len(spans) == 2 - filter = lambda span: span.n == "sdk" + filter = lambda span: span.n == "sdk" # noqa: E731 test_span = get_first_span_by_filter(spans, filter) assert test_span - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span + filter = lambda span: span.n == "s3" # noqa: E731 + s3_span = get_first_span_by_filter(spans, filter) + assert s3_span - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s + assert s3_span.t == test_span.t + assert s3_span.p == test_span.s assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "ListBuckets" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - assert boto_span.data["boto3"]["payload"] == {} - assert boto_span.data["http"]["status"] == 200 - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] == "https://s3.amazonaws.com:443/ListBuckets" - ) + assert not s3_span.ec + + assert s3_span.data["s3"]["op"] == "ListBuckets" + assert not s3_span.data["s3"]["bucket"] def test_s3_vanilla_upload_file(self) -> None: object_name = "aws_key_name" @@ -155,33 +132,22 @@ def test_s3_upload_file(self) -> None: spans = self.recorder.queued_spans() assert len(spans) == 2 - filter = lambda span: span.n == "sdk" + filter = lambda span: span.n == "sdk" # noqa: E731 test_span = get_first_span_by_filter(spans, filter) assert test_span - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span + filter = lambda span: span.n == "s3" # noqa: E731 + s3_span = get_first_span_by_filter(spans, filter) + assert s3_span - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s + assert s3_span.t == test_span.t + assert s3_span.p == test_span.s assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "upload_file" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - payload = { - "Filename": upload_filename, - "Bucket": "aws_bucket_name", - "Key": "aws_key_name", - } - assert boto_span.data["boto3"]["payload"] == payload - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] == "https://s3.amazonaws.com:443/upload_file" - ) + assert not s3_span.ec + + assert s3_span.data["s3"]["op"] == "UploadFile" + assert s3_span.data["s3"]["bucket"] == "aws_bucket_name" def test_s3_upload_file_obj(self) -> None: object_name = "aws_key_name" @@ -196,30 +162,22 @@ def test_s3_upload_file_obj(self) -> None: spans = self.recorder.queued_spans() assert len(spans) == 2 - filter = lambda span: span.n == "sdk" + filter = lambda span: span.n == "sdk" # noqa: E731 test_span = get_first_span_by_filter(spans, filter) assert test_span - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span + filter = lambda span: span.n == "s3" # noqa: E731 + s3_span = get_first_span_by_filter(spans, filter) + assert s3_span - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s + assert s3_span.t == test_span.t + assert s3_span.p == test_span.s assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "upload_fileobj" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - payload = {"Bucket": "aws_bucket_name", "Key": "aws_key_name"} - assert boto_span.data["boto3"]["payload"] == payload - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] - == "https://s3.amazonaws.com:443/upload_fileobj" - ) + assert not s3_span.ec + + assert s3_span.data["s3"]["op"] == "UploadFileObj" + assert s3_span.data["s3"]["bucket"] == "aws_bucket_name" def test_s3_download_file(self) -> None: object_name = "aws_key_name" @@ -234,34 +192,22 @@ def test_s3_download_file(self) -> None: spans = self.recorder.queued_spans() assert len(spans) == 2 - filter = lambda span: span.n == "sdk" + filter = lambda span: span.n == "sdk" # noqa: E731 test_span = get_first_span_by_filter(spans, filter) assert test_span - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span + filter = lambda span: span.n == "s3" # noqa: E731 + s3_span = get_first_span_by_filter(spans, filter) + assert s3_span - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s + assert s3_span.t == test_span.t + assert s3_span.p == test_span.s assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "download_file" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - payload = { - "Bucket": "aws_bucket_name", - "Key": "aws_key_name", - "Filename": "%s" % download_target_filename, - } - assert boto_span.data["boto3"]["payload"] == payload - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] - == "https://s3.amazonaws.com:443/download_file" - ) + assert not s3_span.ec + + assert s3_span.data["s3"]["op"] == "DownloadFile" + assert s3_span.data["s3"]["bucket"] == "aws_bucket_name" def test_s3_download_file_obj(self) -> None: object_name = "aws_key_name" @@ -277,204 +223,47 @@ def test_s3_download_file_obj(self) -> None: spans = self.recorder.queued_spans() assert len(spans) == 2 - filter = lambda span: span.n == "sdk" + filter = lambda span: span.n == "sdk" # noqa: E731 test_span = get_first_span_by_filter(spans, filter) assert test_span - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span + filter = lambda span: span.n == "s3" # noqa: E731 + s3_span = get_first_span_by_filter(spans, filter) + assert s3_span - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s + assert s3_span.t == test_span.t + assert s3_span.p == test_span.s assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "download_fileobj" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] - == "https://s3.amazonaws.com:443/download_fileobj" - ) - - def test_request_header_capture_before_call(self) -> None: - original_extra_http_headers = agent.options.extra_http_headers - agent.options.extra_http_headers = ["X-Capture-This", "X-Capture-That"] - - # Access the event system on the S3 client - event_system = self.s3.meta.events - - request_headers = {"X-Capture-This": "this", "X-Capture-That": "that"} - - # Create a function that adds custom headers - def add_custom_header_before_call(params, **kwargs): - params["headers"].update(request_headers) - - # Register the function to before-call event. - event_system.register( - "before-call.s3.CreateBucket", add_custom_header_before_call - ) - - with tracer.start_as_current_span("test"): - self.s3.create_bucket(Bucket="aws_bucket_name") - - result = self.s3.list_buckets() - assert len(result["Buckets"]) == 1 - assert result["Buckets"][0]["Name"] == "aws_bucket_name" - - spans = self.recorder.queued_spans() - assert len(spans) == 2 - - filter = lambda span: span.n == "sdk" - test_span = get_first_span_by_filter(spans, filter) - assert test_span - - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span - - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s - - assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "CreateBucket" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - assert boto_span.data["boto3"]["payload"] == {"Bucket": "aws_bucket_name"} - assert boto_span.data["http"]["status"] == 200 - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] == "https://s3.amazonaws.com:443/CreateBucket" - ) + assert not s3_span.ec - assert "X-Capture-This" in boto_span.data["http"]["header"] - assert boto_span.data["http"]["header"]["X-Capture-This"] == "this" - assert "X-Capture-That" in boto_span.data["http"]["header"] - assert boto_span.data["http"]["header"]["X-Capture-That"] == "that" + assert s3_span.data["s3"]["op"] == "DownloadFileObj" + assert s3_span.data["s3"]["bucket"] == "aws_bucket_name" - agent.options.extra_http_headers = original_extra_http_headers - - def test_request_header_capture_before_sign(self) -> None: - original_extra_http_headers = agent.options.extra_http_headers - agent.options.extra_http_headers = ["X-Custom-1", "X-Custom-2"] - - # Access the event system on the S3 client - event_system = self.s3.meta.events - - request_headers = {"X-Custom-1": "Value1", "X-Custom-2": "Value2"} - - # Create a function that adds custom headers - def add_custom_header_before_sign(request, **kwargs): - for name, value in request_headers.items(): - request.headers.add_header(name, value) + def test_s3_list_obj(self) -> None: + bucket_name = "aws_bucket_name" - # Register the function to before-sign event. - event_system.register_first( - "before-sign.s3.CreateBucket", add_custom_header_before_sign - ) + self.s3.create_bucket(Bucket=bucket_name) with tracer.start_as_current_span("test"): - self.s3.create_bucket(Bucket="aws_bucket_name") - - result = self.s3.list_buckets() - assert len(result["Buckets"]) == 1 - assert result["Buckets"][0]["Name"] == "aws_bucket_name" + self.s3.list_objects(Bucket=bucket_name) spans = self.recorder.queued_spans() assert len(spans) == 2 - filter = lambda span: span.n == "sdk" + filter = lambda span: span.n == "sdk" # noqa: E731 test_span = get_first_span_by_filter(spans, filter) assert test_span - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span + filter = lambda span: span.n == "s3" # noqa: E731 + s3_span = get_first_span_by_filter(spans, filter) + assert s3_span - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s + assert s3_span.t == test_span.t + assert s3_span.p == test_span.s assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "CreateBucket" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - assert boto_span.data["boto3"]["payload"] == {"Bucket": "aws_bucket_name"} - assert boto_span.data["http"]["status"] == 200 - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] == "https://s3.amazonaws.com:443/CreateBucket" - ) - - assert "X-Custom-1" in boto_span.data["http"]["header"] - assert boto_span.data["http"]["header"]["X-Custom-1"] == "Value1" - assert "X-Custom-2" in boto_span.data["http"]["header"] - assert boto_span.data["http"]["header"]["X-Custom-2"] == "Value2" - - agent.options.extra_http_headers = original_extra_http_headers - - def test_response_header_capture(self) -> None: - original_extra_http_headers = agent.options.extra_http_headers - agent.options.extra_http_headers = ["X-Capture-This-Too", "X-Capture-That-Too"] - - # Access the event system on the S3 client - event_system = self.s3.meta.events - - response_headers = { - "X-Capture-This-Too": "this too", - "X-Capture-That-Too": "that too", - } - - # Create a function that sets the custom headers in the after-call event. - def modify_after_call_args(parsed, **kwargs): - parsed["ResponseMetadata"]["HTTPHeaders"].update(response_headers) + assert not s3_span.ec - # Register the function to an event - event_system.register("after-call.s3.CreateBucket", modify_after_call_args) - - with tracer.start_as_current_span("test"): - self.s3.create_bucket(Bucket="aws_bucket_name") - - result = self.s3.list_buckets() - assert len(result["Buckets"]) == 1 - assert result["Buckets"][0]["Name"] == "aws_bucket_name" - - spans = self.recorder.queued_spans() - assert len(spans) == 2 - - filter = lambda span: span.n == "sdk" - test_span = get_first_span_by_filter(spans, filter) - assert test_span - - filter = lambda span: span.n == "boto3" - boto_span = get_first_span_by_filter(spans, filter) - assert boto_span - - assert boto_span.t == test_span.t - assert boto_span.p == test_span.s - - assert not test_span.ec - assert not boto_span.ec - - assert boto_span.data["boto3"]["op"] == "CreateBucket" - assert boto_span.data["boto3"]["ep"] == "https://s3.amazonaws.com" - assert boto_span.data["boto3"]["reg"] == "us-east-1" - assert boto_span.data["boto3"]["payload"] == {"Bucket": "aws_bucket_name"} - assert boto_span.data["http"]["status"] == 200 - assert boto_span.data["http"]["method"] == "POST" - assert ( - boto_span.data["http"]["url"] == "https://s3.amazonaws.com:443/CreateBucket" - ) - - assert "X-Capture-This-Too" in boto_span.data["http"]["header"] - assert boto_span.data["http"]["header"]["X-Capture-This-Too"] == "this too" - assert "X-Capture-That-Too" in boto_span.data["http"]["header"] - assert boto_span.data["http"]["header"]["X-Capture-That-Too"] == "that too" - - agent.options.extra_http_headers = original_extra_http_headers + assert s3_span.data["s3"]["op"] == "ListObjects" + assert s3_span.data["s3"]["bucket"] == "aws_bucket_name"