update aws instrumentation to use updated access key method

newrelic · Nov 19, 2024 · a33550c · a33550c
1 parent 8ba43d4
commit a33550c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 10 deletions.
diff --git a/lib/new_relic/agent/instrumentation/aws_sdk_lambda/instrumentation.rb b/lib/new_relic/agent/instrumentation/aws_sdk_lambda/instrumentation.rb
@@ -63,10 +63,9 @@ def process_function_error(response)
     def generate_segment(action, options = {})
       function = function_name(options)
       region = aws_region
-      account_id = aws_account_id
       arn = aws_arn(function, region)
       segment = NewRelic::Agent::Tracer.start_segment(name: "Lambda/#{action}/#{function}")
-      segment.add_agent_attribute('cloud.account.id', account_id)
+      segment.add_agent_attribute('cloud.account.id', nr_account_id)
       segment.add_agent_attribute('cloud.platform', CLOUD_PLATFORM)
       segment.add_agent_attribute('cloud.region', region)
       segment.add_agent_attribute('cloud.resource_id', arn) if arn
@@ -77,18 +76,18 @@ def function_name(options = {})
       (options.fetch(:function_name, nil) if options.respond_to?(:fetch)) || NewRelic::UNKNOWN
     end
 
-    def aws_account_id
-      return unless self.respond_to?(:config)
-
-      config&.account_id || NewRelic::Agent.config[:'cloud.aws.account_id']
-    end
-
     def aws_region
       config&.region if self.respond_to?(:config)
     end
 
     def aws_arn(function, region)
-      NewRelic::Agent::Aws.create_arn(AWS_SERVICE, "function:#{function}", region)
+      NewRelic::Agent::Aws.create_arn(AWS_SERVICE, "function:#{function}", region, nr_account_id)
+    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
   end
 end
diff --git a/lib/new_relic/agent/instrumentation/dynamodb/instrumentation.rb b/lib/new_relic/agent/instrumentation/dynamodb/instrumentation.rb
@@ -49,10 +49,16 @@ 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]
 
-      NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config&.region)
+      NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config&.region, nr_account_id)
     end
   end
 end