From b43929ab908b74ff3be39ccd4628679c285b5caa Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Mon, 11 Nov 2024 16:56:59 -0800 Subject: [PATCH] Add support for AWS DynamoDB entity --- newrelic/core/attribute.py | 2 + newrelic/hooks/external_botocore.py | 130 ++++++++- .../test_botocore_dynamodb.py | 257 ++++++++++-------- 3 files changed, 262 insertions(+), 127 deletions(-) diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py index 16dacb18a..ef04ecb64 100644 --- a/newrelic/core/attribute.py +++ b/newrelic/core/attribute.py @@ -49,6 +49,7 @@ "aws.requestId", "cloud.account.id", "cloud.region", + "cloud.resource_id", "code.filepath", "code.function", "code.lineno", @@ -57,6 +58,7 @@ "db.instance", "db.operation", "db.statement", + "db.system", "enduser.id", "error.class", "error.expected", diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index dc25d6318..37bb65d79 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -22,7 +22,7 @@ from botocore.response import StreamingBody -from newrelic.api.datastore_trace import datastore_trace +from newrelic.api.datastore_trace import DatastoreTrace from newrelic.api.external_trace import ExternalTrace from newrelic.api.function_trace import FunctionTrace from newrelic.api.message_trace import MessageTrace, message_trace @@ -841,6 +841,118 @@ def handle_chat_completion_event(transaction, bedrock_attrs): ) +def dynamodb_datastore_trace( + product, + target, + operation, + host=None, + port_path_or_id=None, + database_name=None, + async_wrapper=None, +): + @function_wrapper + def _nr_dynamodb_datastore_trace_wrapper_(wrapped, instance, args, kwargs): + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) + if not wrapper: + parent = current_trace() + if not parent: + return wrapped(*args, **kwargs) + else: + parent = None + + if callable(product): + if instance is not None: + _product = product(instance, *args, **kwargs) + else: + _product = product(*args, **kwargs) + else: + _product = product + + if callable(target): + if instance is not None: + _target = target(instance, *args, **kwargs) + else: + _target = target(*args, **kwargs) + else: + _target = target + + if callable(operation): + if instance is not None: + _operation = operation(instance, *args, **kwargs) + else: + _operation = operation(*args, **kwargs) + else: + _operation = operation + + if callable(host): + if instance is not None: + _host = host(instance, *args, **kwargs) + else: + _host = host(*args, **kwargs) + else: + _host = host + + if callable(port_path_or_id): + if instance is not None: + _port_path_or_id = port_path_or_id(instance, *args, **kwargs) + else: + _port_path_or_id = port_path_or_id(*args, **kwargs) + else: + _port_path_or_id = port_path_or_id + + if callable(database_name): + if instance is not None: + _database_name = database_name(instance, *args, **kwargs) + else: + _database_name = database_name(*args, **kwargs) + else: + _database_name = database_name + + trace = DatastoreTrace( + _product, _target, _operation, _host, _port_path_or_id, _database_name, parent=parent, source=wrapped + ) + + # Try to capture AWS DynamoDB info as agent attributes. Log any exception to debug. + agent_attrs = {} + try: + region = None + if hasattr(instance, "_client_config") and hasattr(instance._client_config, "region_name"): + region = instance._client_config.region_name + + transaction = current_transaction() + settings = transaction.settings if transaction.settings else global_settings() + account_id = settings.cloud.aws.account_id if settings and settings.cloud.aws.account_id else None + + # There are 3 different partition options. + # See https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html for details. + partition = None + if hasattr(instance, "_endpoint") and hasattr(instance._endpoint, "host"): + _db_host = instance._endpoint.host + partition = "aws" + if "amazonaws.cn" in _db_host: + partition = "aws-cn" + elif "amazonaws-us-gov.com" in _db_host: + partition = "aws-us-gov" + + if partition and region and account_id and _target: + agent_attrs["cloud.resource_id"] = ( + f"arn:{partition}:dynamodb:{region}:{account_id:012d}:table/{_target}" + ) + agent_attrs["db.system"] = "DynamoDB" + + except Exception as e: + _logger.debug("Failed to capture AWS DynamoDB info.", exc_info=True) + trace.agent_attributes.update(agent_attrs) + + if wrapper: # pylint: disable=W0125,W0126 + return wrapper(wrapped, trace)(*args, **kwargs) + + with trace: + return wrapped(*args, **kwargs) + + return _nr_dynamodb_datastore_trace_wrapper_ + + def sqs_message_trace( operation, destination_type, @@ -891,14 +1003,14 @@ def _nr_sqs_message_trace_wrapper_(wrapped, instance, args, kwargs): CUSTOM_TRACE_POINTS = { ("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"), - ("dynamodb", "update_item"): datastore_trace("DynamoDB", extract("TableName"), "update_item"), - ("dynamodb", "delete_item"): datastore_trace("DynamoDB", extract("TableName"), "delete_item"), - ("dynamodb", "create_table"): datastore_trace("DynamoDB", extract("TableName"), "create_table"), - ("dynamodb", "delete_table"): datastore_trace("DynamoDB", extract("TableName"), "delete_table"), - ("dynamodb", "query"): datastore_trace("DynamoDB", extract("TableName"), "query"), - ("dynamodb", "scan"): datastore_trace("DynamoDB", extract("TableName"), "scan"), + ("dynamodb", "put_item"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "put_item"), + ("dynamodb", "get_item"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "get_item"), + ("dynamodb", "update_item"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "update_item"), + ("dynamodb", "delete_item"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "delete_item"), + ("dynamodb", "create_table"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "create_table"), + ("dynamodb", "delete_table"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "delete_table"), + ("dynamodb", "query"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "query"), + ("dynamodb", "scan"): dynamodb_datastore_trace("DynamoDB", extract("TableName"), "scan"), ("sqs", "send_message"): sqs_message_trace( "Produce", "Queue", extract_sqs, extract_agent_attrs=extract_sqs_agent_attrs ), diff --git a/tests/external_botocore/test_botocore_dynamodb.py b/tests/external_botocore/test_botocore_dynamodb.py index c031f543f..7404f425d 100644 --- a/tests/external_botocore/test_botocore_dynamodb.py +++ b/tests/external_botocore/test_botocore_dynamodb.py @@ -15,8 +15,9 @@ import uuid import botocore.session +import pytest from moto import mock_aws -from testing_support.fixtures import dt_enabled +from testing_support.fixtures import dt_enabled, override_application_settings from testing_support.validators.validate_span_events import validate_span_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, @@ -55,121 +56,141 @@ ] -@dt_enabled -@validate_span_events(expected_agents=("aws.requestId",), count=8) -@validate_span_events(exact_agents={"aws.operation": "PutItem"}, count=1) -@validate_span_events(exact_agents={"aws.operation": "GetItem"}, count=1) -@validate_span_events(exact_agents={"aws.operation": "DeleteItem"}, count=1) -@validate_span_events(exact_agents={"aws.operation": "CreateTable"}, count=1) -@validate_span_events(exact_agents={"aws.operation": "DeleteTable"}, count=1) -@validate_span_events(exact_agents={"aws.operation": "Query"}, count=1) -@validate_span_events(exact_agents={"aws.operation": "Scan"}, count=1) -@validate_tt_segment_params(present_params=("aws.requestId",)) -@validate_transaction_metrics( - "test_botocore_dynamodb:test_dynamodb", - scoped_metrics=_dynamodb_scoped_metrics, - rollup_metrics=_dynamodb_rollup_metrics, - background_task=True, -) -@background_task() -@mock_aws -def test_dynamodb(): - session = botocore.session.get_session() - client = session.create_client( - "dynamodb", - region_name=AWS_REGION, - aws_access_key_id=AWS_ACCESS_KEY_ID, - aws_secret_access_key=AWS_SECRET_ACCESS_KEY, - ) - - # Create table - resp = client.create_table( - TableName=TEST_TABLE, - AttributeDefinitions=[ - {"AttributeName": "Id", "AttributeType": "N"}, - {"AttributeName": "Foo", "AttributeType": "S"}, - ], - KeySchema=[ - {"AttributeName": "Id", "KeyType": "HASH"}, - {"AttributeName": "Foo", "KeyType": "RANGE"}, - ], - ProvisionedThroughput={ - "ReadCapacityUnits": 5, - "WriteCapacityUnits": 5, - }, - ) - assert resp["TableDescription"]["TableName"] == TEST_TABLE - # moto response is ACTIVE, AWS response is CREATING - # assert resp['TableDescription']['TableStatus'] == 'ACTIVE' - - # # AWS needs time to create the table - # import time - # time.sleep(15) - - # Put item - resp = client.put_item( - TableName=TEST_TABLE, - Item={ - "Id": {"N": "101"}, - "Foo": {"S": "hello_world"}, - "SomeValue": {"S": "some_random_attribute"}, - }, - ) - # No checking response, due to inconsistent return values. - # moto returns resp['Attributes']. AWS returns resp['ResponseMetadata'] - - # Get item - resp = client.get_item( - TableName=TEST_TABLE, - Key={ - "Id": {"N": "101"}, - "Foo": {"S": "hello_world"}, - }, - ) - assert resp["Item"]["SomeValue"]["S"] == "some_random_attribute" - - # Update item - resp = client.update_item( - TableName=TEST_TABLE, - Key={ - "Id": {"N": "101"}, - "Foo": {"S": "hello_world"}, - }, - AttributeUpdates={ - "Foo2": {"Value": {"S": "hello_world2"}, "Action": "PUT"}, - }, - ReturnValues="ALL_NEW", - ) - assert resp["Attributes"]["Foo2"] - - # Query for item - resp = client.query( - TableName=TEST_TABLE, - Select="ALL_ATTRIBUTES", - KeyConditionExpression="#Id = :v_id", - ExpressionAttributeNames={"#Id": "Id"}, - ExpressionAttributeValues={":v_id": {"N": "101"}}, - ) - assert len(resp["Items"]) == 1 - assert resp["Items"][0]["SomeValue"]["S"] == "some_random_attribute" - - # Scan - resp = client.scan(TableName=TEST_TABLE) - assert len(resp["Items"]) == 1 - - # Delete item - resp = client.delete_item( - TableName=TEST_TABLE, - Key={ - "Id": {"N": "101"}, - "Foo": {"S": "hello_world"}, - }, +@pytest.mark.parametrize("account_id", (None, 12345678901)) +def test_dynamodb(account_id): + expected_aws_agent_attrs = {} + if account_id: + expected_aws_agent_attrs = { + "cloud.resource_id": f"arn:aws:dynamodb:{AWS_REGION}:{account_id:012d}:table/{TEST_TABLE}", + "db.system": "DynamoDB", + } + + @override_application_settings({"cloud.aws.account_id": account_id}) + @dt_enabled + @validate_span_events(expected_agents=("aws.requestId",), count=8) + @validate_span_events(exact_agents={"aws.operation": "PutItem"}, count=1) + @validate_span_events(exact_agents={"aws.operation": "GetItem"}, count=1) + @validate_span_events(exact_agents={"aws.operation": "DeleteItem"}, count=1) + @validate_span_events(exact_agents={"aws.operation": "CreateTable"}, count=1) + @validate_span_events(exact_agents={"aws.operation": "DeleteTable"}, count=1) + @validate_span_events(exact_agents={"aws.operation": "Query"}, count=1) + @validate_span_events(exact_agents={"aws.operation": "Scan"}, count=1) + @validate_tt_segment_params(present_params=("aws.requestId",)) + @validate_transaction_metrics( + "test_botocore_dynamodb:test_dynamodb.._test", + scoped_metrics=_dynamodb_scoped_metrics, + rollup_metrics=_dynamodb_rollup_metrics, + background_task=True, ) - # No checking response, due to inconsistent return values. - # moto returns resp['Attributes']. AWS returns resp['ResponseMetadata'] - - # Delete table - resp = client.delete_table(TableName=TEST_TABLE) - assert resp["TableDescription"]["TableName"] == TEST_TABLE - # moto response is ACTIVE, AWS response is DELETING - # assert resp['TableDescription']['TableStatus'] == 'DELETING' + @background_task() + @mock_aws + def _test(): + session = botocore.session.get_session() + client = session.create_client( + "dynamodb", + region_name=AWS_REGION, + aws_access_key_id=AWS_ACCESS_KEY_ID, + aws_secret_access_key=AWS_SECRET_ACCESS_KEY, + ) + + # Create table + resp = client.create_table( + TableName=TEST_TABLE, + AttributeDefinitions=[ + {"AttributeName": "Id", "AttributeType": "N"}, + {"AttributeName": "Foo", "AttributeType": "S"}, + ], + KeySchema=[ + {"AttributeName": "Id", "KeyType": "HASH"}, + {"AttributeName": "Foo", "KeyType": "RANGE"}, + ], + ProvisionedThroughput={ + "ReadCapacityUnits": 5, + "WriteCapacityUnits": 5, + }, + ) + assert resp["TableDescription"]["TableName"] == TEST_TABLE + # moto response is ACTIVE, AWS response is CREATING + # assert resp['TableDescription']['TableStatus'] == 'ACTIVE' + + # # AWS needs time to create the table + # import time + # time.sleep(15) + + # Put item + resp = client.put_item( + TableName=TEST_TABLE, + Item={ + "Id": {"N": "101"}, + "Foo": {"S": "hello_world"}, + "SomeValue": {"S": "some_random_attribute"}, + }, + ) + # No checking response, due to inconsistent return values. + # moto returns resp['Attributes']. AWS returns resp['ResponseMetadata'] + + # Get item + resp = client.get_item( + TableName=TEST_TABLE, + Key={ + "Id": {"N": "101"}, + "Foo": {"S": "hello_world"}, + }, + ) + assert resp["Item"]["SomeValue"]["S"] == "some_random_attribute" + + # Update item + resp = client.update_item( + TableName=TEST_TABLE, + Key={ + "Id": {"N": "101"}, + "Foo": {"S": "hello_world"}, + }, + AttributeUpdates={ + "Foo2": {"Value": {"S": "hello_world2"}, "Action": "PUT"}, + }, + ReturnValues="ALL_NEW", + ) + assert resp["Attributes"]["Foo2"] + + # Query for item + resp = client.query( + TableName=TEST_TABLE, + Select="ALL_ATTRIBUTES", + KeyConditionExpression="#Id = :v_id", + ExpressionAttributeNames={"#Id": "Id"}, + ExpressionAttributeValues={":v_id": {"N": "101"}}, + ) + assert len(resp["Items"]) == 1 + assert resp["Items"][0]["SomeValue"]["S"] == "some_random_attribute" + + # Scan + resp = client.scan(TableName=TEST_TABLE) + assert len(resp["Items"]) == 1 + + # Delete item + resp = client.delete_item( + TableName=TEST_TABLE, + Key={ + "Id": {"N": "101"}, + "Foo": {"S": "hello_world"}, + }, + ) + # No checking response, due to inconsistent return values. + # moto returns resp['Attributes']. AWS returns resp['ResponseMetadata'] + + # Delete table + resp = client.delete_table(TableName=TEST_TABLE) + assert resp["TableDescription"]["TableName"] == TEST_TABLE + # moto response is ACTIVE, AWS response is DELETING + # assert resp['TableDescription']['TableStatus'] == 'DELETING' + + if account_id: + + @validate_span_events(exact_agents=expected_aws_agent_attrs, count=8) + def _test_apply_validator(): + _test() + + _test_apply_validator() + else: + _test()