|
| 1 | +# frozen-string-literal: true |
| 2 | +# |
| 3 | +# The pg_auto_parameterize_duplicate_query_detection extension builds on the |
| 4 | +# pg_auto_parameterize extension, adding support for detecting duplicate |
| 5 | +# queries inside a block that occur at the same location. This is designed |
| 6 | +# mostly to catch duplicate query issues (e.g. N+1 queries) during testing. |
| 7 | +# |
| 8 | +# To detect duplicate queries inside a block of code, wrap the code with |
| 9 | +# +detect_duplicate_queries+: |
| 10 | +# |
| 11 | +# DB.detect_duplicate_queries{your_code} |
| 12 | +# |
| 13 | +# With this approach, if the test runs code where the same query is executed |
| 14 | +# more than once with the same call stack, a |
| 15 | +# Sequel::Postgres::AutoParameterizeDuplicateQueryDetection::DuplicateQueries |
| 16 | +# exception will be raised. |
| 17 | +# |
| 18 | +# You can pass the +:warn+ option to +detect_duplicate_queries+ to warn |
| 19 | +# instead of raising. Note that if the block passed to +detect_duplicate_queries+ |
| 20 | +# raises, this extension will warn, and raise the original exception. |
| 21 | +# |
| 22 | +# For more control, you can pass the +:handler+ option to |
| 23 | +# +detect_duplicate_queries+. If the +:handler+ option is provided, this |
| 24 | +# extension will call the +:handler+ option with the hash of duplicate |
| 25 | +# query information, and will not raise or warn. This can be useful in |
| 26 | +# production environments, to record duplicate queries for later analysis. |
| 27 | +# |
| 28 | +# For accuracy, the entire call stack is always used as part of the hash key |
| 29 | +# to determine whether a query is duplicate. However, you can filter the |
| 30 | +# displayed backtrace by using the +:backtrace_filter+ option. |
| 31 | +# |
| 32 | +# +detect_duplicate_queries+ is concurrency aware, it uses the same approach |
| 33 | +# that Sequel's default connection pools use. So if you are running multiple |
| 34 | +# threads, +detect_duplicate_queries+ will only report duplicate queries for |
| 35 | +# the current thread (or fiber if the fiber_concurrency extension is used). |
| 36 | +# |
| 37 | +# When testing applications, it's probably best to use this to wrap the |
| 38 | +# application being tested. For example, testing with rack-test, if your app |
| 39 | +# is +App+, you would want to wrap it: |
| 40 | +# |
| 41 | +# include Rack::Test::Methods |
| 42 | +# |
| 43 | +# WrappedApp = lambda do |env| |
| 44 | +# DB.detect_duplicate_queries{App.call(env)} |
| 45 | +# end |
| 46 | +# |
| 47 | +# def app |
| 48 | +# WrappedApp |
| 49 | +# end |
| 50 | +# |
| 51 | +# It is possible to use this to wrap each separate spec using an around hook, |
| 52 | +# but that can result in false positives when using libraries that have |
| 53 | +# implicit retrying (such as Capybara), as you can have the same call stack |
| 54 | +# for multiple requests. |
| 55 | +# |
| 56 | +# Related module: Sequel::Postgres::AutoParameterizeDuplicateQueryDetection |
| 57 | + |
| 58 | +module Sequel |
| 59 | + module Postgres |
| 60 | + # Enable detecting duplicate queries inside a block |
| 61 | + module AutoParameterizeDuplicateQueryDetection |
| 62 | + def self.extended(db) |
| 63 | + db.instance_exec do |
| 64 | + @duplicate_query_detection_contexts = {} |
| 65 | + @duplicate_query_detection_mutex = Mutex.new |
| 66 | + end |
| 67 | + end |
| 68 | + |
| 69 | + # Exception class raised when duplicate queries are detected. |
| 70 | + class DuplicateQueries < Sequel::Error |
| 71 | + # A hash of queries that were duplicate. Keys are arrays |
| 72 | + # with 2 entries, the first being the query SQL, and the |
| 73 | + # second being the related caller line. |
| 74 | + # The values are the number of query executions. |
| 75 | + attr_reader :queries |
| 76 | + |
| 77 | + def initialize(message, queries) |
| 78 | + @queries = queries |
| 79 | + super(message) |
| 80 | + end |
| 81 | + end |
| 82 | + |
| 83 | + # Record each query executed so duplicates can be detected, |
| 84 | + # if queries are being recorded. |
| 85 | + def execute(sql, opts=OPTS, &block) |
| 86 | + record, queries = duplicate_query_recorder_state |
| 87 | + |
| 88 | + if record |
| 89 | + queries[[sql.is_a?(Symbol) ? sql : sql.to_s, caller].freeze] += 1 |
| 90 | + end |
| 91 | + |
| 92 | + super |
| 93 | + end |
| 94 | + |
| 95 | + # Ignore (do not record) queries inside given block. This can |
| 96 | + # be useful in situations where you want to run your entire test suite |
| 97 | + # with duplicate query detection, but you have duplicate queries in |
| 98 | + # some parts of your application where it is not trivial to use a |
| 99 | + # different approach. You can mark those specific sections with |
| 100 | + # +ignore_duplicate_queries+, and still get duplicate query detection |
| 101 | + # for the rest of the application. |
| 102 | + def ignore_duplicate_queries(&block) |
| 103 | + if state = duplicate_query_recorder_state |
| 104 | + change_duplicate_query_recorder_state(state, false, &block) |
| 105 | + else |
| 106 | + # If we are not inside a detect_duplicate_queries block, there is |
| 107 | + # no need to do anything, since we are not recording queries. |
| 108 | + yield |
| 109 | + end |
| 110 | + end |
| 111 | + |
| 112 | + # Run the duplicate query detector during the block. |
| 113 | + # Options: |
| 114 | + # |
| 115 | + # :backtrace_filter :: Regexp used to filter the displayed backtrace. |
| 116 | + # :handler :: If present, called with hash of duplicate query information, |
| 117 | + # instead of raising or warning. |
| 118 | + # :warn :: Always warn instead of raising for duplicate queries. |
| 119 | + # |
| 120 | + # Note that if you nest calls to this method, only the top |
| 121 | + # level call will respect the passed options. |
| 122 | + def detect_duplicate_queries(opts=OPTS, &block) |
| 123 | + current = Sequel.current |
| 124 | + if state = duplicate_query_recorder_state(current) |
| 125 | + return change_duplicate_query_recorder_state(state, true, &block) |
| 126 | + end |
| 127 | + |
| 128 | + @duplicate_query_detection_mutex.synchronize do |
| 129 | + @duplicate_query_detection_contexts[current] = [true, Hash.new(0)] |
| 130 | + end |
| 131 | + |
| 132 | + begin |
| 133 | + yield |
| 134 | + rescue Exception => e |
| 135 | + raise |
| 136 | + ensure |
| 137 | + _, queries = @duplicate_query_detection_mutex.synchronize do |
| 138 | + @duplicate_query_detection_contexts.delete(current) |
| 139 | + end |
| 140 | + queries.delete_if{|_,v| v < 2} |
| 141 | + |
| 142 | + unless queries.empty? |
| 143 | + if handler = opts[:handler] |
| 144 | + handler.call(queries) |
| 145 | + else |
| 146 | + backtrace_filter = opts[:backtrace_filter] |
| 147 | + backtrace_filter_note = backtrace_filter ? " (filtered)" : "" |
| 148 | + query_info = queries.map do |k,v| |
| 149 | + backtrace = k[1] |
| 150 | + backtrace = backtrace.grep(backtrace_filter) if backtrace_filter |
| 151 | + "times:#{v}\nsql:#{k[0]}\nbacktrace#{backtrace_filter_note}:\n#{backtrace.join("\n")}\n" |
| 152 | + end |
| 153 | + message = "duplicate queries detected:\n\n#{query_info.join("\n")}" |
| 154 | + |
| 155 | + if e || opts[:warn] |
| 156 | + warn(message) |
| 157 | + else |
| 158 | + raise DuplicateQueries.new(message, queries) |
| 159 | + end |
| 160 | + end |
| 161 | + end |
| 162 | + end |
| 163 | + end |
| 164 | + |
| 165 | + private |
| 166 | + |
| 167 | + # Get the query record state for the given context. |
| 168 | + def duplicate_query_recorder_state(current=Sequel.current) |
| 169 | + @duplicate_query_detection_mutex.synchronize{@duplicate_query_detection_contexts[current]} |
| 170 | + end |
| 171 | + |
| 172 | + # Temporarily change whether to record queries for the block, resetting the |
| 173 | + # previous setting after the block exits. |
| 174 | + def change_duplicate_query_recorder_state(state, setting) |
| 175 | + prev = state[0] |
| 176 | + state[0] = setting |
| 177 | + |
| 178 | + begin |
| 179 | + yield |
| 180 | + ensure |
| 181 | + state[0] = prev |
| 182 | + end |
| 183 | + end |
| 184 | + end |
| 185 | + end |
| 186 | + |
| 187 | + Database.register_extension(:pg_auto_parameterize_duplicate_query_detection) do |db| |
| 188 | + db.extension(:pg_auto_parameterize) |
| 189 | + db.extend(Postgres::AutoParameterizeDuplicateQueryDetection) |
| 190 | + end |
| 191 | +end |
0 commit comments