Skip to content

Commit d5b895f

Browse files
tomchipchasetomchipchase
authored
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.
1 parent 5e42a57 commit d5b895f

File tree

4 files changed

+245
-119
lines changed

4 files changed

+245
-119
lines changed

‎lib/patches/db/pg.rb

Lines changed: 4 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -1,122 +1,7 @@
11
# frozen_string_literal: true
22

3-
# PG patches, keep in mind exec and async_exec have a exec{|r| } semantics that is yet to be implemented
4-
class PG::Result
5-
alias_method :each_without_profiling, :each
6-
alias_method :values_without_profiling, :values
7-
8-
def values(*args, &blk)
9-
return values_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
10-
mp_report_sql do
11-
values_without_profiling(*args , &blk)
12-
end
13-
end
14-
15-
def each(*args, &blk)
16-
return each_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
17-
mp_report_sql do
18-
each_without_profiling(*args, &blk)
19-
end
20-
end
21-
22-
def mp_report_sql(&block)
23-
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
24-
result = yield
25-
elapsed_time = SqlPatches.elapsed_time(start)
26-
@miniprofiler_sql_id.report_reader_duration(elapsed_time)
27-
result
28-
end
29-
end
30-
31-
class PG::Connection
32-
alias_method :exec_without_profiling, :exec
33-
alias_method :async_exec_without_profiling, :async_exec
34-
alias_method :exec_prepared_without_profiling, :exec_prepared
35-
alias_method :send_query_prepared_without_profiling, :send_query_prepared
36-
alias_method :prepare_without_profiling, :prepare
37-
38-
if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
39-
alias_method :exec_params_without_profiling, :exec_params
40-
end
41-
42-
def prepare(*args, &blk)
43-
# we have no choice but to do this here,
44-
# if we do the check for profiling first, our cache may miss critical stuff
45-
46-
@prepare_map ||= {}
47-
@prepare_map[args[0]] = args[1]
48-
# dont leak more than 10k ever
49-
@prepare_map = {} if @prepare_map.length > 1000
50-
51-
return prepare_without_profiling(*args, &blk) unless SqlPatches.should_measure?
52-
prepare_without_profiling(*args, &blk)
53-
end
54-
55-
def exec(*args, &blk)
56-
return exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
57-
58-
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
59-
result = exec_without_profiling(*args, &blk)
60-
elapsed_time = SqlPatches.elapsed_time(start)
61-
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
62-
result.instance_variable_set("@miniprofiler_sql_id", record) if result
63-
64-
result
65-
end
66-
67-
if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
68-
def exec_params(*args, &blk)
69-
return exec_params_without_profiling(*args, &blk) unless SqlPatches.should_measure?
70-
71-
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
72-
result = exec_params_without_profiling(*args, &blk)
73-
elapsed_time = SqlPatches.elapsed_time(start)
74-
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
75-
result.instance_variable_set("@miniprofiler_sql_id", record) if result
76-
77-
result
78-
end
79-
end
80-
81-
def exec_prepared(*args, &blk)
82-
return exec_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
83-
84-
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
85-
result = exec_prepared_without_profiling(*args, &blk)
86-
elapsed_time = SqlPatches.elapsed_time(start)
87-
mapped = args[0]
88-
mapped = @prepare_map[mapped] || args[0] if @prepare_map
89-
record = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
90-
result.instance_variable_set("@miniprofiler_sql_id", record) if result
91-
92-
result
93-
end
94-
95-
def send_query_prepared(*args, &blk)
96-
return send_query_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
97-
98-
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
99-
result = send_query_prepared_without_profiling(*args, &blk)
100-
elapsed_time = SqlPatches.elapsed_time(start)
101-
mapped = args[0]
102-
mapped = @prepare_map[mapped] || args[0] if @prepare_map
103-
record = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
104-
result.instance_variable_set("@miniprofiler_sql_id", record) if result
105-
106-
result
107-
end
108-
109-
def async_exec(*args, &blk)
110-
return async_exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
111-
112-
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
113-
result = exec_without_profiling(*args, &blk)
114-
elapsed_time = SqlPatches.elapsed_time(start)
115-
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
116-
result.instance_variable_set("@miniprofiler_sql_id", record) if result
117-
118-
result
119-
end
120-
121-
alias_method :query, :exec
3+
if defined?(Rack::MINI_PROFILER_PREPEND_PG_PATCH)
4+
require "patches/db/pg/prepend"
5+
else
6+
require "patches/db/pg/alias_method"
1227
end

