Add JSONB scan detector

Adds the necessary script to detect a WHERE or JOIN on a JSONB column.
Please see linked issue for more info.

scripts/database/query_analyzers.rb

@@ -2,7 +2,7 @@
 
 require 'yaml'
 
-class Database
+module Database
   class QueryAnalyzers
     attr_reader :analyzers
 

scripts/database/query_analyzers.yml

@@ -9,3 +9,5 @@ MultiplePartitionScanDetector:
     # These fingerprints can be found in the auto_explain pipeline artifacts.
     # Example:
     # - c2cfe803a497101b
+JSONBScanDetector:
+  todos:

scripts/database/query_analyzers/base.rb

@@ -3,7 +3,7 @@
 require 'json'
 require 'zlib'
 
-class Database
+module Database
   class QueryAnalyzers
     class Base
       attr_accessor :output

scripts/database/query_analyzers/jsonb_scan_detector.rb

+# frozen_string_literal: true
+
+require_relative 'base'
+
+module Database
+  class QueryAnalyzers
+    class JSONBScanDetector < Database::QueryAnalyzers::Base
+      JSONB_MATCH_OPERATOR_EXPRESSION = /<@|@>/
+
+      def initialize(*args)
+        super
+        output[:bad_queries] = []
+      end
+
+      def analyze(query)
+        super
+        return if config['todos']&.include?(query['fingerprint'])
+
+        output[:bad_queries] << query if has_operator_in_where?(query['query'])
+      end
+
+      def save!
+        return if output[:bad_queries].empty?
+
+        Zlib::GzipWriter.open(output_path("jsonb_column_scans.ndjson")) do |file|
+          output[:bad_queries].each do |query|
+            file.puts(JSON.generate(query))
+          end
+        end
+      end
+
+      private
+
+      def has_operator_in_where?(query)
+        return false unless query.match?(JSONB_MATCH_OPERATOR_EXPRESSION)
+
+        clauses = query.split(/\sWHERE\s|\sJOIN\s/)
+        return false if clauses.length < 2
+
+        clauses[1..].each do |c|
+          return true if c.include?('::jsonb')
+        end
+      end
+    end
+  end
+end

scripts/database/query_analyzers/multiple_partition_scan_detector.rb

@@ -2,7 +2,7 @@
 
 require_relative 'base'
 
-class Database
+module Database
   class QueryAnalyzers
     class MultiplePartitionScanDetector < Database::QueryAnalyzers::Base
       def analyze(query)

spec/scripts/database/query_analyzers/jsonb_scan_detector_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+
+require_relative '../../../../scripts/database/query_analyzers/jsonb_scan_detector'
+
+RSpec.describe Database::QueryAnalyzers::JSONBScanDetector, feature_category: :database do
+  let(:invalid_query_string) do
+    <<~SQL
+      SELECT
+          *
+      FROM
+          member_roles
+      WHERE
+          member_roles.permissions @> ('{"admin_merge_request":true}')::jsonb
+    SQL
+  end
+
+  let(:invalid_query) { { 'query' => invalid_query_string, 'fingerprint' => '0000000000000001' } }
+
+  let(:valid_query_string) { "SELECT * FROM users WHERE name = 'bob'" }
+
+  let(:valid_query) { { 'query' => valid_query_string, 'fingerprint' => '0000000000000002' } }
+
+  let(:config) { {} }
+
+  subject(:analyzer) { described_class.new(config) }
+
+  it "initalizes" do
+    expect { analyzer }.not_to raise_error
+  end
+
+  context 'when no TODOs are defined' do
+    it 'finds the invalid query' do
+      [valid_query, invalid_query].each { |q| analyzer.analyze q }
+      expect(analyzer.output[:bad_queries].length).to eq 1
+    end
+  end
+
+  context 'when a TODO is defined' do
+    let(:config) do
+      {
+        "todos" => [
+          invalid_query['fingerprint']
+        ]
+      }
+    end
+
+    it 'does not find the invalid query' do
+      [valid_query, invalid_query].each { |q| analyzer.analyze q }
+      expect(analyzer.output[:bad_queries].length).to eq 0
+    end
+  end
+end