From 7064d944b98b0efd21cf0d0840b8515139266151 Mon Sep 17 00:00:00 2001
From: Hari Prasanna <hari.prasanna@linkedin.com>
Date: Tue, 28 Jun 2016 11:39:44 -0700
Subject: [PATCH] Allow SqlTagger to be able to set a custom tag prefix

---
 lib/sql_tagger.rb       | 19 ++++++++++++++++++-
 spec/sql_tagger_spec.rb | 25 ++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/lib/sql_tagger.rb b/lib/sql_tagger.rb
index bee83cd..e0152d5 100644
--- a/lib/sql_tagger.rb
+++ b/lib/sql_tagger.rb
@@ -7,6 +7,10 @@ class SqlTagger
     File.join(File.dirname(__FILE__), '..', 'VERSION')
   ).chomp.freeze
 
+  # @return [String, #call] a string or a proc which resolves to a string
+  #   to prefix the call trace tag
+  attr_accessor :custom_tag_prefix
+
   # @return [Regexp]
   #   regular expression used to match stack strings we should skip (usually
   #   because such stack strings aren't specific enough, like stack strings
@@ -33,7 +37,7 @@ def tag(sql)
     caller(2).each do |string|
       next if @exclusion_cache.member?(string)
       if string !~ @exclusion_pattern
-        return "/* #{string} */ #{sql}"
+        return "/* #{custom_tag_prefix_string} #{string} */ #{sql}"
       else
         @exclusion_cache.add(string)
       end
@@ -101,4 +105,17 @@ def included(base)
       end
     end
   end
+
+  private
+
+  # Returns a freshly resolved custom tag prefix string
+  #
+  # @return [String] string to prefix the call trace tag
+  def custom_tag_prefix_string
+    if @custom_tag_prefix.respond_to?(:call)
+      @custom_tag_prefix.call
+    else
+      @custom_tag_prefix
+    end
+  end
 end
diff --git a/spec/sql_tagger_spec.rb b/spec/sql_tagger_spec.rb
index b05cb4b..c1c858c 100644
--- a/spec/sql_tagger_spec.rb
+++ b/spec/sql_tagger_spec.rb
@@ -21,7 +21,7 @@
     end
 
     it 'skips stack strings that match @exclusion_pattern' do
-      expect(sql_tagger.tag(sql)).to eq("/* #{valid_stack_string} */ #{sql}")
+      expect(sql_tagger.tag(sql)).to eq("/*  #{valid_stack_string} */ #{sql}")
     end
 
     it 'returns the 1st stack string that does not match @exclusion_pattern' do
@@ -29,7 +29,7 @@
         '/home/app/myapp/lib/document.rb:788',
         '/home/app/myapp/runner.rb:29'
       )
-      expect(sql_tagger.tag(sql)).to eq("/* #{valid_stack_string} */ #{sql}")
+      expect(sql_tagger.tag(sql)).to eq("/*  #{valid_stack_string} */ #{sql}")
     end
 
     it 'adds skipped stack strings into @exclusion_cache' do
@@ -45,7 +45,26 @@
       correct_string = '/home/myapp/i.rb:2890'
       caller_result.push(correct_string)
       sql_tagger.exclusion_cache.add(valid_stack_string)
-      expect(sql_tagger.tag(sql)).to eq("/* #{correct_string} */ #{sql}")
+      expect(sql_tagger.tag(sql)).to eq("/*  #{correct_string} */ #{sql}")
+    end
+
+    context 'when @custom_tag_prefix is set as a proc' do
+      it 'prefixes the tag with a freshly resolved @custom_tag_prefix proc' do
+        string_stack = ['a', 'b']
+        sql_tagger.custom_tag_prefix = proc { string_stack.pop }
+        expect(sql_tagger.tag(sql)).
+          to eq("/* b #{valid_stack_string} */ #{sql}")
+        expect(sql_tagger.tag(sql)).
+          to eq("/* a #{valid_stack_string} */ #{sql}")
+      end
+    end
+
+    context 'when @custom_tag_prefix is set as a string' do
+      it 'prefixes the tag with the set @custom_tag_prefix' do
+        sql_tagger.custom_tag_prefix = 'a'
+        expect(sql_tagger.tag(sql)).
+          to eq("/* a #{valid_stack_string} */ #{sql}")
+      end
     end
   end