2424 */
2525package co .elastic .apm .agent .jdbc .signature ;
2626
27+ import javax .annotation .Nullable ;
2728import java .util .concurrent .ConcurrentHashMap ;
2829import java .util .concurrent .ConcurrentMap ;
2930
3031import static co .elastic .apm .agent .jdbc .signature .Scanner .Token .EOF ;
3132import static co .elastic .apm .agent .jdbc .signature .Scanner .Token .FROM ;
3233import static co .elastic .apm .agent .jdbc .signature .Scanner .Token .IDENT ;
34+ import static co .elastic .apm .agent .jdbc .signature .Scanner .Token .INTO ;
3335import static co .elastic .apm .agent .jdbc .signature .Scanner .Token .LPAREN ;
3436import static co .elastic .apm .agent .jdbc .signature .Scanner .Token .RPAREN ;
3537
@@ -54,55 +56,61 @@ public class SignatureParser {
5456 * When relying on weak keys, we would not leverage any caching benefits if the query string is collected.
5557 * That means that we are leaking Strings but as the size of the map is limited that should not be an issue.
5658 */
57- private final static ConcurrentMap <String , String > signatureCache = new ConcurrentHashMap <String , String >(DISABLE_CACHE_THRESHOLD , 0.5f , Runtime .getRuntime ().availableProcessors ());
59+ private final static ConcurrentMap <String , String []> signatureCache = new ConcurrentHashMap <String , String []>(DISABLE_CACHE_THRESHOLD ,
60+ 0.5f , Runtime .getRuntime ().availableProcessors ());
5861
5962 private final Scanner scanner = new Scanner ();
6063
6164 public void querySignature (String query , StringBuilder signature , boolean preparedStatement ) {
65+ querySignature (query , signature , null , preparedStatement );
66+ }
6267
68+ public void querySignature (String query , StringBuilder signature , @ Nullable StringBuilder dbLink , boolean preparedStatement ) {
6369 final boolean cacheable = preparedStatement // non-prepared statements are likely to be dynamic strings
6470 && QUERY_LENGTH_CACHE_LOWER_THRESHOLD < query .length ()
6571 && query .length () < QUERY_LENGTH_CACHE_UPPER_THRESHOLD ;
6672 if (cacheable ) {
67- final String cachedSignature = signatureCache .get (query );
73+ final String [] cachedSignature = signatureCache .get (query );
6874 if (cachedSignature != null ) {
69- signature .append (cachedSignature );
75+ signature .append (cachedSignature [0 ]);
76+ if (dbLink != null ) {
77+ dbLink .append (cachedSignature [1 ]);
78+ }
7079 return ;
7180 }
7281 }
7382
7483 scanner .setQuery (query );
75- parse (query , signature );
84+ parse (query , signature , dbLink );
7685
7786 if (cacheable && signatureCache .size () <= DISABLE_CACHE_THRESHOLD ) {
7887 // we don't mind a small overshoot due to race conditions
79- signatureCache .put (query , signature .toString ());
88+ signatureCache .put (query , new String []{ signature .toString (), dbLink != null ? dbLink . toString () : "" } );
8089 }
8190 }
8291
83- private void parse (String query , StringBuilder signature ) {
92+ private void parse (String query , StringBuilder signature , @ Nullable StringBuilder dbLink ) {
8493 final Scanner .Token firstToken = scanner .scanWhile (Scanner .Token .COMMENT );
8594 switch (firstToken ) {
8695 case CALL :
8796 signature .append ("CALL" );
8897 if (scanner .scanUntil (Scanner .Token .IDENT )) {
89- signature .append (' ' );
90- scanner .appendCurrentTokenText (signature );
98+ appendIdentifiers (signature , dbLink );
9199 }
92100 return ;
93101 case DELETE :
94102 signature .append ("DELETE" );
95103 if (scanner .scanUntil (FROM ) && scanner .scanUntil (Scanner .Token .IDENT )) {
96- signature .append (" FROM " );
97- appendIdentifiers (signature );
104+ signature .append (" FROM" );
105+ appendIdentifiers (signature , dbLink );
98106 }
99107 return ;
100108 case INSERT :
101109 case REPLACE :
102110 signature .append (firstToken .name ());
103111 if (scanner .scanUntil (Scanner .Token .INTO ) && scanner .scanUntil (Scanner .Token .IDENT )) {
104- signature .append (" INTO " );
105- appendIdentifiers (signature );
112+ signature .append (" INTO" );
113+ appendIdentifiers (signature , dbLink );
106114 }
107115 return ;
108116 case SELECT :
@@ -116,8 +124,8 @@ private void parse(String query, StringBuilder signature) {
116124 } else if (t == FROM ) {
117125 if (level == 0 ) {
118126 if (scanner .scanToken (Scanner .Token .IDENT )) {
119- signature .append (" FROM " );
120- appendIdentifiers (signature );
127+ signature .append (" FROM" );
128+ appendIdentifiers (signature , dbLink );
121129 } else {
122130 return ;
123131 }
@@ -128,7 +136,7 @@ private void parse(String query, StringBuilder signature) {
128136 case UPDATE :
129137 signature .append ("UPDATE" );
130138 // Scan for the table name
131- boolean hasPeriod = false , hasFirstPeriod = false ;
139+ boolean hasPeriod = false , hasFirstPeriod = false , isDbLink = false ;
132140 if (scanner .scanToken (IDENT )) {
133141 signature .append (' ' );
134142 scanner .appendCurrentTokenText (signature );
@@ -145,6 +153,11 @@ private void parse(String query, StringBuilder signature) {
145153 signature .setLength (0 );
146154 signature .append ("UPDATE " );
147155 scanner .appendCurrentTokenText (signature );
156+ } else if (isDbLink ) {
157+ if (dbLink != null ) {
158+ scanner .appendCurrentTokenText (dbLink );
159+ }
160+ isDbLink = false ;
148161 }
149162 // Two adjacent identifiers found after the first period.
150163 // Ignore the secondary ones, in case they are unknown keywords.
@@ -155,23 +168,62 @@ private void parse(String query, StringBuilder signature) {
155168 signature .append ('.' );
156169 break ;
157170 default :
158- return ;
171+ if ("@" .equals (scanner .text ())) {
172+ isDbLink = true ;
173+ break ;
174+ } else {
175+ return ;
176+ }
159177 }
160178 }
161179 }
162180 return ;
181+ case MERGE :
182+ signature .append ("MERGE" );
183+ if (scanner .scanToken (INTO ) && scanner .scanUntil (Scanner .Token .IDENT )) {
184+ signature .append (" INTO" );
185+ appendIdentifiers (signature , dbLink );
186+ }
187+ return ;
163188 default :
164189 query = query .trim ();
165190 final int indexOfWhitespace = query .indexOf (' ' );
166191 signature .append (query , 0 , indexOfWhitespace > 0 ? indexOfWhitespace : query .length ());
167192 }
168193 }
169194
170- private void appendIdentifiers (StringBuilder signature ) {
195+ private void appendIdentifiers (StringBuilder signature , @ Nullable StringBuilder dbLink ) {
196+ signature .append (' ' );
171197 scanner .appendCurrentTokenText (signature );
172- while (scanner .scanToken (Scanner .Token .PERIOD ) && scanner .scanToken (Scanner .Token .IDENT )) {
173- signature .append ('.' );
174- scanner .appendCurrentTokenText (signature );
198+ boolean connectedIdents = false , isDbLink = false ;
199+ for (Scanner .Token t = scanner .scan (); t != EOF ; t = scanner .scan ()) {
200+ switch (t ) {
201+ case IDENT :
202+ // do not add tokens which are separated by a space
203+ if (connectedIdents ) {
204+ scanner .appendCurrentTokenText (signature );
205+ connectedIdents = false ;
206+ } else {
207+ if (isDbLink ) {
208+ if (dbLink != null ) {
209+ scanner .appendCurrentTokenText (dbLink );
210+ }
211+ }
212+ return ;
213+ }
214+ break ;
215+ case PERIOD :
216+ signature .append ('.' );
217+ connectedIdents = true ;
218+ break ;
219+ case USING :
220+ return ;
221+ default :
222+ if ("@" .equals (scanner .text ())) {
223+ isDbLink = true ;
224+ }
225+ break ;
226+ }
175227 }
176228 }
177229}
0 commit comments