‎lib/patches/db/pg/alias_method.rb

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
# frozen_string_literal: true
2+
3+
class PG::Result
4+
alias_method :each_without_profiling, :each
5+
alias_method :values_without_profiling, :values
6+
7+
def values(*args, &blk)
8+
return values_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
9+
mp_report_sql do
10+
values_without_profiling(*args , &blk)
11+
end
12+
end
13+
14+
def each(*args, &blk)
15+
return each_without_profiling(*args, &blk) unless defined?(@miniprofiler_sql_id)
16+
mp_report_sql do
17+
each_without_profiling(*args, &blk)
18+
end
19+
end
20+
21+
def mp_report_sql(&block)
22+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
23+
result = yield
24+
elapsed_time = SqlPatches.elapsed_time(start)
25+
@miniprofiler_sql_id.report_reader_duration(elapsed_time)
26+
result
27+
end
28+
end
29+
30+
class PG::Connection
31+
alias_method :exec_without_profiling, :exec
32+
alias_method :async_exec_without_profiling, :async_exec
33+
alias_method :exec_prepared_without_profiling, :exec_prepared
34+
alias_method :send_query_prepared_without_profiling, :send_query_prepared
35+
alias_method :prepare_without_profiling, :prepare
36+
37+
if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
38+
alias_method :exec_params_without_profiling, :exec_params
39+
end
40+
41+
def prepare(*args, &blk)
42+
# we have no choice but to do this here,
43+
# if we do the check for profiling first, our cache may miss critical stuff
44+
45+
@prepare_map ||= {}
46+
@prepare_map[args[0]] = args[1]
47+
# dont leak more than 10k ever
48+
@prepare_map = {} if @prepare_map.length > 1000
49+
50+
return prepare_without_profiling(*args, &blk) unless SqlPatches.should_measure?
51+
prepare_without_profiling(*args, &blk)
52+
end
53+
54+
def exec(*args, &blk)
55+
return exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
56+
57+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
58+
result = exec_without_profiling(*args, &blk)
59+
elapsed_time = SqlPatches.elapsed_time(start)
60+
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
61+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
62+
63+
result
64+
end
65+
66+
if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
67+
def exec_params(*args, &blk)
68+
return exec_params_without_profiling(*args, &blk) unless SqlPatches.should_measure?
69+
70+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
71+
result = exec_params_without_profiling(*args, &blk)
72+
elapsed_time = SqlPatches.elapsed_time(start)
73+
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
74+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
75+
76+
result
77+
end
78+
end
79+
80+
def exec_prepared(*args, &blk)
81+
return exec_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
82+
83+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
84+
result = exec_prepared_without_profiling(*args, &blk)
85+
elapsed_time = SqlPatches.elapsed_time(start)
86+
mapped = args[0]
87+
mapped = @prepare_map[mapped] || args[0] if @prepare_map
88+
record = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
89+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
90+
91+
result
92+
end
93+
94+
def send_query_prepared(*args, &blk)
95+
return send_query_prepared_without_profiling(*args, &blk) unless SqlPatches.should_measure?
96+
97+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
98+
result = send_query_prepared_without_profiling(*args, &blk)
99+
elapsed_time = SqlPatches.elapsed_time(start)
100+
mapped = args[0]
101+
mapped = @prepare_map[mapped] || args[0] if @prepare_map
102+
record = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
103+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
104+
105+
result
106+
end
107+
108+
def async_exec(*args, &blk)
109+
return async_exec_without_profiling(*args, &blk) unless SqlPatches.should_measure?
110+
111+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
112+
result = exec_without_profiling(*args, &blk)
113+
elapsed_time = SqlPatches.elapsed_time(start)
114+
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
115+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
116+
117+
result
118+
end
119+
120+
alias_method :query, :exec
121+
end

