-
Notifications
You must be signed in to change notification settings - Fork 56
/
TracingDriver.java
268 lines (225 loc) · 8.01 KB
/
TracingDriver.java
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
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
/*
* Copyright 2017-2021 The OpenTracing Authors
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
* in compliance with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software distributed under the License
* is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
* or implied. See the License for the specific language governing permissions and limitations under
* the License.
*/
package io.opentracing.contrib.jdbc;
import io.opentracing.Tracer;
import io.opentracing.contrib.common.WrapperProxy;
import io.opentracing.contrib.jdbc.parser.URLParser;
import io.opentracing.util.GlobalTracer;
import java.sql.Connection;
import java.sql.Driver;
import java.sql.DriverManager;
import java.sql.DriverPropertyInfo;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.List;
import java.util.Properties;
import java.util.Set;
import java.util.logging.Logger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
public class TracingDriver implements Driver {
private static final Driver INSTANCE = new TracingDriver();
protected static final String TRACE_WITH_ACTIVE_SPAN_ONLY = "traceWithActiveSpanOnly";
protected static final String WITH_ACTIVE_SPAN_ONLY = TRACE_WITH_ACTIVE_SPAN_ONLY + "=true";
public static final String IGNORE_FOR_TRACING_REGEX = "ignoreForTracing=\"((?:\\\\\"|[^\"])*)\"[;]*";
protected static final Pattern PATTERN_FOR_IGNORING = Pattern.compile(IGNORE_FOR_TRACING_REGEX);
static {
try {
DriverManager.registerDriver(INSTANCE);
} catch (SQLException e) {
throw new IllegalStateException("Could not register TracingDriver with DriverManager", e);
}
}
/**
* @return The singleton instance of the {@code TracingDriver}.
*/
public static Driver load() {
return INSTANCE;
}
/**
* Ensure {@code TracingDriver} be the first driver of {@link DriverManager} to make sure
* "interceptor mode" works. WARNING: Driver like Oracle JDBC may fail since it's destroyed
* forever after deregistration.
*/
public synchronized static void ensureRegisteredAsTheFirstDriver() {
try {
Enumeration<Driver> enumeration = DriverManager.getDrivers();
List<Driver> drivers = null;
for (int i = 0; enumeration.hasMoreElements(); ++i) {
Driver driver = enumeration.nextElement();
if (i == 0) {
if (driver == INSTANCE) {
return;
}
drivers = new ArrayList<>();
}
if (driver != INSTANCE) {
drivers.add(driver);
DriverManager.deregisterDriver(driver);
}
}
for (Driver driver : drivers) {
DriverManager.registerDriver(driver);
}
} catch (SQLException e) {
throw new IllegalStateException("Could not register TracingDriver with DriverManager", e);
}
}
/**
* Sets the {@code traceEnabled} property to enable or disable traces.
*
* @param traceEnabled The {@code traceEnabled} value.
*/
public static void setTraceEnabled(boolean traceEnabled) {
JdbcTracing.setTraceEnabled(traceEnabled);
}
public static boolean isTraceEnabled() {
return JdbcTracing.isTraceEnabled();
}
private static boolean interceptorMode = false;
/**
* Turns "interceptor mode" on or off.
*
* @param interceptorMode The {@code interceptorMode} value.
*/
public static void setInterceptorMode(final boolean interceptorMode) {
TracingDriver.interceptorMode = interceptorMode;
}
private static boolean withActiveSpanOnly;
/**
* Sets the {@code withActiveSpanOnly} property for "interceptor mode".
*
* @param withActiveSpanOnly The {@code withActiveSpanOnly} value.
*/
public static void setInterceptorProperty(final boolean withActiveSpanOnly) {
TracingDriver.withActiveSpanOnly = withActiveSpanOnly;
}
private static Set<String> ignoreStatements;
/**
* Sets the {@code ignoreStatements} property for "interceptor mode".
*
* @param ignoreStatements The {@code ignoreStatements} value.
*/
public static void setInterceptorProperty(final Set<String> ignoreStatements) {
TracingDriver.ignoreStatements = ignoreStatements;
}
protected Tracer tracer;
@Override
public Connection connect(String url, Properties info) throws SQLException {
// if there is no url, we have problems
if (url == null) {
throw new SQLException("url is required");
}
final Set<String> ignoreStatements;
final boolean withActiveSpanOnly;
if (interceptorMode) {
withActiveSpanOnly = TracingDriver.withActiveSpanOnly;
ignoreStatements = TracingDriver.ignoreStatements;
} else if (acceptsURL(url)) {
withActiveSpanOnly = url.contains(WITH_ACTIVE_SPAN_ONLY);
ignoreStatements = extractIgnoredStatements(url);
} else {
return null;
}
url = extractRealUrl(url);
// find the real driver for the URL
final Driver wrappedDriver = findDriver(url);
final Tracer currentTracer = getTracer();
final ConnectionInfo connectionInfo = URLParser.parse(url);
final String realUrl = url;
final Connection connection = JdbcTracingUtils.call("AcquireConnection", () ->
wrappedDriver.connect(realUrl, info), null, connectionInfo, withActiveSpanOnly,
null, currentTracer);
return WrapperProxy
.wrap(connection, new TracingConnection(connection, connectionInfo, withActiveSpanOnly,
ignoreStatements, currentTracer));
}
@Override
public boolean acceptsURL(String url) throws SQLException {
return url != null && (
url.startsWith(getUrlPrefix()) ||
(interceptorMode && url.startsWith("jdbc:"))
);
}
@Override
public DriverPropertyInfo[] getPropertyInfo(String url, Properties info) throws SQLException {
return findDriver(url).getPropertyInfo(url, info);
}
@Override
public int getMajorVersion() {
// There is no way to get it from wrapped driver
return 1;
}
@Override
public int getMinorVersion() {
// There is no way to get it from wrapped driver
return 0;
}
@Override
public boolean jdbcCompliant() {
return true;
}
@Override
public Logger getParentLogger() throws SQLFeatureNotSupportedException {
// There is no way to get it from wrapped driver
return null;
}
public void setTracer(Tracer tracer) {
this.tracer = tracer;
}
protected String getUrlPrefix() {
return "jdbc:tracing:";
}
protected Driver findDriver(String realUrl) throws SQLException {
if (realUrl == null || realUrl.trim().length() == 0) {
throw new IllegalArgumentException("url is required");
}
for (Driver candidate : Collections.list(DriverManager.getDrivers())) {
try {
if (!(candidate instanceof TracingDriver) && candidate.acceptsURL(realUrl)) {
return candidate;
}
} catch (SQLException ignored) {
// intentionally ignore exception
}
}
throw new SQLException("Unable to find a driver that accepts url: " + realUrl);
}
protected String extractRealUrl(String url) {
String extracted = url.startsWith(getUrlPrefix()) ? url.replace(getUrlPrefix(), "jdbc:") : url;
return extracted.replaceAll(TRACE_WITH_ACTIVE_SPAN_ONLY + "=(true|false)[;]*", "")
.replaceAll(IGNORE_FOR_TRACING_REGEX, "")
.replaceAll("\\?$", "");
}
protected Set<String> extractIgnoredStatements(String url) {
final Matcher matcher = PATTERN_FOR_IGNORING.matcher(url);
Set<String> results = new HashSet<>(8);
while (matcher.find()) {
String rawValue = matcher.group(1);
String finalValue = rawValue.replace("\\\"", "\"");
results.add(finalValue);
}
return results;
}
Tracer getTracer() {
if (tracer == null) {
return GlobalTracer.get();
}
return tracer;
}
}