further SQL parsing improvement (#676)

apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/signature/Scanner.java

@@ -373,7 +373,9 @@ public enum Token {
         SET,
         TABLE,
         TRUNCATE, // Cassandra/CQL-specific
-        UPDATE;
+        UPDATE,
+        MERGE,
+        USING;
 
         private static final Token[] EMPTY = {};
         private static final Token[][] KEYWORDS_BY_LENGTH = {
@@ -382,7 +384,7 @@ public enum Token {
             {AS, OR},
             {SET},
             {CALL, FROM, INTO},
-            {TABLE},
+            {TABLE, MERGE, USING},
             {DELETE, INSERT, SELECT, UPDATE},
             {REPLACE},
             {TRUNCATE}

apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/signature/SignatureParser.java

@@ -27,11 +27,14 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
+import co.elastic.apm.agent.jdbc.signature.filter.FilterChain;
+
 import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.EOF;
 import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.FROM;
 import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.IDENT;
 import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.LPAREN;
 import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.RPAREN;
+import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.INTO;
 
 public class SignatureParser {
 
@@ -54,55 +57,56 @@ public class SignatureParser {
      * When relying on weak keys, we would not leverage any caching benefits if the query string is collected.
      * That means that we are leaking Strings but as the size of the map is limited that should not be an issue.
      */
-    private final static ConcurrentMap<String, String> signatureCache = new ConcurrentHashMap<String, String>(DISABLE_CACHE_THRESHOLD, 0.5f, Runtime.getRuntime().availableProcessors());
+    private final static ConcurrentMap<String, String[]> signatureCache = new ConcurrentHashMap<String, String[]>(DISABLE_CACHE_THRESHOLD, 0.5f, Runtime.getRuntime().availableProcessors());
 
     private final Scanner scanner = new Scanner();
 
     public void querySignature(String query, StringBuilder signature, boolean preparedStatement) {
-
+    	querySignature(query, signature, new StringBuilder(), preparedStatement);
+    }
+    public void querySignature(String query, StringBuilder signature, StringBuilder dbLink, boolean preparedStatement) {
         final boolean cacheable = preparedStatement // non-prepared statements are likely to be dynamic strings
             && QUERY_LENGTH_CACHE_LOWER_THRESHOLD < query.length()
             && query.length() < QUERY_LENGTH_CACHE_UPPER_THRESHOLD;
         if (cacheable) {
-            final String cachedSignature = signatureCache.get(query);
+            final String[] cachedSignature = signatureCache.get(query);
             if (cachedSignature != null) {
-                signature.append(cachedSignature);
+                signature.append(cachedSignature[0]);
+                dbLink.append(cachedSignature[1]);
                 return;
             }
         }
 
-        scanner.setQuery(query);
-        parse(query, signature);
+        scanner.setQuery(FilterChain.DEFAULT_CHAIN.doFilter(query));
+        parse(query, signature, dbLink);
 
         if (cacheable && signatureCache.size() <= DISABLE_CACHE_THRESHOLD) {
             // we don't mind a small overshoot due to race conditions
-            signatureCache.put(query, signature.toString());
+            signatureCache.put(query, new String[] {signature.toString(), dbLink.toString()});
         }
     }
-
-    private void parse(String query, StringBuilder signature) {
+    private void parse(String query, StringBuilder signature, StringBuilder dbLink) {
         final Scanner.Token firstToken = scanner.scanWhile(Scanner.Token.COMMENT);
         switch (firstToken) {
             case CALL:
                 signature.append("CALL");
-                if (scanner.scanUntil(Scanner.Token.IDENT)) {
-                    signature.append(' ');
-                    scanner.appendCurrentTokenText(signature);
+                if(scanner.scanUntil(Scanner.Token.IDENT)) {
+                	appendIdentifiers(signature, dbLink);
                 }
                 return;
             case DELETE:
                 signature.append("DELETE");
                 if (scanner.scanUntil(FROM) && scanner.scanUntil(Scanner.Token.IDENT)) {
-                    signature.append(" FROM ");
-                    appendIdentifiers(signature);
+                    signature.append(" FROM");
+                    appendIdentifiers(signature, dbLink);
                 }
                 return;
             case INSERT:
             case REPLACE:
                 signature.append(firstToken.name());
                 if (scanner.scanUntil(Scanner.Token.INTO) && scanner.scanUntil(Scanner.Token.IDENT)) {
-                    signature.append(" INTO ");
-                    appendIdentifiers(signature);
+                    signature.append(" INTO");
+                    appendIdentifiers(signature, dbLink);
                 }
                 return;
             case SELECT:
@@ -116,8 +120,8 @@ private void parse(String query, StringBuilder signature) {
                     } else if (t == FROM) {
                         if (level == 0) {
                             if (scanner.scanToken(Scanner.Token.IDENT)) {
-                                signature.append(" FROM ");
-                                appendIdentifiers(signature);
+                                signature.append(" FROM");
+                                appendIdentifiers(signature, dbLink);
                             } else {
                                 return;
                             }
@@ -128,7 +132,7 @@ private void parse(String query, StringBuilder signature) {
             case UPDATE:
                 signature.append("UPDATE");
                 // Scan for the table name
-                boolean hasPeriod = false, hasFirstPeriod = false;
+                boolean hasPeriod = false, hasFirstPeriod = false, isDbLink = false;
                 if (scanner.scanToken(IDENT)) {
                     signature.append(' ');
                     scanner.appendCurrentTokenText(signature);
@@ -142,10 +146,14 @@ private void parse(String query, StringBuilder signature) {
                                 if (!hasFirstPeriod) {
                                     // Some dialects allow option keywords before the table name
                                     // example: UPDATE IGNORE foo.bar
-                                    signature.setLength(0);
-                                    signature.append("UPDATE ");
+                                	signature.setLength(0);
+                                	signature.append("UPDATE ");
                                     scanner.appendCurrentTokenText(signature);
                                 }
+                                else if(isDbLink) {
+            						scanner.appendCurrentTokenText(dbLink);
+            						isDbLink = false;
+            					}
                                 // Two adjacent identifiers found after the first period.
                                 // Ignore the secondary ones, in case they are unknown keywords.
                                 break;
@@ -155,23 +163,61 @@ private void parse(String query, StringBuilder signature) {
                                 signature.append('.');
                                 break;
                             default:
-                                return;
+                				if("@".equals(scanner.text())) {
+                					isDbLink = true;
+                					break;
+                				}else {
+                                    return;
+                				}
                         }
                     }
                 }
                 return;
+            case MERGE:
+                signature.append("MERGE");
+                if(scanner.scanToken(INTO) && scanner.scanUntil(Scanner.Token.IDENT)) {
+                    signature.append(" INTO");
+                    appendIdentifiers(signature, dbLink);
+                }
+                return;
             default:
                 query = query.trim();
                 final int indexOfWhitespace = query.indexOf(' ');
                 signature.append(query, 0, indexOfWhitespace > 0 ? indexOfWhitespace : query.length());
         }
     }
 
-    private void appendIdentifiers(StringBuilder signature) {
-        scanner.appendCurrentTokenText(signature);
-        while (scanner.scanToken(Scanner.Token.PERIOD) && scanner.scanToken(Scanner.Token.IDENT)) {
-            signature.append('.');
-            scanner.appendCurrentTokenText(signature);
-        }
+    private void appendIdentifiers(StringBuilder signature, StringBuilder dbLink) {
+    	signature.append(' ');
+		scanner.appendCurrentTokenText(signature);
+		boolean connectedIdents = false, isDbLink = false;
+		for (Scanner.Token t = scanner.scan(); t != EOF; t = scanner.scan()) {
+			switch (t) {
+			case IDENT:
+				// do not add tokens which are separated by a space
+				if (connectedIdents) {
+					scanner.appendCurrentTokenText(signature);
+					connectedIdents = false;
+				} else {
+					if(isDbLink) {
+						scanner.appendCurrentTokenText(dbLink);
+						isDbLink = false;
+					}
+					return;
+				}
+				break;
+			case PERIOD:
+				signature.append('.');
+				connectedIdents = true;
+				break;
+			case USING:
+				return;
+			default:
+				if("@".equals(scanner.text())) {
+					isDbLink = true;
+				}
+				break;
+			}
+		}
     }
 }

apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/signature/filter/Filter.java

@@ -0,0 +1,29 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2019 Elastic and contributors
+ * %%
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ * #L%
+ */
+package co.elastic.apm.agent.jdbc.signature.filter;
+
+public interface Filter {
+	String doFilter(String sql);
+}

apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/signature/filter/FilterChain.java

@@ -0,0 +1,51 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2019 Elastic and contributors
+ * %%
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ * #L%
+ */
+package co.elastic.apm.agent.jdbc.signature.filter;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class FilterChain implements Filter {
+	public static final FilterChain DEFAULT_CHAIN = new FilterChain(new JdbcFilter());
+	private final List<Filter> filter=new ArrayList<>();
+	public FilterChain() {}
+	public FilterChain(Filter... filter) {
+		for(int i=0;i<filter.length;i++) {
+			this.filter.add(filter[i]);
+		}
+	}
+	public void addFilter(Filter f) {
+		filter.add(f);
+	}
+
+	@Override
+	public String doFilter(String sql) {
+		String current = sql;
+		for(Filter f : filter) {
+			current=f.doFilter(current);
+		}
+		return current;
+	}
+}

apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/signature/filter/JdbcFilter.java

@@ -0,0 +1,96 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2019 Elastic and contributors
+ * %%
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ * #L%
+ */
+package co.elastic.apm.agent.jdbc.signature.filter;
+
+public class JdbcFilter implements Filter {
+	private static final ThreadLocal<StringBuilder> cache=new ThreadLocal<StringBuilder>() {
+	        @Override protected StringBuilder initialValue() {
+	            return new StringBuilder();
+	    }
+	};
+	private int findFirstNonWhitespace(String str, int pos) {
+		for(int i = pos+1; i < str.length(); i++){
+	        if(!Character.isWhitespace(str.charAt(i))){
+	        	return i;
+	        }
+	    }
+		return -1;
+	}
+	@Override
+	public String doFilter(String sql) {
+		if(!sql.contains("{")) {
+			return sql;
+		}
+		StringBuilder sb = cache.get();
+		sb.setLength(0);
+		boolean inQuote = false, inJdbcEscape = false;
+		char c;
+		int i = 0,j = 0;
+		while(i<sql.length()) {
+			c = sql.charAt(i);
+			switch(c) {
+				case '{':
+					if(inQuote) {
+						sb.append(c);
+					}else {
+						inJdbcEscape = true;
+					}
+					j=findFirstNonWhitespace(sql, i);
+					if(j>i) {
+						i=j;
+					}
+					if(sql.toLowerCase().indexOf("oj", i) == i) {
+						i+=2;
+					}
+					break;
+				case '}':
+					if(inQuote) {
+						sb.append(c);
+					}else {
+						inJdbcEscape = false;
+					}
+					i++;
+					break;
+				case '?':
+				case '=':
+					if(inQuote || !inJdbcEscape) {
+						sb.append(c);
+					}
+					i++;
+					break;
+				case '\'':
+					sb.append(c);
+					inQuote = !inQuote;
+					i++;
+					break;
+				default:
+					sb.append(c);
+					i++;
+					break;
+			}
+		}
+		return sb.toString();
+	}
+}