Merge pull request #2904 from newrelic/add_aws_account_id_config

Add config option for AWS account id

CHANGELOG.md

@@ -2,7 +2,12 @@
 
 ## dev
 
-Version <dev> updates View Componment instrumentation to use a default metric name when one is unavailable, resolves a bug in rdkafka instrumentation when using the karafka-rdkafka gem, resolves a bug in the ruby-kafka instrumentation, and fixes a bug with Grape instrumentation.
+Version <dev> updates View Componment instrumentation to use a default metric name when one is unavailable, adds a configuration option to associate the AWS account ID with the DynamoDB calls from the AWS SDK, resolves a bug in rdkafka instrumentation when using the karafka-rdkafka gem, resolves a bug in the ruby-kafka instrumentation, and fixes a bug with Grape instrumentation.
+
+
+- **Feature: New configuration option cloud.aws.account_id**
+
+  A new configuration option has been added, `cloud.aws.account_id`, that will allow New Relic to provide more details about certain calls made using the AWS SDK. One example, is that relationships between AWS services instrumented with New Relic's CloudWatch Metric Streams will have relationships formed in the service map with APM applications. Currently, the DynamoDB instrumentation is the only instrumentation that will make use of this configuration option, but this will be used in future instrumentation as well. [PR#2904](https://github.com/newrelic/newrelic-ruby-agent/pull/2904)
 
 - **Feature: Use default `View/component` metric name for unidentified View Components**
 

lib/new_relic/agent/aws.rb

@@ -5,58 +5,13 @@
 module NewRelic
   module Agent
     module Aws
-      CHARACTERS = %w[A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 2 3 4 5 6 7].freeze
-      HEX_MASK = '7fffffffff80'
+      def self.create_arn(service, resource, region)
+        return unless NewRelic::Agent.config[:'cloud.aws.account_id']
 
-      def self.create_arn(service, resource, region, account_id)
-        "arn:aws:#{service}:#{region}:#{account_id}:#{resource}"
+        "arn:aws:#{service}:#{region}:#{NewRelic::Agent.config[:'cloud.aws.account_id']}:#{resource}"
       rescue => e
         NewRelic::Agent.logger.warn("Failed to create ARN: #{e}")
       end
-
-      def self.get_account_id(config)
-        access_key_id = config.credentials.credentials.access_key_id if config&.credentials&.credentials&.respond_to?(:access_key_id)
-        return unless access_key_id
-
-        NewRelic::Agent::Aws.convert_access_key_to_account_id(access_key_id)
-      rescue => e
-        NewRelic::Agent.logger.debug("Failed to create account id: #{e}")
-      end
-
-      def self.convert_access_key_to_account_id(access_key)
-        decoded_key = Integer(decode_to_hex(access_key[4..-1]), 16)
-        mask = Integer(HEX_MASK, 16)
-        (decoded_key & mask) >> 7
-      end
-
-      def self.decode_to_hex(access_key)
-        bytes = access_key.delete('=').each_char.map { |c| CHARACTERS.index(c) }
-
-        bytes.each_slice(8).map do |section|
-          convert_section(section)
-        end.flatten[0...6].join
-      end
-
-      def self.convert_section(section)
-        buffer = 0
-        section.each do |chunk|
-          buffer = (buffer << 5) + chunk
-        end
-
-        chunk_count = (section.length * 5.0 / 8.0).floor
-
-        if section.length < 8
-          buffer >>= (5 - (chunk_count * 8)) % 5
-        end
-
-        decoded = []
-        chunk_count.times do |i|
-          shift = 8 * (chunk_count - 1 - i)
-          decoded << ((buffer >> shift) & 255).to_s(16)
-        end
-
-        decoded
-      end
     end
   end
 end

lib/new_relic/agent/configuration/default_source.rb

