From ef41a6e503f218c00c16914aa9c1433d9b26db13 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Wed, 22 Apr 2020 12:18:14 +0200 Subject: [PATCH] fix: StatementParser did not accept multiple query hints (#170) Fixes #163 --- .../spanner/connection/StatementParser.java | 38 ++++++++++++------- .../connection/StatementParserTest.java | 18 +++++++++ 2 files changed, 43 insertions(+), 13 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementParser.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementParser.java index 60bbac813d..335eb12534 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementParser.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementParser.java @@ -212,9 +212,9 @@ ClientSideStatement getClientSideStatement() { } } - private final Set ddlStatements = ImmutableSet.of("CREATE", "DROP", "ALTER"); - private final Set selectStatements = ImmutableSet.of("SELECT", "WITH"); - private final Set dmlStatements = ImmutableSet.of("INSERT", "UPDATE", "DELETE"); + private static final Set ddlStatements = ImmutableSet.of("CREATE", "DROP", "ALTER"); + private static final Set selectStatements = ImmutableSet.of("SELECT", "WITH"); + private static final Set dmlStatements = ImmutableSet.of("INSERT", "UPDATE", "DELETE"); private final Set 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; } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java index 63c8c5111f..edde69a05b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java +++ b/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