diff --git a/src/main/java/com/google/cloud/spanner/jdbc/StatementParser.java b/src/main/java/com/google/cloud/spanner/jdbc/StatementParser.java index 9ed1ecd2..e43a0974 100644 --- a/src/main/java/com/google/cloud/spanner/jdbc/StatementParser.java +++ b/src/main/java/com/google/cloud/spanner/jdbc/StatementParser.java @@ -161,7 +161,7 @@ ClientSideStatement getClientSideStatement() { } private final Set ddlStatements = ImmutableSet.of("CREATE", "DROP", "ALTER"); - private final Set selectStatements = ImmutableSet.of("SELECT"); + private final Set selectStatements = ImmutableSet.of("SELECT", "WITH"); private final Set dmlStatements = ImmutableSet.of("INSERT", "UPDATE", "DELETE"); private final Set statements; diff --git a/src/test/java/com/google/cloud/spanner/jdbc/StatementParserTest.java b/src/test/java/com/google/cloud/spanner/jdbc/StatementParserTest.java index 6da65c15..1344a0f3 100644 --- a/src/test/java/com/google/cloud/spanner/jdbc/StatementParserTest.java +++ b/src/test/java/com/google/cloud/spanner/jdbc/StatementParserTest.java @@ -266,6 +266,27 @@ public void testIsQuery() { assertTrue(parser.isQuery("select * from foo")); assertFalse(parser.isQuery("INSERT INTO FOO (ID, NAME) SELECT ID+1, NAME FROM FOO")); + assertTrue( + parser.isQuery( + "WITH subQ1 AS (SELECT SchoolID FROM Roster),\n" + + " subQ2 AS (SELECT OpponentID FROM PlayerStats)\n" + + "SELECT * FROM subQ1\n" + + "UNION ALL\n" + + "SELECT * FROM subQ2")); + assertTrue( + parser.isQuery( + "with subQ1 AS (SELECT SchoolID FROM Roster),\n" + + " subQ2 AS (SELECT OpponentID FROM PlayerStats)\n" + + "select * FROM subQ1\n" + + "UNION ALL\n" + + "SELECT * FROM subQ2")); + assertTrue( + parser + .parse( + Statement.of( + "-- this is a comment\nwith foo as (select * from bar)\nselect * from foo")) + .isQuery()); + assertTrue(parser.parse(Statement.of("-- this is a comment\nselect * from foo")).isQuery()); assertTrue( parser @@ -310,6 +331,13 @@ public void testQueryHints() { assertTrue( parser.isQuery( "@{JOIN_METHOD=HASH_JOIN}\n /* Multi line comment\n with more comments\n */SELECT * FROM PersonsTable")); + assertTrue( + parser.isQuery( + "@{JOIN_METHOD=HASH_JOIN} WITH 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")); diff --git a/src/test/resources/com/google/cloud/spanner/jdbc/ITSqlMusicScriptTest.sql b/src/test/resources/com/google/cloud/spanner/jdbc/ITSqlMusicScriptTest.sql index 93218da9..243cdda3 100644 --- a/src/test/resources/com/google/cloud/spanner/jdbc/ITSqlMusicScriptTest.sql +++ b/src/test/resources/com/google/cloud/spanner/jdbc/ITSqlMusicScriptTest.sql @@ -668,3 +668,8 @@ WHERE SingerId=2; SELECT COUNT(*) AS ACTUAL, 0 AS EXPECTED FROM Songs WHERE SingerId=2; + +@EXPECT RESULT_SET +WITH Song2 AS (SELECT * FROM Songs WHERE SingerId=2) +SELECT COUNT(*) AS ACTUAL, 0 AS EXPECTED +FROM Song2;