@@ -472,6 +472,14 @@ def self.enforce_fallback(allowed_values: nil, fallback: nil)
           :allowed_from_server => false,
           :description => 'If `true`, the agent will clear `Tracer::State` in `Agent.drop_buffered_data`.'
         },
+        :'cloud.aws.account_id' => {
+          :default => nil,
+          :public => true,
+          :type => String,
+          :allow_nil => true,
+          :allowed_from_server => false,
+          :description => 'The AWS account ID for the AWS account associated with this app'
+        },
         :config_path => {
           :default => DefaultSource.config_path,
           :public => true,

lib/new_relic/agent/instrumentation/dynamodb/instrumentation.rb

@@ -31,9 +31,8 @@ def instrument_method_with_new_relic(method_name, *args)
         collection: args[0][:table_name]
       )
 
-      # TODO: Update this when it has been decided how to handle account id for ARN
-      # arn = get_arn(args[0])
-      # segment&.add_agent_attribute('cloud.resource_id', arn) if arn
+      arn = get_arn(args[0])
+      segment&.add_agent_attribute('cloud.resource_id', arn) if arn
 
       @nr_captured_request = nil # clear request just in case
       begin
@@ -50,16 +49,10 @@ def build_request_with_new_relic(*args)
       @nr_captured_request = yield
     end
 
-    def nr_account_id
-      return @nr_account_id if defined?(@nr_account_id)
-
-      @nr_account_id = NewRelic::Agent::Aws.get_account_id(config)
-    end
-
     def get_arn(params)
-      return unless params[:table_name] && nr_account_id
+      return unless params[:table_name]
 
-      NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config&.region, nr_account_id)
+      NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config&.region)
     end
   end
 end

test/multiverse/suites/dynamodb/dynamodb_instrumentation_test.rb

@@ -8,7 +8,6 @@ class DynamodbInstrumentationTest < Minitest::Test
   def setup
     Aws.config.update(stub_responses: true)
     NewRelic::Agent::Aws.stubs(:create_arn).returns('test-arn')
-    NewRelic::Agent::Aws.stubs(:get_account_id).returns('123456789')
     @stats_engine = NewRelic::Agent.instance.stats_engine
   end
 
@@ -42,8 +41,7 @@ def test_all_attributes_added_to_segment
     assert_equal 'us-east-2', span[2]['aws.region']
     assert_equal 'query', span[2]['aws.operation']
     assert_equal '1234321', span[2]['aws.requestId']
-    # TODO: Uncomment this when the ARN is added to the segment
-    # assert_equal 'test-arn', span[2]['cloud.resource_id']
+    assert_equal 'test-arn', span[2]['cloud.resource_id']
   end
 
   def test_create_table_table_name_operation

test/new_relic/agent/aws_test.rb

@@ -10,22 +10,16 @@ def test_create_arn
     region = 'us-test-region-1'
     account_id = '123456789'
     resource = 'test/test-resource'
-    arn = NewRelic::Agent::Aws.create_arn(service, resource, region, account_id)
-
     expected = 'arn:aws:test-service:us-test-region-1:123456789:test/test-resource'
 
-    assert_equal expected, arn
-  end
-
-  def test_get_account_id
-    config = mock
-    mock_credentials = mock
-    mock_credentials.stubs(:credentials).returns(mock_credentials)
-    mock_credentials.stubs(:access_key_id).returns('AKIAIOSFODNN7EXAMPLE') # this is a fake access key id from aws docs
-    config.stubs(:credentials).returns(mock_credentials)
+    with_config('cloud.aws.account_id': account_id) do
+      arn = NewRelic::Agent::Aws.create_arn(service, resource, region)
 
-    account_id = NewRelic::Agent::Aws.get_account_id(config)
+      assert_equal expected, arn
+    end
+  end
 
-    assert_equal 36315003739, account_id
+  def test_doesnt_create_arn_no_account_id
+    assert_nil NewRelic::Agent::Aws.create_arn('service', 'resource', 'region')
   end
 end