fix: add support for WITH clauses (#42)

googleapis · Jan 28, 2020 · 7f4bea4 · 7f4bea4
1 parent b368b84
commit 7f4bea4
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 1 deletion.
diff --git 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<String> ddlStatements = ImmutableSet.of("CREATE", "DROP", "ALTER");
-  private final Set<String> selectStatements = ImmutableSet.of("SELECT");
+  private final Set<String> selectStatements = ImmutableSet.of("SELECT", "WITH");
   private final Set<String> dmlStatements = ImmutableSet.of("INSERT", "UPDATE", "DELETE");
   private final Set<ClientSideStatementImpl> 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
@@ -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
@@ -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;