Skip to content

Commit

Permalink
Add attrs for aws lambda entity
Browse files Browse the repository at this point in the history
  • Loading branch information
hmstepanek committed Oct 8, 2024
1 parent 1807857 commit 52e28fb
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 0 deletions.
2 changes: 2 additions & 0 deletions newrelic/core/attribute.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,9 @@
"aws.operation",
"aws.requestId",
"cloud.account.id",
"cloud.platform",
"cloud.region",
"cloud.resource_id",
"code.filepath",
"code.function",
"code.lineno",
Expand Down
24 changes: 24 additions & 0 deletions newrelic/hooks/external_botocore.py
Original file line number Diff line number Diff line change
Expand Up @@ -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+)/([^/]+)")
Expand Down Expand Up @@ -885,7 +887,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"),
Expand Down Expand Up @@ -947,6 +966,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

Expand Down
123 changes: 123 additions & 0 deletions tests/external_botocore/test_boto3_lambda.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions tox.ini
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,7 @@ deps =
external_botocore-botocore128: botocore<1.29
external_botocore-botocore0125: botocore<1.26
external_botocore-{py38,py39,py310,py311,py312}: moto
external_botocore-{py38,py39,py310,py311,py312}: docker
external_feedparser-feedparser06: feedparser<7
external_httplib2: httplib2<1.0
external_httpx: httpx<0.17
Expand Down

0 comments on commit 52e28fb

Please sign in to comment.