From 61c9ebb10b0bc3b2752b212acdf68eb2da421817 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Mon, 7 Oct 2024 13:27:16 -0700 Subject: [PATCH] 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