‎lib/patches/db/pg/prepend.rb

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# frozen_string_literal: true
2+
3+
class PG::Result
4+
module MiniProfiler
5+
def values(*args, &blk)
6+
return super unless defined?(@miniprofiler_sql_id)
7+
mp_report_sql do
8+
super
9+
end
10+
end
11+
12+
def each(*args, &blk)
13+
return super unless defined?(@miniprofiler_sql_id)
14+
mp_report_sql do
15+
super
16+
end
17+
end
18+
19+
def mp_report_sql(&block)
20+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
21+
result = yield
22+
elapsed_time = SqlPatches.elapsed_time(start)
23+
@miniprofiler_sql_id.report_reader_duration(elapsed_time)
24+
result
25+
end
26+
end
27+
28+
prepend MiniProfiler
29+
end
30+
31+
class PG::Connection
32+
module MiniProfiler
33+
def prepare(*args, &blk)
34+
# we have no choice but to do this here,
35+
# if we do the check for profiling first, our cache may miss critical stuff
36+
37+
@prepare_map ||= {}
38+
@prepare_map[args[0]] = args[1]
39+
# dont leak more than 10k ever
40+
@prepare_map = {} if @prepare_map.length > 1000
41+
42+
return super unless SqlPatches.should_measure?
43+
super
44+
end
45+
46+
def exec(*args, &blk)
47+
return super unless SqlPatches.should_measure?
48+
49+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
50+
result = super
51+
elapsed_time = SqlPatches.elapsed_time(start)
52+
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
53+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
54+
55+
result
56+
end
57+
58+
if Gem::Version.new(PG::VERSION) >= Gem::Version.new("1.1.0")
59+
def exec_params(*args, &blk)
60+
return super unless SqlPatches.should_measure?
61+
62+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
63+
result = super
64+
elapsed_time = SqlPatches.elapsed_time(start)
65+
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
66+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
67+
68+
result
69+
end
70+
end
71+
72+
def exec_prepared(*args, &blk)
73+
return super unless SqlPatches.should_measure?
74+
75+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
76+
result = super
77+
elapsed_time = SqlPatches.elapsed_time(start)
78+
mapped = args[0]
79+
mapped = @prepare_map[mapped] || args[0] if @prepare_map
80+
record = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
81+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
82+
83+
result
84+
end
85+
86+
def send_query_prepared(*args, &blk)
87+
return super unless SqlPatches.should_measure?
88+
89+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
90+
result = super
91+
elapsed_time = SqlPatches.elapsed_time(start)
92+
mapped = args[0]
93+
mapped = @prepare_map[mapped] || args[0] if @prepare_map
94+
record = ::Rack::MiniProfiler.record_sql(mapped, elapsed_time)
95+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
96+
97+
result
98+
end
99+
100+
def async_exec(*args, &blk)
101+
return super unless SqlPatches.should_measure?
102+
103+
start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
104+
result = super
105+
elapsed_time = SqlPatches.elapsed_time(start)
106+
record = ::Rack::MiniProfiler.record_sql(args[0], elapsed_time)
107+
result.instance_variable_set("@miniprofiler_sql_id", record) if result
108+
109+
result
110+
end
111+
end
112+
113+
prepend MiniProfiler
114+
alias_method :query, :exec
115+
end

‎lib/prepend_pg_patch.rb

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
# frozen_string_literal: true
2+
3+
module Rack
4+
MINI_PROFILER_PREPEND_PG_PATCH = true
5+
end

0 commit comments

Comments
 (0)