perf: qualify statements without removing comments

Determine the type of statement without first removing all comments and hints.
This prevents the creation of new strings and stepping through the entire SQL
string for each statement that is not found in the statement cache.

Benchmark                                Mode  Cnt       Score      Error  Units
StatementParserBenchmark.isQueryTest    thrpt    5  547904.501 ± 1970.170  ops/s
StatementParserBenchmark.longDmlTest    thrpt    5  114806.782 ±  826.881  ops/s
StatementParserBenchmark.longQueryTest  thrpt    5  112666.992 ±  700.783  ops/s

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java

@@ -27,24 +27,28 @@
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractBaseUnitOfWork.InterceptorsUsage;
+import com.google.cloud.spanner.connection.SimpleParser.Result;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
 import com.google.cloud.spanner.connection.UnitOfWork.CallType;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Suppliers;
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheStats;
 import com.google.common.cache.Weigher;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
+import java.nio.CharBuffer;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.function.Supplier;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -179,24 +183,24 @@ public static class ParsedStatement {
     private final StatementType type;
     private final ClientSideStatementImpl clientSideStatement;
     private final Statement statement;
-    private final String sqlWithoutComments;
-    private final boolean returningClause;
+    private final Supplier<String> sqlWithoutComments;
+    private final Supplier<Boolean> returningClause;
     private final ReadQueryUpdateTransactionOption[] optionsFromHints;
 
     private static ParsedStatement clientSideStatement(
         ClientSideStatementImpl clientSideStatement,
         Statement statement,
-        String sqlWithoutComments) {
+        Supplier<String> sqlWithoutComments) {
       return new ParsedStatement(clientSideStatement, statement, sqlWithoutComments);
     }
 
-    private static ParsedStatement ddl(Statement statement, String sqlWithoutComments) {
+    private static ParsedStatement ddl(Statement statement, Supplier<String> sqlWithoutComments) {
       return new ParsedStatement(StatementType.DDL, statement, sqlWithoutComments);
     }
 
     private static ParsedStatement query(
         Statement statement,
-        String sqlWithoutComments,
+        Supplier<String> sqlWithoutComments,
         QueryOptions defaultQueryOptions,
         ReadQueryUpdateTransactionOption[] optionsFromHints) {
       return new ParsedStatement(
@@ -205,57 +209,66 @@ private static ParsedStatement query(
           statement,
           sqlWithoutComments,
           defaultQueryOptions,
-          false,
+          Suppliers.ofInstance(false),
           optionsFromHints);
     }
 
     private static ParsedStatement update(
         Statement statement,
-        String sqlWithoutComments,
-        boolean returningClause,
+        Supplier<String> sqlWithoutComments,
+        Supplier<Boolean> returningClause,
         ReadQueryUpdateTransactionOption[] optionsFromHints) {
       return new ParsedStatement(
           StatementType.UPDATE, statement, sqlWithoutComments, returningClause, optionsFromHints);
     }
 
-    private static ParsedStatement unknown(Statement statement, String sqlWithoutComments) {
+    private static ParsedStatement unknown(
+        Statement statement, Supplier<String> sqlWithoutComments) {
       return new ParsedStatement(StatementType.UNKNOWN, statement, sqlWithoutComments);
     }
 
     private ParsedStatement(
         ClientSideStatementImpl clientSideStatement,
         Statement statement,
-        String sqlWithoutComments) {
+        Supplier<String> sqlWithoutComments) {
       Preconditions.checkNotNull(clientSideStatement);
       Preconditions.checkNotNull(statement);
       this.type = StatementType.CLIENT_SIDE;
       this.clientSideStatement = clientSideStatement;
       this.statement = statement;
-      this.sqlWithoutComments = Preconditions.checkNotNull(sqlWithoutComments);
-      this.returningClause = false;
+      this.sqlWithoutComments = sqlWithoutComments;
+      this.returningClause = Suppliers.ofInstance(false);
       this.optionsFromHints = EMPTY_OPTIONS;
     }
 
     private ParsedStatement(
         StatementType type,
         Statement statement,
-        String sqlWithoutComments,
-        boolean returningClause,
+        Supplier<String> sqlWithoutComments,
+        Supplier<Boolean> returningClause,
         ReadQueryUpdateTransactionOption[] optionsFromHints) {
       this(type, null, statement, sqlWithoutComments, null, returningClause, optionsFromHints);
     }
 
-    private ParsedStatement(StatementType type, Statement statement, String sqlWithoutComments) {
-      this(type, null, statement, sqlWithoutComments, null, false, EMPTY_OPTIONS);
+    private ParsedStatement(
+        StatementType type, Statement statement, Supplier<String> sqlWithoutComments) {
+      this(
+          type,
+          null,
+          statement,
+          sqlWithoutComments,
+          null,
+          Suppliers.ofInstance(false),
+          EMPTY_OPTIONS);
     }
 
     private ParsedStatement(
         StatementType type,
         ClientSideStatementImpl clientSideStatement,
         Statement statement,
-        String sqlWithoutComments,
+        Supplier<String> sqlWithoutComments,
         QueryOptions defaultQueryOptions,
-        boolean returningClause,
+        Supplier<Boolean> returningClause,
         ReadQueryUpdateTransactionOption[] optionsFromHints) {
       Preconditions.checkNotNull(type);
       this.type = type;
@@ -315,7 +328,7 @@ public StatementType getType() {
     /** @return whether the statement has a returning clause or not. */
     @InternalApi
     public boolean hasReturningClause() {
-      return this.returningClause;
+      return this.returningClause.get();
     }
 
     @InternalApi
@@ -413,7 +426,7 @@ Statement mergeQueryOptions(Statement statement, QueryOptions defaultQueryOption
     /** @return the SQL statement with all comments removed from the SQL string. */
     @InternalApi
     public String getSqlWithoutComments() {
-      return sqlWithoutComments;
+      return sqlWithoutComments.get();
     }
 
     ClientSideStatement getClientSideStatement() {
@@ -464,7 +477,7 @@ private static boolean isRecordStatementCacheStats() {
               // We do length*2 because Java uses 2 bytes for each char.
               .weigher(
                   (Weigher<String, ParsedStatement>)
-                      (key, value) -> 2 * key.length() + 2 * value.sqlWithoutComments.length())
+                      (key, value) -> 2 * key.length() + 2 * value.statement.getSql().length())
               .concurrencyLevel(Runtime.getRuntime().availableProcessors());
       if (isRecordStatementCacheStats()) {
         cacheBuilder.recordStats();
@@ -511,28 +524,56 @@ ParsedStatement parse(Statement statement, QueryOptions defaultQueryOptions) {
     return parsedStatement.copy(statement, defaultQueryOptions);
   }
 
-  private ParsedStatement internalParse(Statement statement, QueryOptions defaultQueryOptions) {
-    StatementHintParser statementHintParser =
-        new StatementHintParser(getDialect(), statement.getSql());
+  ParsedStatement internalParse(Statement statement, QueryOptions defaultQueryOptions) {
+    String sql = statement.getSql();
+    StatementHintParser statementHintParser = new StatementHintParser(getDialect(), sql);
     ReadQueryUpdateTransactionOption[] optionsFromHints = EMPTY_OPTIONS;
     if (statementHintParser.hasStatementHints()
         && !statementHintParser.getClientSideStatementHints().isEmpty()) {
       statement =
           statement.toBuilder().replace(statementHintParser.getSqlWithoutClientSideHints()).build();
       optionsFromHints = convertHintsToOptions(statementHintParser.getClientSideStatementHints());
     }
-    String sql = removeCommentsAndTrim(statement.getSql());
-    ClientSideStatementImpl client = parseClientSideStatement(sql);
+    // Create a supplier that will actually remove all comments and hints from the SQL string to be
+    // backwards compatible with anything that really needs the SQL string without comments.
+    Supplier<String> sqlWithoutCommentsSupplier =
+        Suppliers.memoize(() -> removeCommentsAndTrim(sql));
+
+    // Get rid of any spaces/comments at the start of the string.
+    SimpleParser simpleParser = new SimpleParser(getDialect(), sql);
+    simpleParser.skipWhitespaces();
+    // Create a wrapper around the SQL string from the point after the first whitespace.
+    CharBuffer charBuffer = CharBuffer.wrap(sql, simpleParser.getPos(), sql.length());
+    ClientSideStatementImpl client = parseClientSideStatement(charBuffer);
+
     if (client != null) {
-      return ParsedStatement.clientSideStatement(client, statement, sql);
-    } else if (isQuery(sql)) {
-      return ParsedStatement.query(statement, sql, defaultQueryOptions, optionsFromHints);
-    } else if (isUpdateStatement(sql)) {
-      return ParsedStatement.update(statement, sql, checkReturningClause(sql), optionsFromHints);
-    } else if (isDdlStatement(sql)) {
-      return ParsedStatement.ddl(statement, sql);
+      return ParsedStatement.clientSideStatement(client, statement, sqlWithoutCommentsSupplier);
+    } else {
+      Result keywordResult = simpleParser.eatNextKeyword();
+      if (keywordResult.isValid()) {
+        String keyword = keywordResult.getValue().toUpperCase();
+        if (keywordResult.isInParenthesis()) {
+          if (SELECT_STATEMENTS_ALLOWING_PRECEDING_BRACKETS.contains(keyword)) {
+            return ParsedStatement.query(
+                statement, sqlWithoutCommentsSupplier, defaultQueryOptions, optionsFromHints);
+          }
+        } else {
+          if (selectStatements.contains(keyword)) {
+            return ParsedStatement.query(
+                statement, sqlWithoutCommentsSupplier, defaultQueryOptions, optionsFromHints);
+          } else if (dmlStatements.contains(keyword)) {
+            return ParsedStatement.update(
+                statement,
+                sqlWithoutCommentsSupplier,
+                Suppliers.memoize(() -> checkReturningClause(sqlWithoutCommentsSupplier.get())),
+                optionsFromHints);
+          } else if (ddlStatements.contains(keyword)) {
+            return ParsedStatement.ddl(statement, sqlWithoutCommentsSupplier);
+          }
+        }
+      }
     }
-    return ParsedStatement.unknown(statement, sql);
+    return ParsedStatement.unknown(statement, sqlWithoutCommentsSupplier);
   }
 
   /**
@@ -546,7 +587,7 @@ private ParsedStatement internalParse(Statement statement, QueryOptions defaultQ
    *     statement.
    */
   @VisibleForTesting
-  ClientSideStatementImpl parseClientSideStatement(String sql) {
+  ClientSideStatementImpl parseClientSideStatement(CharSequence sql) {
     for (ClientSideStatementImpl css : statements) {
       if (css.matches(sql)) {
         return css;

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ClientSideStatementImpl.java

@@ -193,7 +193,7 @@ public ClientSideStatementType getStatementType() {
     return statementType;
   }
 
-  boolean matches(String statement) {
+  boolean matches(CharSequence statement) {
     Preconditions.checkState(pattern != null, "This statement has not been compiled");
     return pattern.matcher(statement).matches();
   }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SimpleParser.java

@@ -28,21 +28,31 @@
 
 /** A very simple token-based parser for extracting relevant information from SQL strings. */
 class SimpleParser {
+  static final char DOLLAR = '$';
+  static final char BACKSLASH = '\\';
+
   /**
    * An immutable result from a parse action indicating whether the parse action was successful, and
    * if so, what the value was.
    */
   static class Result {
-    static final Result NOT_FOUND = new Result(null);
+    static final Result NOT_FOUND = new Result(null, false);
 
     static Result found(String value) {
-      return new Result(Preconditions.checkNotNull(value));
+      return new Result(Preconditions.checkNotNull(value), false);
+    }
+
+    static Result found(String value, boolean inParenthesis) {
+      return new Result(Preconditions.checkNotNull(value), inParenthesis);
     }
 
     private final String value;
 
-    private Result(String value) {
+    private final boolean inParenthesis;
+
+    private Result(String value, boolean inParenthesis) {
       this.value = value;
+      this.inParenthesis = inParenthesis;
     }
 
     @Override
@@ -55,7 +65,8 @@ public boolean equals(Object o) {
       if (!(o instanceof Result)) {
         return false;
       }
-      return Objects.equals(this.value, ((Result) o).value);
+      return Objects.equals(this.value, ((Result) o).value)
+          && Objects.equals(this.inParenthesis, ((Result) o).inParenthesis);
     }
 
     @Override
@@ -73,6 +84,10 @@ boolean isValid() {
     String getValue() {
       return this.value;
     }
+
+    boolean isInParenthesis() {
+      return this.inParenthesis;
+    }
   }
 
   // TODO: Replace this with a direct reference to the dialect, and move the isXYZSupported methods
@@ -81,6 +96,8 @@ String getValue() {
 
   private final String sql;
 
+  private final int length;
+
   private final boolean treatHintCommentsAsTokens;
 
   private int pos;
@@ -100,6 +117,7 @@ String getValue() {
         !(treatHintCommentsAsTokens && dialect != Dialect.POSTGRESQL),
         "treatHintCommentsAsTokens can only be enabled for PostgreSQL");
     this.sql = sql;
+    this.length = sql.length();
     this.pos = pos;
     this.statementParser = AbstractStatementParser.getInstance(dialect);
     this.treatHintCommentsAsTokens = treatHintCommentsAsTokens;
@@ -117,12 +135,50 @@ int getPos() {
     return this.pos;
   }
 
+  void skipHint() {
+    // We don't need to do anything special for PostgreSQL, as hints in PostgreSQL are inside
+    // comments.
+    if (getDialect() == Dialect.GOOGLE_STANDARD_SQL && eatTokens('@', '{')) {
+      while (pos < length && !eatToken('}')) {
+        pos += statementParser.skip(sql, pos, /*result=*/ null);
+      }
+    }
+  }
+
+  Result eatNextKeyword() {
+    skipHint();
+    boolean inParenthesis = false;
+    while (pos < length && eatToken('(')) {
+      inParenthesis = true;
+    }
+    return eatKeyword(inParenthesis);
+  }
+
   /** Returns true if this parser has more tokens. Advances the position to the first next token. */
   boolean hasMoreTokens() {
     skipWhitespaces();
     return pos < sql.length();
   }
 
+  /** Eats and returns the keyword at the current position. */
+  Result eatKeyword() {
+    return eatKeyword(false);
+  }
+
+  Result eatKeyword(boolean inParenthesis) {
+    if (!hasMoreTokens()) {
+      return Result.NOT_FOUND;
+    }
+    if (!Character.isLetter(sql.charAt(pos))) {
+      return Result.NOT_FOUND;
+    }
+    int startPos = pos;
+    while (pos < length && Character.isLetter(sql.charAt(pos))) {
+      pos++;
+    }
+    return Result.found(sql.substring(startPos, pos), inParenthesis);
+  }
+
   /**
    * Eats and returns the identifier at the current position. This implementation does not support
    * quoted identifiers.