From 61c9ebb10b0bc3b2752b212acdf68eb2da421817 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Mon, 7 Oct 2024 13:27:16 -0700 Subject: [PATCH 1/7] Add attrs for aws lambda entity --- newrelic/core/attribute.py | 2 + newrelic/hooks/external_botocore.py | 24 ++++ tests/external_botocore/test_boto3_lambda.py | 123 +++++++++++++++++++ tox.ini | 1 + 4 files changed, 150 insertions(+) create mode 100644 tests/external_botocore/test_boto3_lambda.py diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py index 16dacb18a..c705bc652 100644 --- a/newrelic/core/attribute.py +++ b/newrelic/core/attribute.py @@ -48,7 +48,9 @@ "aws.operation", "aws.requestId", "cloud.account.id", + "cloud.platform", "cloud.region", + "cloud.resource_id", "code.filepath", "code.function", "code.lineno", diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index dc25d6318..28703658a 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -30,11 +30,13 @@ from newrelic.api.transaction import current_transaction from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import ( + FunctionWrapper, ObjectProxy, function_wrapper, wrap_function_wrapper, ) from newrelic.common.package_version_utils import get_package_version +from newrelic.common.signature import bind_args from newrelic.core.config import global_settings QUEUE_URL_PATTERN = re.compile(r"https://sqs.([\w\d-]+).amazonaws.com/(\d+)/([^/]+)") @@ -889,7 +891,24 @@ def _nr_sqs_message_trace_wrapper_(wrapped, instance, args, kwargs): return _nr_sqs_message_trace_wrapper_ +def wrap_lambda_invoke(wrapped): + def _nr_wrap_lambda_invoke(wrapped, instance, args, kwargs): + transaction = current_transaction() + if not transaction: + return wrapped(*args, **kwargs) + + bound_args = bind_args(wrapped, args, kwargs) + arn = bound_args["kwargs"].get("FunctionName") + if arn: + transaction._nr_lambda_arn = arn + + return wrapped(*args, **kwargs) + + return FunctionWrapper(wrapped, _nr_wrap_lambda_invoke) + + CUSTOM_TRACE_POINTS = { + ("lambda", "invoke"): wrap_lambda_invoke, ("sns", "publish"): message_trace("SNS", "Produce", "Topic", extract(("TopicArn", "TargetArn"), "PhoneNumber")), ("dynamodb", "put_item"): datastore_trace("DynamoDB", extract("TableName"), "put_item"), ("dynamodb", "get_item"): datastore_trace("DynamoDB", extract("TableName"), "get_item"), @@ -951,6 +970,11 @@ def _nr_endpoint_make_request_(wrapped, instance, args, kwargs): with ExternalTrace(library="botocore", url=url, method=method, source=wrapped) as trace: try: trace._add_agent_attribute("aws.operation", operation_model.name) + lambda_arn = getattr(trace.transaction, "_nr_lambda_arn", None) + if lambda_arn: + trace._add_agent_attribute("cloud.platform", "aws_lambda") + trace._add_agent_attribute("cloud.resource_id", lambda_arn) + del trace.transaction._nr_lambda_arn except: pass diff --git a/tests/external_botocore/test_boto3_lambda.py b/tests/external_botocore/test_boto3_lambda.py new file mode 100644 index 000000000..15f7f48a8 --- /dev/null +++ b/tests/external_botocore/test_boto3_lambda.py @@ -0,0 +1,123 @@ +# Copyright 2010 New Relic, Inc. +# +# 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 io +import json +import zipfile + +import boto3 +import pytest +from moto import mock_aws +from testing_support.fixtures import dt_enabled +from testing_support.validators.validate_span_events import validate_span_events +from testing_support.validators.validate_transaction_metrics import ( + validate_transaction_metrics, +) + +from newrelic.api.background_task import background_task +from newrelic.common.package_version_utils import get_package_version_tuple + +MOTO_VERSION = get_package_version_tuple("moto") +BOTOCORE_VERSION = get_package_version_tuple("botocore") + +AWS_ACCESS_KEY_ID = "AAAAAAAAAAAACCESSKEY" +AWS_SECRET_ACCESS_KEY = "AAAAAASECRETKEY" # nosec +AWS_REGION_NAME = "us-west-2" + +LAMBDA_URL = "lambda.us-west-2.amazonaws.com" +EXPECTED_LAMBDA_URL = f"https://{LAMBDA_URL}/2015-03-31/functions" +LAMBDA_ARN = f"arn:aws:lambda:{AWS_REGION_NAME}:383735328703:function:lambdaFunction" + + +_lambda_scoped_metrics = [ + (f"External/{LAMBDA_URL}/botocore/POST", 2), +] + +_lambda_rollup_metrics = [ + ("External/all", 3), + ("External/allOther", 3), + (f"External/{LAMBDA_URL}/all", 2), + (f"External/{LAMBDA_URL}/botocore/POST", 2), +] + + +@dt_enabled +@validate_span_events(exact_agents={"aws.operation": "CreateFunction"}, count=1) +@validate_span_events( + exact_agents={"aws.operation": "Invoke", "cloud.platform": "aws_lambda", "cloud.resource_id": LAMBDA_ARN}, count=1 +) +@validate_span_events(exact_agents={"aws.operation": "Invoke"}, count=1) +@validate_span_events(exact_agents={"http.url": EXPECTED_LAMBDA_URL}, count=1) +@validate_transaction_metrics( + "test_boto3_lambda:test_lambda", + scoped_metrics=_lambda_scoped_metrics, + rollup_metrics=_lambda_rollup_metrics, + background_task=True, +) +@background_task() +@mock_aws +def test_lambda(iam_role_arn, lambda_zip): + role_arn = iam_role_arn() + + client = boto3.client( + "lambda", + aws_access_key_id=AWS_ACCESS_KEY_ID, + aws_secret_access_key=AWS_SECRET_ACCESS_KEY, + region_name=AWS_REGION_NAME, + ) + + # Create lambda + resp = client.create_function( + FunctionName="lambdaFunction", Runtime="python3.9", Role=role_arn, Code={"ZipFile": lambda_zip} + ) + assert resp["ResponseMetadata"]["HTTPStatusCode"] == 201 + + # Invoke lambda + client.invoke(FunctionName=LAMBDA_ARN, InvocationType="RequestResponse", Payload=json.dumps({})) + assert resp["ResponseMetadata"]["HTTPStatusCode"] == 201 + + +@pytest.fixture +def lambda_zip(): + code = """ + def lambda_handler(event, context): + return event + """ + zip_output = io.BytesIO() + zip_file = zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED) + zip_file.writestr("lambda_function.py", code) + zip_file.close() + zip_output.seek(0) + return zip_output.read() + + +@pytest.fixture +def iam_role_arn(): + def create_role(): + iam = boto3.client( + "iam", + aws_access_key_id=AWS_ACCESS_KEY_ID, + aws_secret_access_key=AWS_SECRET_ACCESS_KEY, + region_name=AWS_REGION_NAME, + ) + # Create IAM role + return iam.create_role( + RoleName="my-role", + AssumeRolePolicyDocument="some policy", + Path="/my-path/", + )[ + "Role" + ]["Arn"] + + return create_role diff --git a/tox.ini b/tox.ini index df74a12a1..afc363b19 100644 --- a/tox.ini +++ b/tox.ini @@ -300,6 +300,7 @@ deps = external_botocore-botocore128: botocore<1.29 external_botocore-botocore0125: botocore<1.26 external_botocore: moto + external_botocore: docker external_feedparser-feedparser06: feedparser<7 external_httplib2: httplib2<1.0 external_httpx: httpx<0.17 From 84d977a4f72c4ebe48f7e6fa8eaac19705594858 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Tue, 19 Nov 2024 12:41:56 -0800 Subject: [PATCH 2/7] Store arn on events obj instead of transaction --- newrelic/hooks/external_botocore.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 28703658a..6ca7f1bd8 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -899,8 +899,8 @@ def _nr_wrap_lambda_invoke(wrapped, instance, args, kwargs): bound_args = bind_args(wrapped, args, kwargs) arn = bound_args["kwargs"].get("FunctionName") - if arn: - transaction._nr_lambda_arn = arn + if arn and hasattr(instance, "meta") and hasattr(instance.meta, "events"): + instance.meta.events._nr_lambda_arn = arn return wrapped(*args, **kwargs) @@ -970,11 +970,10 @@ def _nr_endpoint_make_request_(wrapped, instance, args, kwargs): with ExternalTrace(library="botocore", url=url, method=method, source=wrapped) as trace: try: trace._add_agent_attribute("aws.operation", operation_model.name) - lambda_arn = getattr(trace.transaction, "_nr_lambda_arn", None) + lambda_arn = getattr(instance._event_emitter, "_nr_lambda_arn", None) if lambda_arn: trace._add_agent_attribute("cloud.platform", "aws_lambda") trace._add_agent_attribute("cloud.resource_id", lambda_arn) - del trace.transaction._nr_lambda_arn except: pass From 678a35ca92b64c5d0972a3e09da650cc04e75ed1 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Fri, 22 Nov 2024 10:59:57 -0800 Subject: [PATCH 3/7] Reformat iam create role Co-authored-by: Timothy Pansino <11214426+TimPansino@users.noreply.github.com> --- tests/external_botocore/test_boto3_lambda.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/external_botocore/test_boto3_lambda.py b/tests/external_botocore/test_boto3_lambda.py index 15f7f48a8..edc292c93 100644 --- a/tests/external_botocore/test_boto3_lambda.py +++ b/tests/external_botocore/test_boto3_lambda.py @@ -112,12 +112,11 @@ def create_role(): region_name=AWS_REGION_NAME, ) # Create IAM role - return iam.create_role( + role = iam.create_role( RoleName="my-role", AssumeRolePolicyDocument="some policy", Path="/my-path/", - )[ - "Role" - ]["Arn"] + ) + return role["Role"]["Arn"] return create_role From 5c97248f89afdb97bb4c26de0ba68f9c3f0df096 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Mon, 7 Oct 2024 13:27:16 -0700 Subject: [PATCH 4/7] Add attrs for aws lambda entity --- newrelic/core/attribute.py | 2 + newrelic/hooks/external_botocore.py | 24 ++++ tests/external_botocore/test_boto3_lambda.py | 123 +++++++++++++++++++ tox.ini | 1 + 4 files changed, 150 insertions(+) create mode 100644 tests/external_botocore/test_boto3_lambda.py diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py index 16dacb18a..c705bc652 100644 --- a/newrelic/core/attribute.py +++ b/newrelic/core/attribute.py @@ -48,7 +48,9 @@ "aws.operation", "aws.requestId", "cloud.account.id", + "cloud.platform", "cloud.region", + "cloud.resource_id", "code.filepath", "code.function", "code.lineno", diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index dc25d6318..28703658a 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -30,11 +30,13 @@ from newrelic.api.transaction import current_transaction from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import ( + FunctionWrapper, ObjectProxy, function_wrapper, wrap_function_wrapper, ) from newrelic.common.package_version_utils import get_package_version +from newrelic.common.signature import bind_args from newrelic.core.config import global_settings QUEUE_URL_PATTERN = re.compile(r"https://sqs.([\w\d-]+).amazonaws.com/(\d+)/([^/]+)") @@ -889,7 +891,24 @@ def _nr_sqs_message_trace_wrapper_(wrapped, instance, args, kwargs): return _nr_sqs_message_trace_wrapper_ +def wrap_lambda_invoke(wrapped): + def _nr_wrap_lambda_invoke(wrapped, instance, args, kwargs): + transaction = current_transaction() + if not transaction: + return wrapped(*args, **kwargs) + + bound_args = bind_args(wrapped, args, kwargs) + arn = bound_args["kwargs"].get("FunctionName") + if arn: + transaction._nr_lambda_arn = arn + + return wrapped(*args, **kwargs) + + return FunctionWrapper(wrapped, _nr_wrap_lambda_invoke) + + CUSTOM_TRACE_POINTS = { + ("lambda", "invoke"): wrap_lambda_invoke, ("sns", "publish"): message_trace("SNS", "Produce", "Topic", extract(("TopicArn", "TargetArn"), "PhoneNumber")), ("dynamodb", "put_item"): datastore_trace("DynamoDB", extract("TableName"), "put_item"), ("dynamodb", "get_item"): datastore_trace("DynamoDB", extract("TableName"), "get_item"), @@ -951,6 +970,11 @@ def _nr_endpoint_make_request_(wrapped, instance, args, kwargs): with ExternalTrace(library="botocore", url=url, method=method, source=wrapped) as trace: try: trace._add_agent_attribute("aws.operation", operation_model.name) + lambda_arn = getattr(trace.transaction, "_nr_lambda_arn", None) + if lambda_arn: + trace._add_agent_attribute("cloud.platform", "aws_lambda") + trace._add_agent_attribute("cloud.resource_id", lambda_arn) + del trace.transaction._nr_lambda_arn except: pass diff --git a/tests/external_botocore/test_boto3_lambda.py b/tests/external_botocore/test_boto3_lambda.py new file mode 100644 index 000000000..15f7f48a8 --- /dev/null +++ b/tests/external_botocore/test_boto3_lambda.py @@ -0,0 +1,123 @@ +# Copyright 2010 New Relic, Inc. +# +# 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 io +import json +import zipfile + +import boto3 +import pytest +from moto import mock_aws +from testing_support.fixtures import dt_enabled +from testing_support.validators.validate_span_events import validate_span_events +from testing_support.validators.validate_transaction_metrics import ( + validate_transaction_metrics, +) + +from newrelic.api.background_task import background_task +from newrelic.common.package_version_utils import get_package_version_tuple + +MOTO_VERSION = get_package_version_tuple("moto") +BOTOCORE_VERSION = get_package_version_tuple("botocore") + +AWS_ACCESS_KEY_ID = "AAAAAAAAAAAACCESSKEY" +AWS_SECRET_ACCESS_KEY = "AAAAAASECRETKEY" # nosec +AWS_REGION_NAME = "us-west-2" + +LAMBDA_URL = "lambda.us-west-2.amazonaws.com" +EXPECTED_LAMBDA_URL = f"https://{LAMBDA_URL}/2015-03-31/functions" +LAMBDA_ARN = f"arn:aws:lambda:{AWS_REGION_NAME}:383735328703:function:lambdaFunction" + + +_lambda_scoped_metrics = [ + (f"External/{LAMBDA_URL}/botocore/POST", 2), +] + +_lambda_rollup_metrics = [ + ("External/all", 3), + ("External/allOther", 3), + (f"External/{LAMBDA_URL}/all", 2), + (f"External/{LAMBDA_URL}/botocore/POST", 2), +] + + +@dt_enabled +@validate_span_events(exact_agents={"aws.operation": "CreateFunction"}, count=1) +@validate_span_events( + exact_agents={"aws.operation": "Invoke", "cloud.platform": "aws_lambda", "cloud.resource_id": LAMBDA_ARN}, count=1 +) +@validate_span_events(exact_agents={"aws.operation": "Invoke"}, count=1) +@validate_span_events(exact_agents={"http.url": EXPECTED_LAMBDA_URL}, count=1) +@validate_transaction_metrics( + "test_boto3_lambda:test_lambda", + scoped_metrics=_lambda_scoped_metrics, + rollup_metrics=_lambda_rollup_metrics, + background_task=True, +) +@background_task() +@mock_aws +def test_lambda(iam_role_arn, lambda_zip): + role_arn = iam_role_arn() + + client = boto3.client( + "lambda", + aws_access_key_id=AWS_ACCESS_KEY_ID, + aws_secret_access_key=AWS_SECRET_ACCESS_KEY, + region_name=AWS_REGION_NAME, + ) + + # Create lambda + resp = client.create_function( + FunctionName="lambdaFunction", Runtime="python3.9", Role=role_arn, Code={"ZipFile": lambda_zip} + ) + assert resp["ResponseMetadata"]["HTTPStatusCode"] == 201 + + # Invoke lambda + client.invoke(FunctionName=LAMBDA_ARN, InvocationType="RequestResponse", Payload=json.dumps({})) + assert resp["ResponseMetadata"]["HTTPStatusCode"] == 201 + + +@pytest.fixture +def lambda_zip(): + code = """ + def lambda_handler(event, context): + return event + """ + zip_output = io.BytesIO() + zip_file = zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED) + zip_file.writestr("lambda_function.py", code) + zip_file.close() + zip_output.seek(0) + return zip_output.read() + + +@pytest.fixture +def iam_role_arn(): + def create_role(): + iam = boto3.client( + "iam", + aws_access_key_id=AWS_ACCESS_KEY_ID, + aws_secret_access_key=AWS_SECRET_ACCESS_KEY, + region_name=AWS_REGION_NAME, + ) + # Create IAM role + return iam.create_role( + RoleName="my-role", + AssumeRolePolicyDocument="some policy", + Path="/my-path/", + )[ + "Role" + ]["Arn"] + + return create_role diff --git a/tox.ini b/tox.ini index 335ea0776..9096c2d67 100644 --- a/tox.ini +++ b/tox.ini @@ -293,6 +293,7 @@ deps = external_botocore-botocore128: botocore<1.29 external_botocore-botocore0125: botocore<1.26 external_botocore: moto + external_botocore: docker external_feedparser-feedparser06: feedparser<7 external_httplib2: httplib2<1.0 external_httpx: httpx[http2] From 71ce40807ca3ba9e8687f72afb47f92f998c22f4 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Tue, 19 Nov 2024 12:41:56 -0800 Subject: [PATCH 5/7] Store arn on events obj instead of transaction --- newrelic/hooks/external_botocore.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 28703658a..6ca7f1bd8 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -899,8 +899,8 @@ def _nr_wrap_lambda_invoke(wrapped, instance, args, kwargs): bound_args = bind_args(wrapped, args, kwargs) arn = bound_args["kwargs"].get("FunctionName") - if arn: - transaction._nr_lambda_arn = arn + if arn and hasattr(instance, "meta") and hasattr(instance.meta, "events"): + instance.meta.events._nr_lambda_arn = arn return wrapped(*args, **kwargs) @@ -970,11 +970,10 @@ def _nr_endpoint_make_request_(wrapped, instance, args, kwargs): with ExternalTrace(library="botocore", url=url, method=method, source=wrapped) as trace: try: trace._add_agent_attribute("aws.operation", operation_model.name) - lambda_arn = getattr(trace.transaction, "_nr_lambda_arn", None) + lambda_arn = getattr(instance._event_emitter, "_nr_lambda_arn", None) if lambda_arn: trace._add_agent_attribute("cloud.platform", "aws_lambda") trace._add_agent_attribute("cloud.resource_id", lambda_arn) - del trace.transaction._nr_lambda_arn except: pass From 86c83620e91bddd9ce04381812c178e026607005 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Fri, 22 Nov 2024 10:59:57 -0800 Subject: [PATCH 6/7] Reformat iam create role Co-authored-by: Timothy Pansino <11214426+TimPansino@users.noreply.github.com> --- tests/external_botocore/test_boto3_lambda.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/external_botocore/test_boto3_lambda.py b/tests/external_botocore/test_boto3_lambda.py index 15f7f48a8..5ae657e03 100644 --- a/tests/external_botocore/test_boto3_lambda.py +++ b/tests/external_botocore/test_boto3_lambda.py @@ -112,12 +112,11 @@ def create_role(): region_name=AWS_REGION_NAME, ) # Create IAM role - return iam.create_role( + role = iam.create_role( RoleName="my-role", AssumeRolePolicyDocument="some policy", Path="/my-path/", - )[ - "Role" - ]["Arn"] + ) + return role["Role"]["Arn"] return create_role From d6347c90137e49d2bf5c430bb182a9868cc0f6e7 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Fri, 22 Nov 2024 15:24:50 -0800 Subject: [PATCH 7/7] Attach arn to api_params & request_dict These are regenerated for each request as opposed to the events obj that persists across multiple requests. Unfortunately, we have to pass the arn from api_params to the request_dict because: 1. api_params has validation key checks that exclude unknown params upon generating the request_dict 2. the arn is not available after api_params are emitted because emit drops the param that contains the arn So...we capture the arn before it is lost when the api_params are emitted and pop it off before generating the request dict and put it onto the output request dict to be popped off later by the make request method and added to the trace to link the entity. --- newrelic/hooks/external_botocore.py | 49 ++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 6ca7f1bd8..fd0c1aa6a 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -30,7 +30,6 @@ from newrelic.api.transaction import current_transaction from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import ( - FunctionWrapper, ObjectProxy, function_wrapper, wrap_function_wrapper, @@ -891,24 +890,39 @@ def _nr_sqs_message_trace_wrapper_(wrapped, instance, args, kwargs): return _nr_sqs_message_trace_wrapper_ -def wrap_lambda_invoke(wrapped): - def _nr_wrap_lambda_invoke(wrapped, instance, args, kwargs): - transaction = current_transaction() - if not transaction: - return wrapped(*args, **kwargs) +def wrap_emit_api_params(wrapped, instance, args, kwargs): + transaction = current_transaction() + if not transaction: + return wrapped(*args, **kwargs) + + bound_args = bind_args(wrapped, args, kwargs) + + api_params = wrapped(*args, **kwargs) - bound_args = bind_args(wrapped, args, kwargs) - arn = bound_args["kwargs"].get("FunctionName") - if arn and hasattr(instance, "meta") and hasattr(instance.meta, "events"): - instance.meta.events._nr_lambda_arn = arn + arn = bound_args.get("api_params").get("FunctionName") + if arn and hasattr(arn, "startswith") and arn.startswith("arn:"): + api_params["_nr_arn"] = arn + return api_params + + +def wrap_convert_to_request_dict(wrapped, instance, args, kwargs): + transaction = current_transaction() + if not transaction: return wrapped(*args, **kwargs) - return FunctionWrapper(wrapped, _nr_wrap_lambda_invoke) + bound_args = bind_args(wrapped, args, kwargs) + arn = bound_args.get("api_params").pop("_nr_arn", None) + + request_dict = wrapped(*args, **kwargs) + + if arn: + request_dict["_nr_arn"] = arn + + return request_dict CUSTOM_TRACE_POINTS = { - ("lambda", "invoke"): wrap_lambda_invoke, ("sns", "publish"): message_trace("SNS", "Produce", "Topic", extract(("TopicArn", "TargetArn"), "PhoneNumber")), ("dynamodb", "put_item"): datastore_trace("DynamoDB", extract("TableName"), "put_item"), ("dynamodb", "get_item"): datastore_trace("DynamoDB", extract("TableName"), "get_item"), @@ -970,7 +984,8 @@ def _nr_endpoint_make_request_(wrapped, instance, args, kwargs): with ExternalTrace(library="botocore", url=url, method=method, source=wrapped) as trace: try: trace._add_agent_attribute("aws.operation", operation_model.name) - lambda_arn = getattr(instance._event_emitter, "_nr_lambda_arn", None) + bound_args = bind_args(wrapped, args, kwargs) + lambda_arn = bound_args.get("request_dict").pop("_nr_arn", None) if lambda_arn: trace._add_agent_attribute("cloud.platform", "aws_lambda") trace._add_agent_attribute("cloud.resource_id", lambda_arn) @@ -991,5 +1006,9 @@ def instrument_botocore_endpoint(module): def instrument_botocore_client(module): - wrap_function_wrapper(module, "ClientCreator._create_api_method", _nr_clientcreator__create_api_method_) - wrap_function_wrapper(module, "ClientCreator._create_methods", _nr_clientcreator__create_methods) + if hasattr(module, "ClientCreator"): + wrap_function_wrapper(module, "ClientCreator._create_api_method", _nr_clientcreator__create_api_method_) + wrap_function_wrapper(module, "ClientCreator._create_methods", _nr_clientcreator__create_methods) + if hasattr(module, "BaseClient"): + wrap_function_wrapper(module, "BaseClient._convert_to_request_dict", wrap_convert_to_request_dict) + wrap_function_wrapper(module, "BaseClient._emit_api_params", wrap_emit_api_params)