further SQL parsing improvement (#676)

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

@@ -0,0 +1,75 @@
+/*-
+ * #%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;
+
+class JdbcFilter {
+
+    private boolean inQuote = false;
+    private boolean inJdbcEscape = false;
+    private boolean jdbcKeyWord = false;
+
+    boolean skip(Scanner s, char c) {
+        switch (c) {
+            case '{':
+                if (!inQuote) {
+                    inJdbcEscape = true;
+                    jdbcKeyWord = true;
+                    return true;
+                }
+                break;
+            case 'o':
+            case 'O':
+                if (!inQuote && inJdbcEscape && jdbcKeyWord && s.isNextCharIgnoreCase('j')) {
+                    s.next();
+                    jdbcKeyWord = false;
+                    return true;
+                }
+                break;
+            case '}':
+                if (!inQuote) {
+                    inJdbcEscape = false;
+                    return true;
+                }
+                break;
+            case '?':
+            case '=':
+                if (!inQuote && inJdbcEscape) {
+                    return true;
+                }
+                break;
+            case '\'':
+                inQuote = !inQuote;
+                break;
+        }
+        jdbcKeyWord = false;
+        return false;
+    }
+
+    void reset() {
+        inQuote = false;
+        inJdbcEscape = false;
+        jdbcKeyWord = false;
+    }
+}

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

@@ -31,9 +31,11 @@ public class Scanner {
     private int end; // text end char offset
     private int pos; // read position char offset
     private int inputLength;
+    private final JdbcFilter filter = new JdbcFilter();
 
     public void setQuery(String sql) {
         this.input = sql;
+        filter.reset();
         inputLength = sql.length();
         start = 0;
         end = 0;
@@ -74,7 +76,7 @@ public Token scan() {
             return Token.EOF;
         }
         char c = next();
-        while (Character.isSpaceChar(c)) {
+        while (Character.isSpaceChar(c) || filter.skip(this, c)) {
             if (hasNext()) {
                 c = next();
             } else {
@@ -301,7 +303,7 @@ private char peek() {
         return input.charAt(pos);
     }
 
-    private char next() {
+    char next() {
         final char c = peek();
         pos++;
         end = pos;
@@ -347,6 +349,10 @@ private boolean isNextChar(char c) {
         return hasNext() && peek() == c;
     }
 
+    boolean isNextCharIgnoreCase(char c) {
+        return hasNext() && Character.toLowerCase(peek()) == Character.toLowerCase(c);
+    }
+
     public enum Token {
 
         OTHER,
@@ -373,7 +379,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 +390,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

@@ -24,12 +24,14 @@
  */
 package co.elastic.apm.agent.jdbc.signature;
 
+import javax.annotation.Nullable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
 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.INTO;
 import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.LPAREN;
 import static co.elastic.apm.agent.jdbc.signature.Scanner.Token.RPAREN;
 
@@ -54,55 +56,61 @@ 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, null, preparedStatement);
+    }
 
+    public void querySignature(String query, StringBuilder signature, @Nullable 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]);
+                if (dbLink != null) {
+                    dbLink.append(cachedSignature[1]);
+                }
                 return;
             }
         }
 
         scanner.setQuery(query);
-        parse(query, signature);
+        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 != null ? dbLink.toString() : ""});
         }
     }
 
-    private void parse(String query, StringBuilder signature) {
+    private void parse(String query, StringBuilder signature, @Nullable 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);
+                    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 +124,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 +136,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);
@@ -145,6 +153,11 @@ private void parse(String query, StringBuilder signature) {
                                     signature.setLength(0);
                                     signature.append("UPDATE ");
                                     scanner.appendCurrentTokenText(signature);
+                                } else if (isDbLink) {
+                                    if (dbLink != null) {
+                                        scanner.appendCurrentTokenText(dbLink);
+                                    }
+                                    isDbLink = false;
                                 }
                                 // Two adjacent identifiers found after the first period.
                                 // Ignore the secondary ones, in case they are unknown keywords.
@@ -155,23 +168,62 @@ 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) {
+    private void appendIdentifiers(StringBuilder signature, @Nullable StringBuilder dbLink) {
+        signature.append(' ');
         scanner.appendCurrentTokenText(signature);
-        while (scanner.scanToken(Scanner.Token.PERIOD) && scanner.scanToken(Scanner.Token.IDENT)) {
-            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) {
+                            if (dbLink != null) {
+                                scanner.appendCurrentTokenText(dbLink);
+                            }
+                        }
+                        return;
+                    }
+                    break;
+                case PERIOD:
+                    signature.append('.');
+                    connectedIdents = true;
+                    break;
+                case USING:
+                    return;
+                default:
+                    if ("@".equals(scanner.text())) {
+                        isDbLink = true;
+                    }
+                    break;
+            }
         }
     }
 }