Merge branch 'sh-fix-active-record-regex' into 'master'

Fix ActiveRecord metrics to handle comments

See merge request gitlab-org/gitlab!54364

lib/gitlab/metrics/subscribers/active_record.rb

@@ -9,6 +9,7 @@ module Gitlab
 
         IGNORABLE_SQL = %w{BEGIN COMMIT}.freeze
         DB_COUNTERS = %i{db_count db_write_count db_cached_count}.freeze
+        SQL_COMMANDS_WITH_COMMENTS_REGEX = /\A(\/\*.*\*\/\s)?((?!(.*[^\w'"](DELETE|UPDATE|INSERT INTO)[^\w'"])))(WITH.*)?(SELECT)((?!(FOR UPDATE|FOR SHARE)).)*$/i.freeze
 
         def sql(event)
           # Mark this thread as requiring a database connection. This is used
@@ -37,7 +38,7 @@ module Gitlab
         private
 
         def select_sql_command?(payload)
-          payload[:sql].match(/\A((?!(.*[^\w'"](DELETE|UPDATE|INSERT INTO)[^\w'"])))(WITH.*)?(SELECT)((?!(FOR UPDATE|FOR SHARE)).)*$/i)
+          payload[:sql].match(SQL_COMMANDS_WITH_COMMENTS_REGEX)
         end
 
         def increment_db_counters(payload)

spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -17,6 +17,15 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
     )
   end
 
+  # Emulate Marginalia pre-pending comments
+  def sql(query, comments: true)
+    if comments
+      "/*application:web,controller:badges,action:pipeline,correlation_id:01EYN39K9VMJC56Z7808N7RSRH*/ #{query}"
+    else
+      query
+    end
+  end
+
   describe '#sql' do
     shared_examples 'track query in metrics' do
       before do
@@ -101,7 +110,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         expect { subscriber.sql(event) }.to change { Thread.current[:uses_db_connection] }.from(nil).to(true)
       end
 
-      context 'with read query' do
+      shared_examples 'read queries' do
         let(:expected_counters) do
           {
             db_count: 1,
@@ -114,14 +123,14 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         it_behaves_like 'track query in RequestStore'
 
         context 'with only select' do
-          let(:payload) { { sql: 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' } }
+          let(:payload) { { sql: sql('WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones', comments: comments) } }
 
           it_behaves_like 'track query in metrics'
           it_behaves_like 'track query in RequestStore'
         end
       end
 
-      context 'write query' do
+      shared_examples 'write queries' do
         let(:expected_counters) do
           {
             db_count: 1,
@@ -131,7 +140,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         end
 
         context 'with select for update sql event' do
-          let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10 FOR UPDATE' } }
+          let(:payload) { { sql: sql('SELECT * FROM users WHERE id = 10 FOR UPDATE', comments: comments) } }
 
           it_behaves_like 'track query in metrics'
           it_behaves_like 'track query in RequestStore'
@@ -139,7 +148,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
 
         context 'with common table expression' do
           context 'with insert' do
-            let(:payload) { { sql: 'WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows' } }
+            let(:payload) { { sql: sql('WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows', comments: comments) } }
 
             it_behaves_like 'track query in metrics'
             it_behaves_like 'track query in RequestStore'
@@ -147,27 +156,41 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         end
 
         context 'with delete sql event' do
-          let(:payload) { { sql: 'DELETE FROM users where id = 10' } }
+          let(:payload) { { sql: sql('DELETE FROM users where id = 10', comments: comments) } }
 
           it_behaves_like 'track query in metrics'
           it_behaves_like 'track query in RequestStore'
         end
 
         context 'with insert sql event' do
-          let(:payload) { { sql: 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' } }
+          let(:payload) { { sql: sql('INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects', comments: comments) } }
 
           it_behaves_like 'track query in metrics'
           it_behaves_like 'track query in RequestStore'
         end
 
         context 'with update sql event' do
-          let(:payload) { { sql: 'UPDATE users SET admin = true WHERE id = 10' } }
+          let(:payload) { { sql: sql('UPDATE users SET admin = true WHERE id = 10', comments: comments) } }
 
           it_behaves_like 'track query in metrics'
           it_behaves_like 'track query in RequestStore'
         end
       end
 
+      context 'without Marginalia comments' do
+        let(:comments) { false }
+
+        it_behaves_like 'write queries'
+        it_behaves_like 'read queries'
+      end
+
+      context 'with Marginalia comments' do
+        let(:comments) { true }
+
+        it_behaves_like 'write queries'
+        it_behaves_like 'read queries'
+      end
+
       context 'with cached query' do
         let(:expected_counters) do
           {
@@ -180,7 +203,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         context 'with cached payload ' do
           let(:payload) do
             {
-              sql: 'SELECT * FROM users WHERE id = 10',
+              sql: sql('SELECT * FROM users WHERE id = 10'),
               cached: true
             }
           end
@@ -192,7 +215,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         context 'with cached payload name' do
           let(:payload) do
             {
-              sql: 'SELECT * FROM users WHERE id = 10',
+              sql: sql('SELECT * FROM users WHERE id = 10'),
               name: 'CACHE'
             }
           end
@@ -208,7 +231,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
             :event,
             name: 'sql.active_record',
             payload: {
-              sql: "SELECT attr.attname FROM pg_attribute attr INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey) WHERE cons.contype = 'p' AND cons.conrelid = '\"projects\"'::regclass",
+              sql: sql("SELECT attr.attname FROM pg_attribute attr INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey) WHERE cons.contype = 'p' AND cons.conrelid = '\"projects\"'::regclass"),
               name: 'SCHEMA',
               connection_id: 135,
               statement_name: nil,