From 9377e11bd7d787c0f91a16e702fe7703ec89ee8a Mon Sep 17 00:00:00 2001
From: "Chris St. John" <cstjohn@onedoor.com>
Date: Tue, 5 Apr 2022 10:36:31 -0500
Subject: [PATCH] Adds config flag to retain ids in output queries

---
 lib/sql_tracker/config.rb  |  3 ++-
 lib/sql_tracker/handler.rb | 14 ++++++++------
 test/config_test.rb        |  2 ++
 test/handler_test.rb       | 17 +++++++++++++++++
 4 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/lib/sql_tracker/config.rb b/lib/sql_tracker/config.rb
index 76acf7c..c1939dc 100644
--- a/lib/sql_tracker/config.rb
+++ b/lib/sql_tracker/config.rb
@@ -3,7 +3,7 @@
 module SqlTracker
   class Config
     include ActiveSupport::Configurable
-    config_accessor :tracked_paths, :tracked_sql_command, :output_path, :enabled
+    config_accessor :tracked_paths, :tracked_sql_command, :output_path, :enabled, :retain_sql_query_ids
 
     class << self
       def apply_defaults
@@ -17,6 +17,7 @@ def apply_defaults
             'tmp'
           end
         end
+        self.retain_sql_query_ids ||= false
         self
       end
     end
diff --git a/lib/sql_tracker/handler.rb b/lib/sql_tracker/handler.rb
index 9797c36..815c370 100644
--- a/lib/sql_tracker/handler.rb
+++ b/lib/sql_tracker/handler.rb
@@ -62,12 +62,14 @@ def trace_path_matcher
 
     def clean_sql_query(query)
       query.squish!
-      query.gsub!(/(\s(=|>|<|>=|<=|<>|!=)\s)('[^']+'|[\$\+\-\w\.]+)/, '\1xxx')
-      query.gsub!(/(\sIN\s)\([^\(\)]+\)/i, '\1(xxx)')
-      query.gsub!(/(\sBETWEEN\s)('[^']+'|[\+\-\w\.]+)(\sAND\s)('[^']+'|[\+\-\w\.]+)/i, '\1xxx\3xxx')
-      query.gsub!(/(\sVALUES\s)\(.+\)/i, '\1(xxx)')
-      query.gsub!(/(\s(LIKE|ILIKE|SIMILAR TO|NOT SIMILAR TO)\s)('[^']+')/i, '\1xxx')
-      query.gsub!(/(\s(LIMIT|OFFSET)\s)(\d+)/i, '\1xxx')
+      unless @config&.retain_sql_query_ids
+        query.gsub!(/(\s(=|>|<|>=|<=|<>|!=)\s)('[^']+'|[\$\+\-\w\.]+)/, '\1xxx')
+        query.gsub!(/(\sIN\s)\([^\(\)]+\)/i, '\1(xxx)')
+        query.gsub!(/(\sBETWEEN\s)('[^']+'|[\+\-\w\.]+)(\sAND\s)('[^']+'|[\+\-\w\.]+)/i, '\1xxx\3xxx')
+        query.gsub!(/(\sVALUES\s)\(.+\)/i, '\1(xxx)')
+        query.gsub!(/(\s(LIKE|ILIKE|SIMILAR TO|NOT SIMILAR TO)\s)('[^']+')/i, '\1xxx')
+        query.gsub!(/(\s(LIMIT|OFFSET)\s)(\d+)/i, '\1xxx')
+      end
       query
     end
 
diff --git a/test/config_test.rb b/test/config_test.rb
index a40d334..887230c 100644
--- a/test/config_test.rb
+++ b/test/config_test.rb
@@ -11,6 +11,7 @@ def test_defaults_should_not_overwrite_user_configs
       SqlTracker::Config.tracked_paths = %w(app/model)
       SqlTracker::Config.tracked_sql_command = %w(SELECT)
       SqlTracker::Config.output_path = '/usr/local/test'
+      SqlTracker::Config.retain_sql_query_ids = true
 
       config = SqlTracker::Config.apply_defaults
 
@@ -18,6 +19,7 @@ def test_defaults_should_not_overwrite_user_configs
       assert_equal(%w(app/model), config.tracked_paths)
       assert_equal(%w(SELECT), config.tracked_sql_command)
       assert_equal('/usr/local/test', config.output_path)
+      assert_equal(true, config.retain_sql_query_ids)
     end
   end
 end
diff --git a/test/handler_test.rb b/test/handler_test.rb
index a92432f..9a8b8c5 100644
--- a/test/handler_test.rb
+++ b/test/handler_test.rb
@@ -157,6 +157,23 @@ def test_clean_values
       assert_equal(expected, cleaned_query)
     end
 
+    def test_not_clean_values_due_to_config
+      query = %{
+        INSERT INTO users VALUES
+        (nextval('id_seq'), 'a', 105, DEFAULT),
+        (nextval('id_seq'), 'b', 9100, DEFAULT);
+      }
+      expected = %{
+        INSERT INTO users VALUES (nextval('id_seq'), 'a', 105, DEFAULT), (nextval('id_seq'), 'b', 9100, DEFAULT);
+      }.squish
+
+      config = sample_config
+      config.retain_sql_query_ids = true
+      handler = SqlTracker::Handler.new(config)
+      returned_query = handler.clean_sql_query(query)
+      assert_equal(expected, returned_query)
+    end
+
     def test_clean_pattern_matching
       query = %(
         SELECT users.* FROM users