From 3977a4d7814ec7111b0482a4043dde6214fd69c6 Mon Sep 17 00:00:00 2001
From: Christoph Lupprich <clupprich@optoro.com>
Date: Thu, 25 Aug 2022 14:07:52 +0200
Subject: [PATCH] Add default_labels config option

---
 google-cloud-trace/lib/google-cloud-trace.rb        |  1 +
 google-cloud-trace/lib/google/cloud/trace.rb        |  2 ++
 .../lib/google/cloud/trace/middleware.rb            | 13 +++++++++++++
 google-cloud-trace/lib/google/cloud/trace/rails.rb  |  4 ++++
 .../test/google/cloud/trace/rails_test.rb           |  3 +++
 5 files changed, 23 insertions(+)

diff --git a/google-cloud-trace/lib/google-cloud-trace.rb b/google-cloud-trace/lib/google-cloud-trace.rb
index d92865d9c69a..d4549d0c63da 100644
--- a/google-cloud-trace/lib/google-cloud-trace.rb
+++ b/google-cloud-trace/lib/google-cloud-trace.rb
@@ -152,4 +152,5 @@ def self.trace project_id = nil,
   config.add_field! :notifications, nil, match: Array
   config.add_field! :max_data_length, nil, match: Integer
   config.add_field! :on_error, nil, match: Proc
+  config.add_field! :default_labels, {}, match: Hash
 end
diff --git a/google-cloud-trace/lib/google/cloud/trace.rb b/google-cloud-trace/lib/google/cloud/trace.rb
index dfac31d8f43b..33b0d76d64a0 100644
--- a/google-cloud-trace/lib/google/cloud/trace.rb
+++ b/google-cloud-trace/lib/google/cloud/trace.rb
@@ -144,6 +144,8 @@ def self.new project_id: nil,
       # * `on_error` - (Proc) A Proc to be run when an error is encountered
       #   during the reporting of traces by the middleware. The Proc must take
       #   the error object as the single argument.
+      # * `default_labels` - (Hash) A Hash that contains the default labels
+      #   which will be added to every root span.
       #
       # See the {file:INSTRUMENTATION.md Configuration Guide} for full
       # configuration parameters.
diff --git a/google-cloud-trace/lib/google/cloud/trace/middleware.rb b/google-cloud-trace/lib/google/cloud/trace/middleware.rb
index 3a92b2fd1982..8bde7725ac08 100644
--- a/google-cloud-trace/lib/google/cloud/trace/middleware.rb
+++ b/google-cloud-trace/lib/google/cloud/trace/middleware.rb
@@ -301,12 +301,25 @@ def get_url env
         #
         def configure_span span, env
           span.name = get_path env
+          set_default_labels span.labels
           set_basic_labels span.labels, env
           set_extended_labels span.labels,
                               span.trace.trace_context.capture_stack?
           span
         end
 
+        ##
+        # Configures default labels.
+        # @private
+        #
+        # rubocop:disable Naming/AccessorMethodName
+        def set_default_labels labels
+          configuration.default_labels.each do |key, value|
+            set_label labels, key, value
+          end
+        end
+        # rubocop:enable Naming/AccessorMethodName
+
         ##
         # Configures standard labels.
         # @private
diff --git a/google-cloud-trace/lib/google/cloud/trace/rails.rb b/google-cloud-trace/lib/google/cloud/trace/rails.rb
index a267aeaa1c1c..8e4a290d1fa4 100644
--- a/google-cloud-trace/lib/google/cloud/trace/rails.rb
+++ b/google-cloud-trace/lib/google/cloud/trace/rails.rb
@@ -85,6 +85,7 @@ class Railtie < ::Rails::Railtie
         config.google_cloud.trace.notifications = DEFAULT_NOTIFICATIONS.dup
         config.google_cloud.trace.max_data_length =
           Google::Cloud::Trace::Notifications::DEFAULT_MAX_DATA_LENGTH
+        config.google_cloud.trace.default_labels = {}
 
         initializer "Google.Cloud.Trace" do |app|
           self.class.consolidate_rails_config app.config
@@ -168,6 +169,9 @@ def self.merge_rails_config rails_config
             end
             config.sampler ||= trace_config.sampler
             config.span_id_generator ||= trace_config.span_id_generator
+            if trace_config.default_labels.present?
+              config.default_labels = trace_config.default_labels
+            end
           end
         end
 
diff --git a/google-cloud-trace/test/google/cloud/trace/rails_test.rb b/google-cloud-trace/test/google/cloud/trace/rails_test.rb
index f1917ac43849..42d971c7951c 100644
--- a/google-cloud-trace/test/google/cloud/trace/rails_test.rb
+++ b/google-cloud-trace/test/google/cloud/trace/rails_test.rb
@@ -58,6 +58,7 @@
           _(config.max_data_length).must_equal 123
           _(config.sampler).must_equal "test-sampler"
           _(config.span_id_generator).must_equal a_test_generator
+          _(config.default_labels).must_equal({})
         end
       end
     end
@@ -71,6 +72,7 @@
         config.max_data_length = 345
         config.sampler = "another-test-sampler"
         config.span_id_generator = another_generator
+        config.default_labels = { '/component' => 'test-service' }
       end
 
       STDOUT.stub :puts, nil do
@@ -84,6 +86,7 @@
           _(config.max_data_length).must_equal 345
           _(config.sampler).must_equal "another-test-sampler"
           _(config.span_id_generator).must_equal another_generator
+          _(config.default_labels).must_equal({ '/component' => 'test-service' })
         end
       end
     end