Skip to content

Commit cd18597

Browse files
sethsamuelsethsamuel
authored
Merge pull request #6807 from DataDog/seth.samuel/SDBM-876-PSQLException-when-enabling-dbm-and-apm-correlation
Handle comment injection for procedure CALLs
2 parents a598d44 + 8a1b83c commit cd18597

File tree

3 files changed

+114
-12
lines changed

3 files changed

+114
-12
lines changed

‎dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java‎

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,18 @@ public static String prepend(
4444
return inject(sql, dbService, dbType, hostname, dbName, null, false, false);
4545
}
4646

47+
public static String getFirstWord(String sql) {
48+
int beginIndex = 0;
49+
while (beginIndex < sql.length() && Character.isWhitespace(sql.charAt(beginIndex))) {
50+
beginIndex++;
51+
}
52+
int endIndex = beginIndex;
53+
while (endIndex < sql.length() && !Character.isWhitespace(sql.charAt(endIndex))) {
54+
endIndex++;
55+
}
56+
return sql.substring(beginIndex, endIndex);
57+
}
58+
4759
public static String inject(
4860
final String sql,
4961
final String dbService,
@@ -52,26 +64,29 @@ public static String inject(
5264
final String dbName,
5365
final String traceParent,
5466
final boolean injectTrace,
55-
final boolean appendComment) {
67+
boolean appendComment) {
5668
if (sql == null || sql.isEmpty()) {
5769
return sql;
5870
}
5971
if (hasDDComment(sql, appendComment)) {
6072
return sql;
6173
}
6274

63-
if (dbType != null && dbType.startsWith("postgres")) {
64-
// The Postgres JDBC parser doesn't allow SQL comments anywhere in a callable statement
65-
//
75+
if (dbType != null) {
76+
final String firstWord = getFirstWord(sql);
77+
78+
// The Postgres JDBC parser doesn't allow SQL comments anywhere in a JDBC callable statements
6679
// https://github.com/pgjdbc/pgjdbc/blob/master/pgjdbc/src/main/java/org/postgresql/core/Parser.java#L1038
67-
// TODO: Could we inject the comment after the JDBC has been converted to standard SQL?
68-
int charIndex = 0;
69-
while (charIndex < sql.length() && Character.isWhitespace(sql.charAt(charIndex))) {
70-
charIndex++;
71-
}
72-
if (charIndex < sql.length() && sql.charAt(charIndex) == '{') {
80+
// TODO: Could we inject the comment after the JDBC has been converted to standard SQL?
81+
if (firstWord.startsWith("{") && dbType.startsWith("postgres")) {
7382
return sql;
7483
}
84+
85+
// Both Postgres and MySQL are unhappy with anything before CALL in a stored procedure
86+
// invocation but they seem ok with it after so we force append mode
87+
if (firstWord.equalsIgnoreCase("call")) {
88+
appendComment = true;
89+
}
7590
}
7691

7792
final Config config = Config.get();

‎dd-java-agent/instrumentation/jdbc/src/test/groovy/RemoteJDBCInstrumentationTest.groovy‎

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -481,9 +481,10 @@ abstract class RemoteJDBCInstrumentationTest extends VersionedNamingTestBase {
481481
}
482482

483483

484-
def "prepared call with storedproc on #driver with #connection.getClass().getCanonicalName() does not hang"() {
484+
def "prepared procedure call with return value on #driver with #connection.getClass().getCanonicalName() does not hang"() {
485485
setup:
486486
injectSysConfig("dd.dbm.propagation.mode", "full")
487+
487488
CallableStatement upperProc = connection.prepareCall(query)
488489
upperProc.registerOutParameter(1, Types.VARCHAR)
489490
upperProc.setString(2, "hello world")
@@ -505,12 +506,76 @@ abstract class RemoteJDBCInstrumentationTest extends VersionedNamingTestBase {
505506
"mysql" | cpDatasources.get("hikari").get(driver).getConnection() | "{ ? = call upper( ? ) }"
506507
"postgresql" | cpDatasources.get("tomcat").get(driver).getConnection() | " { ? = call upper( ? ) }"
507508
"mysql" | cpDatasources.get("tomcat").get(driver).getConnection() | "{ ? = call upper( ? ) }"
508-
"postgresql" | cpDatasources.get("c3p0").get(driver).getConnection() | "{ ? = call upper( ? ) }"
509+
"postgresql" | cpDatasources.get("c3p0").get(driver).getConnection() | " { ? = call upper( ? ) }"
509510
"mysql" | cpDatasources.get("c3p0").get(driver).getConnection() | "{ ? = call upper( ? ) }"
510511
"postgresql" | connectTo(driver, peerConnectionProps) | " { ? = call upper( ? ) }"
511512
"mysql" | connectTo(driver, peerConnectionProps) | " { ? = call upper( ? ) }"
512513
}
513514

515+
def "prepared procedure call on #driver with #connection.getClass().getCanonicalName() does not hang"() {
516+
setup:
517+
518+
String createSql
519+
if (driver == "postgresql") {
520+
createSql =
521+
"""
522+
CREATE OR REPLACE PROCEDURE dummy(inout res integer)
523+
LANGUAGE SQL
524+
AS \$\$
525+
SELECT 1;
526+
\$\$;
527+
"""
528+
} else if (driver == "mysql") {
529+
createSql =
530+
"""
531+
CREATE PROCEDURE IF NOT EXISTS dummy(inout res int)
532+
BEGIN
533+
SELECT 1;
534+
END
535+
"""
536+
} else {
537+
assert false
538+
}
539+
540+
if (driver.equals("postgresql") && connection.getMetaData().getDatabaseMajorVersion() <= 11) {
541+
// Skip test for older versions of PG that don't support out on procedure
542+
return
543+
}
544+
545+
546+
connection.prepareCall(createSql).execute()
547+
548+
injectSysConfig("dd.dbm.propagation.mode", "full")
549+
CallableStatement proc = connection.prepareCall(query)
550+
proc.setInt(1,1)
551+
proc.registerOutParameter(1, Types.INTEGER)
552+
when:
553+
runUnderTrace("parent") {
554+
return proc.execute()
555+
}
556+
TEST_WRITER.waitForTraces(1)
557+
558+
then:
559+
assert proc.getInt(1) == 1
560+
561+
cleanup:
562+
if (proc != null) {
563+
proc.close()
564+
}
565+
connection.close()
566+
567+
where:
568+
driver | connection | query
569+
"postgresql" | cpDatasources.get("hikari").get(driver).getConnection() | "CALL dummy(?)"
570+
"mysql" | cpDatasources.get("hikari").get(driver).getConnection() | "CALL dummy(?)"
571+
"postgresql" | cpDatasources.get("tomcat").get(driver).getConnection() | " CALL dummy(?)"
572+
"mysql" | cpDatasources.get("tomcat").get(driver).getConnection() | "CALL dummy(?)"
573+
"postgresql" | cpDatasources.get("c3p0").get(driver).getConnection() | " CALL dummy(?)"
574+
"mysql" | cpDatasources.get("c3p0").get(driver).getConnection() | "CALL dummy(?)"
575+
"postgresql" | connectTo(driver, peerConnectionProps) | " CALL dummy(?)"
576+
"mysql" | connectTo(driver, peerConnectionProps) | "CALL dummy(?)"
577+
}
578+
514579

515580
Driver driverFor(String db) {
516581
return newDriver(jdbcDriverClassNames.get(db))

‎dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,26 @@ import datadog.trace.instrumentation.jdbc.SQLCommenter
33

44
class SQLCommenterTest extends AgentTestRunner {
55

6+
def "test find first word"() {
7+
setup:
8+
9+
when:
10+
String word = SQLCommenter.getFirstWord(sql)
11+
12+
then:
13+
word == firstWord
14+
15+
where:
16+
sql | firstWord
17+
"SELECT *" | "SELECT"
18+
" { " | "{"
19+
"{" | "{"
20+
"CALL ( ? )" | "CALL"
21+
"" | ""
22+
" " | ""
23+
}
24+
25+
626
def "test encode Sql Comment"() {
727
setup:
828
injectSysConfig("dd.service", ddService)
@@ -32,6 +52,8 @@ class SQLCommenterTest extends AgentTestRunner {
3252
"SELECT * FROM foo" | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
3353
"{call dogshelterProc(?, ?)}" | "SqlCommenter" | "Test" | "my-service" | "mysql" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "{call dogshelterProc(?, ?)} /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
3454
"{call dogshelterProc(?, ?)}" | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "{call dogshelterProc(?, ?)}"
55+
"CALL dogshelterProc(?, ?)" | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "CALL dogshelterProc(?, ?) /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
56+
"CALL dogshelterProc(?, ?)" | "SqlCommenter" | "Test" | "my-service" | "mysql" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "CALL dogshelterProc(?, ?) /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
3557
"SELECT * FROM foo" | "" | "Test" | "" | "mysql" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "SELECT * FROM foo /*ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
3658
"SELECT * FROM foo" | "" | "Test" | "my-service" | "mysql" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "SELECT * FROM foo /*dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
3759
"SELECT * FROM foo" | "SqlCommenter" | "Test" | "my-service" | "mysql" | "" | "" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"

0 commit comments

Comments
 (0)