1 files changed, 121 insertions, 11 deletions

diff --git a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
index 1624cea8bda..a78d048908d 100644
--- a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
+++ b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
@@ -6,10 +6,15 @@ describe Gitlab::Metrics::Subscribers::ActiveRecord do
   let(:env) { {} }
   let(:transaction) { Gitlab::Metrics::WebTransaction.new(env) }
   let(:subscriber)  { described_class.new }
+  let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10' } }
 
   let(:event) do
-    double(:event, duration: 2,
-                   payload:  { sql: 'SELECT * FROM users WHERE id = 10' })
+    double(
+      :event,
+      name: 'sql.active_record',
+      duration: 2,
+      payload:  payload
+    )
   end
 
   describe '#sql' do
@@ -23,6 +28,63 @@ describe Gitlab::Metrics::Subscribers::ActiveRecord do
     end
 
     describe 'with a current transaction' do
+      shared_examples 'read only query' do
+        it 'increments only db count value' do
+          allow(subscriber).to receive(:current_transaction)
+                                  .at_least(:once)
+                                  .and_return(transaction)
+
+          expect(transaction).to receive(:increment)
+                                   .with(:db_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                       .with(:db_cached_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                       .with(:db_write_count, 1)
+
+          subscriber.sql(event)
+        end
+      end
+
+      shared_examples 'write query' do
+        it 'increments db_write_count and db_count value' do
+          expect(subscriber).to receive(:current_transaction)
+                                  .at_least(:once)
+                                  .and_return(transaction)
+
+          expect(transaction).to receive(:increment)
+                                   .with(:db_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                       .with(:db_cached_count, 1)
+
+          expect(transaction).to receive(:increment)
+                                       .with(:db_write_count, 1)
+
+          subscriber.sql(event)
+        end
+      end
+
+      shared_examples 'cached query' do
+        it 'increments db_cached_count and db_count value' do
+          expect(subscriber).to receive(:current_transaction)
+                                  .at_least(:once)
+                                  .and_return(transaction)
+
+          expect(transaction).to receive(:increment)
+                                   .with(:db_count, 1)
+
+          expect(transaction).to receive(:increment)
+                                       .with(:db_cached_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                   .with(:db_write_count, 1)
+
+          subscriber.sql(event)
+        end
+      end
+
       it 'observes sql_duration metric' do
         expect(subscriber).to receive(:current_transaction)
                                 .at_least(:once)
@@ -31,18 +93,66 @@ describe Gitlab::Metrics::Subscribers::ActiveRecord do
         subscriber.sql(event)
       end
 
-      it 'increments the :sql_duration value' do
-        expect(subscriber).to receive(:current_transaction)
-          .at_least(:once)
-          .and_return(transaction)
+      it_behaves_like 'read only query'
+
+      context 'with select for update sql event' do
+        let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10 FOR UPDATE' } }
 
-        expect(transaction).to receive(:increment)
-          .with(:sql_duration, 2, false)
+        it_behaves_like 'write query'
+      end
 
-        expect(transaction).to receive(:increment)
-          .with(:sql_count, 1, false)
+      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' } }
 
-        subscriber.sql(event)
+          it_behaves_like 'write query'
+        end
+
+        context 'with only select' do
+          let(:payload) { { sql: 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' } }
+
+          it_behaves_like 'read only query'
+        end
+      end
+
+      context 'with delete sql event' do
+        let(:payload) { { sql: 'DELETE FROM users where id = 10' } }
+
+        it_behaves_like 'write query'
+      end
+
+      context 'with insert sql event' do
+        let(:payload) { { sql: 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' } }
+
+        it_behaves_like 'write query'
+      end
+
+      context 'with update sql event' do
+        let(:payload) { { sql: 'UPDATE users SET admin = true WHERE id = 10' } }
+
+        it_behaves_like 'write query'
+      end
+
+      context 'with cached payload ' do
+        let(:payload) do
+          {
+            sql: 'SELECT * FROM users WHERE id = 10',
+            cached: true
+          }
+        end
+
+        it_behaves_like 'cached query'
+      end
+
+      context 'with cached payload name' do
+        let(:payload) do
+          {
+           sql: 'SELECT * FROM users WHERE id = 10',
+           name: 'CACHE'
+          }
+        end
+
+        it_behaves_like 'cached query'
       end
 
       context 'events are internal to Rails or irrelevant' do