From 7f4bea43c42df68258f776944ea744069a1a218e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Tue, 28 Jan 2020 08:30:59 +0100
Subject: [PATCH] fix: add support for WITH clauses (#42)

---
 .../cloud/spanner/jdbc/StatementParser.java   |  2 +-
 .../spanner/jdbc/StatementParserTest.java     | 28 +++++++++++++++++++
 .../spanner/jdbc/ITSqlMusicScriptTest.sql     |  5 ++++
 3 files changed, 34 insertions(+), 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
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<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
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;