/
mysql2.rb
147 lines (120 loc) · 4.01 KB
/
mysql2.rb
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
# frozen_string_literal: true
require "semian/adapter"
require "mysql2"
module Mysql2
Mysql2::Error.include(::Semian::AdapterError)
class SemianError < Mysql2::Error
def initialize(semian_identifier, *args)
super(*args)
@semian_identifier = semian_identifier
end
end
ResourceBusyError = Class.new(SemianError)
CircuitOpenError = Class.new(SemianError)
end
module Semian
module Mysql2
include Semian::Adapter
CONNECTION_ERROR = Regexp.union(
/Can't connect to (?:MySQL )?server on/i,
/Lost connection to (?:MySQL )?server/i,
/MySQL server has gone away/i,
/Too many connections/i,
/closed MySQL connection/i,
/Timeout waiting for a response/i,
/No matching servers with free connections/i,
/Max connect timeout reached while reaching hostgroup/i,
)
ResourceBusyError = ::Mysql2::ResourceBusyError
CircuitOpenError = ::Mysql2::CircuitOpenError
PingFailure = Class.new(::Mysql2::Error)
DEFAULT_HOST = "localhost"
DEFAULT_PORT = 3306
QUERY_ALLOWLIST = %r{\A(?:/\*.*?\*/)?\s*(ROLLBACK|COMMIT|RELEASE\s+SAVEPOINT)}i
class << self
# The naked methods are exposed as `raw_query` and `raw_connect` for instrumentation purpose
def included(base)
base.send(:alias_method, :raw_query, :query)
base.send(:remove_method, :query)
base.send(:alias_method, :raw_connect, :connect)
base.send(:remove_method, :connect)
base.send(:alias_method, :raw_ping, :ping)
base.send(:remove_method, :ping)
end
end
def semian_identifier
@semian_identifier ||= begin
name = semian_options && semian_options[:name]
unless name
host = query_options[:host] || DEFAULT_HOST
port = query_options[:port] || DEFAULT_PORT
name = "#{host}:#{port}"
end
:"mysql_#{name}"
end
end
def ping
return false if closed?
result = nil
acquire_semian_resource(adapter: :mysql, scope: :ping) do
result = raw_ping
raise PingFailure, result.to_s unless result
end
result
rescue ResourceBusyError, CircuitOpenError, PingFailure
false
end
def query(*args)
if query_whitelisted?(*args)
raw_query(*args)
else
acquire_semian_resource(adapter: :mysql, scope: :query) { raw_query(*args) }
end
end
# TODO: write_timeout and connect_timeout can't be configured currently
# dynamically, await https://github.com/brianmario/mysql2/pull/955
def with_resource_timeout(temp_timeout)
prev_read_timeout = @read_timeout
begin
# C-ext reads this directly, writer method will configure
# properly on the client but based on my read--this is good enough
# until we get https://github.com/brianmario/mysql2/pull/955 in
@read_timeout = temp_timeout
yield
ensure
@read_timeout = prev_read_timeout
end
end
private
EXCEPTIONS = [].freeze
def resource_exceptions
EXCEPTIONS
end
def query_whitelisted?(sql, *)
QUERY_ALLOWLIST =~ sql
rescue ArgumentError
# The above regexp match can fail if the input SQL string contains binary
# data that is not recognized as a valid encoding, in which case we just
# return false.
return false unless sql.valid_encoding?
raise
end
def connect(*args)
acquire_semian_resource(adapter: :mysql, scope: :connection) { raw_connect(*args) }
end
def acquire_semian_resource(**)
super
rescue ::Mysql2::Error => error
if error.is_a?(PingFailure) || (!error.is_a?(::Mysql2::SemianError) && error.message.match?(CONNECTION_ERROR))
semian_resource.mark_failed(error)
error.semian_identifier = semian_identifier
end
raise
end
def raw_semian_options
return query_options[:semian] if query_options.key?(:semian)
query_options["semian"] if query_options.key?("semian")
end
end
end
::Mysql2::Client.include(Semian::Mysql2)