fix: StatementParser did not accept multiple query hints (#170)

Fixes #163

Author: olavloite

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

@@ -212,9 +212,9 @@ ClientSideStatement getClientSideStatement() {
  }
  }
 
- private final Set<String> ddlStatements = ImmutableSet.of("CREATE", "DROP", "ALTER");
- private final Set<String> selectStatements = ImmutableSet.of("SELECT", "WITH");
- private final Set<String> dmlStatements = ImmutableSet.of("INSERT", "UPDATE", "DELETE");
+ private static final Set<String> ddlStatements = ImmutableSet.of("CREATE", "DROP", "ALTER");
+ private static final Set<String> selectStatements = ImmutableSet.of("SELECT", "WITH");
+ private static final Set<String> dmlStatements = ImmutableSet.of("INSERT", "UPDATE", "DELETE");
  private final Set<ClientSideStatementImpl> statements;
 
  /** Private constructor for singleton instance. */
@@ -445,19 +445,31 @@ public static String removeCommentsAndTrim(String sql) {
 
  /** Removes any statement hints at the beginning of the statement. */
  static String removeStatementHint(String sql) {
- // Valid statement hints at the beginning of a SQL statement can only contain a fixed set of
+ // Valid statement hints at the beginning of a query statement can only contain a fixed set of
  // possible values. Although it is possible to add a @{FORCE_INDEX=...} as a statement hint, the
  // only allowed value is _BASE_TABLE. This means that we can safely assume that the statement
- // hint will not contain any special characters, for example a closing curly brace, and
- // that we can keep the check simple by just searching for the first occurrence of a closing
- // curly brace at the end of the statement hint.
+ // hint will not contain any special characters, for example a closing curly brace or one of the
+ // keywords SELECT, UPDATE, DELETE, WITH, and that we can keep the check simple by just
+ // searching for the first occurrence of a keyword that should be preceded by a closing curly
+ // brace at the end of the statement hint.
  int startStatementHintIndex = sql.indexOf('{');
- int endStatementHintIndex = sql.indexOf('}');
- if (startStatementHintIndex == -1 || startStatementHintIndex > endStatementHintIndex) {
- // Looks like an invalid statement hint. Just ignore at this point and let the caller handle
- // the invalid query.
- return sql;
+ // Statement hints are only allowed for queries.
+ int startQueryIndex = -1;
+ String upperCaseSql = sql.toUpperCase();
+ for (String keyword : selectStatements) {
+ startQueryIndex = upperCaseSql.indexOf(keyword);
+ if (startQueryIndex > -1) break;
  }
- return removeCommentsAndTrim(sql.substring(endStatementHintIndex + 1));
+ if (startQueryIndex > -1) {
+ int endStatementHintIndex = sql.substring(0, startQueryIndex).lastIndexOf('}');
+ if (startStatementHintIndex == -1 || startStatementHintIndex > endStatementHintIndex) {
+ // Looks like an invalid statement hint. Just ignore at this point and let the caller handle
+ // the invalid query.
+ return sql;
+ }
+ return removeCommentsAndTrim(sql.substring(endStatementHintIndex + 1));
+ }
+ // Seems invalid, just return the original statement.
+ return sql;
  }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java

@@ -339,10 +339,28 @@ public void testQueryHints() {
  + "UNION ALL\n"
  + "SELECT * FROM subQ2"));
 
+ // Multiple query hints.
+ assertTrue(
+ StatementParser.INSTANCE.isQuery(
+ "@{FORCE_INDEX=index_name} @{JOIN_METHOD=HASH_JOIN} SELECT * FROM tbl"));
+ assertTrue(
+ StatementParser.INSTANCE.isQuery(
+ "@{FORCE_INDEX=index_name} @{JOIN_METHOD=HASH_JOIN} Select * FROM tbl"));
+ assertTrue(
+ StatementParser.INSTANCE.isQuery(
+ "@{FORCE_INDEX=index_name}\n@{JOIN_METHOD=HASH_JOIN}\nWITH subQ1 AS (SELECT SchoolID FROM Roster),\n"
+ + " subQ2 AS (SELECT OpponentID FROM PlayerStats)\n"
+ + "SELECT * FROM subQ1\n"
+ + "UNION ALL\n"
+ + "SELECT * FROM subQ2"));
+
  // Invalid query hints.
  assertFalse(parser.isQuery("@{JOIN_METHOD=HASH_JOIN SELECT * FROM PersonsTable"));
  assertFalse(parser.isQuery("@JOIN_METHOD=HASH_JOIN} SELECT * FROM PersonsTable"));
  assertFalse(parser.isQuery("@JOIN_METHOD=HASH_JOIN SELECT * FROM PersonsTable"));
+ assertFalse(
+ StatementParser.INSTANCE.isQuery(
+ "@{FORCE_INDEX=index_name} @{JOIN_METHOD=HASH_JOIN} UPDATE tbl set FOO=1 WHERE ID=2"));
  }
 
  @Test