Implement prepend patch for postgres (#625)

We've had issues running rack-mini-profiler alongside dd-trace (see:
DataDog/dd-trace-rb#2348).

The suggested solution is to add support for prepend style patching for
pg, similar to the one that exists for mysql. This does that.

lib/patches/db/pg.rb

@@ -1,122 +1,7 @@
 # frozen_string_literal: true
 
-# PG patches, keep in mind exec and async_exec have a exec{|r| } semantics that is yet to be implemented
-class PG::Result
-  alias_method :each_without_profiling, :each
-  alias_method :values_without_profiling, :values
-
-  def values(*args, &blk)
-    return values_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
-    mp_report_sql do
-      values_without_profiling(*args , &blk)
-    end
-  end
-
-  def each(*args, &blk)
-    return each_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
-    mp_report_sql do
-      each_without_profiling(*args, &blk)
-    end
-  end
-
-  def mp_report_sql(&block)
-    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-    result       = yield
-    elapsed_time = SqlPatches.elapsed_time(start)
-    @miniprofiler_sql_id.report_reader_duration(elapsed_time)
-    result
-  end
-end
-
-class PG::Connection
-  alias_method :exec_without_profiling, :exec
-  alias_method :async_exec_without_profiling, :async_exec
-  alias_method :exec_prepared_without_profiling, :exec_prepared
-  alias_method :send_query_prepared_without_profiling, :send_query_prepared
-  alias_method :prepare_without_profiling, :prepare
-
-  if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
-    alias_method :exec_params_without_profiling, :exec_params
-  end
-
-  def prepare(*args, &blk)
-    # we have no choice but to do this here,
-    # if we do the check for profiling first, our cache may miss critical stuff
-
-    @prepare_map ||= {}
-    @prepare_map[args[0]] = args[1]
-    # dont leak more than 10k ever
-    @prepare_map = {} if @prepare_map.length > 1000
-
-    return prepare_without_profiling(*args, &blk) unless SqlPatches.should_measure?
-    prepare_without_profiling(*args, &blk)
-  end
-
-  def exec(*args, &blk)
-    return exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
-
-    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-    result       = exec_without_profiling(*args, &blk)
-    elapsed_time = SqlPatches.elapsed_time(start)
-    record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
-    result.instance_variable_set("@miniprofiler_sql_id", record) if result
-
-    result
-  end
-
-  if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
-    def exec_params(*args, &blk)
-      return exec_params_without_profiling(*args, &blk) unless SqlPatches.should_measure?
-
-      start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-      result       = exec_params_without_profiling(*args, &blk)
-      elapsed_time = SqlPatches.elapsed_time(start)
-      record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
-      result.instance_variable_set("@miniprofiler_sql_id", record) if result
-
-      result
-    end
-  end
-
-  def exec_prepared(*args, &blk)
-    return exec_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
-
-    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-    result       = exec_prepared_without_profiling(*args, &blk)
-    elapsed_time = SqlPatches.elapsed_time(start)
-    mapped       = args[0]
-    mapped       = @prepare_map[mapped] || args[0] if @prepare_map
-    record       = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
-    result.instance_variable_set("@miniprofiler_sql_id", record) if result
-
-    result
-  end
-
-  def send_query_prepared(*args, &blk)
-    return send_query_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
-
-    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-    result       = send_query_prepared_without_profiling(*args, &blk)
-    elapsed_time = SqlPatches.elapsed_time(start)
-    mapped       = args[0]
-    mapped       = @prepare_map[mapped] || args[0] if @prepare_map
-    record       = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
-    result.instance_variable_set("@miniprofiler_sql_id", record) if result
-
-    result
-  end
-
-  def async_exec(*args, &blk)
-    return async_exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
-
-    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-    result       = exec_without_profiling(*args, &blk)
-    elapsed_time = SqlPatches.elapsed_time(start)
-    record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
-    result.instance_variable_set("@miniprofiler_sql_id", record) if result
-
-    result
-  end
-
-  alias_method :query, :exec
+if defined?(Rack::MINI_PROFILER_PREPEND_PG_PATCH)
+  require "patches/db/pg/prepend"
+else
+  require "patches/db/pg/alias_method"
 end

lib/patches/db/pg/alias_method.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+class PG::Result
+  alias_method :each_without_profiling, :each
+  alias_method :values_without_profiling, :values
+
+  def values(*args, &blk)
+    return values_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
+    mp_report_sql do
+      values_without_profiling(*args , &blk)
+    end
+  end
+
+  def each(*args, &blk)
+    return each_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
+    mp_report_sql do
+      each_without_profiling(*args, &blk)
+    end
+  end
+
+  def mp_report_sql(&block)
+    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    result       = yield
+    elapsed_time = SqlPatches.elapsed_time(start)
+    @miniprofiler_sql_id.report_reader_duration(elapsed_time)
+    result
+  end
+end
+
+class PG::Connection
+  alias_method :exec_without_profiling, :exec
+  alias_method :async_exec_without_profiling, :async_exec
+  alias_method :exec_prepared_without_profiling, :exec_prepared
+  alias_method :send_query_prepared_without_profiling, :send_query_prepared
+  alias_method :prepare_without_profiling, :prepare
+
+  if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
+    alias_method :exec_params_without_profiling, :exec_params
+  end
+
+  def prepare(*args, &blk)
+    # we have no choice but to do this here,
+    # if we do the check for profiling first, our cache may miss critical stuff
+
+    @prepare_map ||= {}
+    @prepare_map[args[0]] = args[1]
+    # dont leak more than 10k ever
+    @prepare_map = {} if @prepare_map.length > 1000
+
+    return prepare_without_profiling(*args, &blk) unless SqlPatches.should_measure?
+    prepare_without_profiling(*args, &blk)
+  end
+
+  def exec(*args, &blk)
+    return exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
+
+    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    result       = exec_without_profiling(*args, &blk)
+    elapsed_time = SqlPatches.elapsed_time(start)
+    record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
+    result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+    result
+  end
+
+  if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
+    def exec_params(*args, &blk)
+      return exec_params_without_profiling(*args, &blk) unless SqlPatches.should_measure?
+
+      start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result       = exec_params_without_profiling(*args, &blk)
+      elapsed_time = SqlPatches.elapsed_time(start)
+      record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
+      result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+      result
+    end
+  end
+
+  def exec_prepared(*args, &blk)
+    return exec_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
+
+    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    result       = exec_prepared_without_profiling(*args, &blk)
+    elapsed_time = SqlPatches.elapsed_time(start)
+    mapped       = args[0]
+    mapped       = @prepare_map[mapped] || args[0] if @prepare_map
+    record       = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
+    result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+    result
+  end
+
+  def send_query_prepared(*args, &blk)
+    return send_query_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
+
+    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    result       = send_query_prepared_without_profiling(*args, &blk)
+    elapsed_time = SqlPatches.elapsed_time(start)
+    mapped       = args[0]
+    mapped       = @prepare_map[mapped] || args[0] if @prepare_map
+    record       = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
+    result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+    result
+  end
+
+  def async_exec(*args, &blk)
+    return async_exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
+
+    start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    result       = exec_without_profiling(*args, &blk)
+    elapsed_time = SqlPatches.elapsed_time(start)
+    record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
+    result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+    result
+  end
+
+  alias_method :query, :exec
+end

lib/patches/db/pg/prepend.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+class PG::Result
+  module MiniProfiler
+    def values(*args, &blk)
+      return super unless defined?(@miniprofiler_sql_id)
+      mp_report_sql do
+        super
+      end
+    end
+
+    def each(*args, &blk)
+      return super unless defined?(@miniprofiler_sql_id)
+      mp_report_sql do
+        super
+      end
+    end
+
+    def mp_report_sql(&block)
+      start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result       = yield
+      elapsed_time = SqlPatches.elapsed_time(start)
+      @miniprofiler_sql_id.report_reader_duration(elapsed_time)
+      result
+    end
+  end
+
+  prepend MiniProfiler
+end
+
+class PG::Connection
+  module MiniProfiler
+    def prepare(*args, &blk)
+      # we have no choice but to do this here,
+      # if we do the check for profiling first, our cache may miss critical stuff
+
+      @prepare_map ||= {}
+      @prepare_map[args[0]] = args[1]
+      # dont leak more than 10k ever
+      @prepare_map = {} if @prepare_map.length > 1000
+
+      return super unless SqlPatches.should_measure?
+      super
+    end
+
+    def exec(*args, &blk)
+      return super unless SqlPatches.should_measure?
+
+      start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result       = super
+      elapsed_time = SqlPatches.elapsed_time(start)
+      record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
+      result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+      result
+    end
+
+    if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
+      def exec_params(*args, &blk)
+        return super unless SqlPatches.should_measure?
+
+        start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        result       = super
+        elapsed_time = SqlPatches.elapsed_time(start)
+        record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
+        result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+        result
+      end
+    end
+
+    def exec_prepared(*args, &blk)
+      return super unless SqlPatches.should_measure?
+
+      start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result       = super
+      elapsed_time = SqlPatches.elapsed_time(start)
+      mapped       = args[0]
+      mapped       = @prepare_map[mapped] || args[0] if @prepare_map
+      record       = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
+      result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+      result
+    end
+
+    def send_query_prepared(*args, &blk)
+      return super unless SqlPatches.should_measure?
+
+      start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result       = super
+      elapsed_time = SqlPatches.elapsed_time(start)
+      mapped       = args[0]
+      mapped       = @prepare_map[mapped] || args[0] if @prepare_map
+      record       = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
+      result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+      result
+    end
+
+    def async_exec(*args, &blk)
+      return super unless SqlPatches.should_measure?
+
+      start        = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result       = super
+      elapsed_time = SqlPatches.elapsed_time(start)
+      record       = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
+      result.instance_variable_set("@miniprofiler_sql_id", record) if result
+
+      result
+    end
+  end
+
+  prepend MiniProfiler
+  alias_method :query, :exec
+end

lib/prepend_pg_patch.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Rack
+  MINI_PROFILER_PREPEND_PG_PATCH = true
+end