Add config option for AWS account id

CHANGELOG.md

@@ -2,7 +2,11 @@
 
 ## dev
 
-Version <dev> resolves a bug in rdkafka instrumentation when using the karafka-rdkafka gem.
+Version <dev> adds a configuration option to associate the AWS account ID with the DynamoDB calls from the AWS sdk, and resolves a bug in rdkafka instrumentation when using the karafka-rdkafka gem.
+
+- **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. 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)
 
 - **Bugfix: Instrumentation errors when using the karafka-rdkafka gem**
 

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

@@ -471,6 +471,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 associated AWS account'
+        },
         :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