File: pg_auto_parameterize_duplicate_query_detection.rb

package info (click to toggle)
ruby-sequel 5.102.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 11,328 kB
  • sloc: ruby: 124,743; makefile: 3
file content (191 lines) | stat: -rw-r--r-- 7,313 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
# frozen-string-literal: true
#
# The pg_auto_parameterize_duplicate_query_detection extension builds on the
# pg_auto_parameterize extension, adding support for detecting duplicate
# queries inside a block that occur at the same location. This is designed
# mostly to catch duplicate query issues (e.g. N+1 queries) during testing.
#
# To detect duplicate queries inside a block of code, wrap the code with
# +detect_duplicate_queries+:
#
#   DB.detect_duplicate_queries{your_code}
#
# With this approach, if the test runs code where the same query is executed
# more than once with the same call stack, a
# Sequel::Postgres::AutoParameterizeDuplicateQueryDetection::DuplicateQueries
# exception will be raised.
#
# You can pass the +:warn+ option to +detect_duplicate_queries+ to warn
# instead of raising. Note that if the block passed to +detect_duplicate_queries+
# raises, this extension will warn, and raise the original exception.
#
# For more control, you can pass the +:handler+ option to
# +detect_duplicate_queries+. If the +:handler+ option is provided, this
# extension will call the +:handler+ option with the hash of duplicate
# query information, and will not raise or warn. This can be useful in
# production environments, to record duplicate queries for later analysis.
#
# For accuracy, the entire call stack is always used as part of the hash key
# to determine whether a query is duplicate. However, you can filter the
# displayed backtrace by using the +:backtrace_filter+ option.
#
# +detect_duplicate_queries+ is concurrency aware, it uses the same approach
# that Sequel's default connection pools use. So if you are running multiple
# threads, +detect_duplicate_queries+ will only report duplicate queries for
# the current thread (or fiber if the fiber_concurrency extension is used).
#
# When testing applications, it's probably best to use this to wrap the
# application being tested. For example, testing with rack-test, if your app
# is +App+, you would want to wrap it:
#
#   include Rack::Test::Methods
#
#   WrappedApp = lambda do |env|
#     DB.detect_duplicate_queries{App.call(env)}
#   end
#
#   def app
#     WrappedApp
#   end
#
# It is possible to use this to wrap each separate spec using an around hook,
# but that can result in false positives when using libraries that have
# implicit retrying (such as Capybara), as you can have the same call stack
# for multiple requests.
#
# Related module: Sequel::Postgres::AutoParameterizeDuplicateQueryDetection

module Sequel
  module Postgres
    # Enable detecting duplicate queries inside a block
    module AutoParameterizeDuplicateQueryDetection
      def self.extended(db)
        db.instance_exec do
          @duplicate_query_detection_contexts = {}
          @duplicate_query_detection_mutex = Mutex.new
        end
      end

      # Exception class raised when duplicate queries are detected.
      class DuplicateQueries < Sequel::Error
        # A hash of queries that were duplicate. Keys are arrays
        # with 2 entries, the first being the query SQL, and the
        # second being the related call stack (backtrace).
        # The values are the number of query executions.
        attr_reader :queries

        def initialize(message, queries)
          @queries = queries
          super(message)
        end
      end

      # Record each query executed so duplicates can be detected,
      # if queries are being recorded.
      def execute(sql, opts=OPTS, &block)
        record, queries = duplicate_query_recorder_state

        if record
          queries[[sql.is_a?(Symbol) ? sql : sql.to_s, caller].freeze] += 1
        end

        super
      end

      # Ignore (do not record) queries inside given block. This can
      # be useful in situations where you want to run your entire test suite
      # with duplicate query detection, but you have duplicate queries in
      # some parts of your application where it is not trivial to use a
      # different approach. You can mark those specific sections with
      # +ignore_duplicate_queries+, and still get duplicate query detection
      # for the rest of the application.
      def ignore_duplicate_queries(&block)
        if state = duplicate_query_recorder_state
          change_duplicate_query_recorder_state(state, false, &block)
        else
          # If we are not inside a detect_duplicate_queries block, there is
          # no need to do anything, since we are not recording queries.
          yield
        end
      end

      # Run the duplicate query detector during the block.
      # Options:
      #
      # :backtrace_filter :: Regexp used to filter the displayed backtrace.
      # :handler :: If present, called with hash of duplicate query information,
      #             instead of raising or warning.
      # :warn :: Always warn instead of raising for duplicate queries.
      #
      # Note that if you nest calls to this method, only the top
      # level call will respect the passed options.
      def detect_duplicate_queries(opts=OPTS, &block)
        current = Sequel.current
        if state = duplicate_query_recorder_state(current)
          return change_duplicate_query_recorder_state(state, true, &block)
        end

        @duplicate_query_detection_mutex.synchronize do
          @duplicate_query_detection_contexts[current] = [true, Hash.new(0)]
        end

        begin
          yield
        rescue Exception => e
          raise
        ensure
          _, queries = @duplicate_query_detection_mutex.synchronize do
            @duplicate_query_detection_contexts.delete(current)
          end
          queries.delete_if{|_,v| v < 2}

          unless queries.empty?
            if handler = opts[:handler]
              handler.call(queries)
            else
              backtrace_filter = opts[:backtrace_filter]
              backtrace_filter_note = backtrace_filter ? " (filtered)" : ""
              query_info = queries.map do |k,v|
                backtrace = k[1]
                backtrace = backtrace.grep(backtrace_filter) if backtrace_filter
                "times:#{v}\nsql:#{k[0]}\nbacktrace#{backtrace_filter_note}:\n#{backtrace.join("\n")}\n"
              end
              message = "duplicate queries detected:\n\n#{query_info.join("\n")}"

              if e || opts[:warn]
                warn(message)
              else
                raise DuplicateQueries.new(message, queries)
              end
            end
          end
        end
      end

      private

      # Get the query record state for the given context.
      def duplicate_query_recorder_state(current=Sequel.current)
        @duplicate_query_detection_mutex.synchronize{@duplicate_query_detection_contexts[current]}
      end

      # Temporarily change whether to record queries for the block, resetting the
      # previous setting after the block exits.
      def change_duplicate_query_recorder_state(state, setting)
        prev = state[0]
        state[0] = setting
        
        begin
          yield
        ensure
          state[0] = prev
        end
      end
    end
  end

  Database.register_extension(:pg_auto_parameterize_duplicate_query_detection) do |db|
    db.extension(:pg_auto_parameterize)
    db.extend(Postgres::AutoParameterizeDuplicateQueryDetection)
  end
end