Issue#576 Add support for insert,update,merge,delete and create table statemets in WITH statement

h2/src/main/org/h2/command/Parser.java

@@ -482,8 +482,8 @@ public class Parser {
                 break;
             case 'w':
             case 'W':
-                if (isToken("WITH")) {
-                    c = parseSelect();
+                if (readIf("WITH")) {
+                    c = parseWithStatementOrQuery();
                 }
                 break;
             case ';':
@@ -1758,6 +1758,21 @@ public class Parser {
         return command;
     }
 
+    private Prepared parseWithStatementOrQuery() {
+        int paramIndex = parameters.size();
+        Prepared command = parseWith();
+        ArrayList<Parameter> params = New.arrayList();
+        for (int i = paramIndex, size = parameters.size(); i < size; i++) {
+            params.add(parameters.get(i));
+        }
+        command.setParameterList(params);
+        if(command instanceof Query){
+            Query query = (Query) command;
+            query.init();
+        }
+        return command;
+    }
+
     private Query parseSelectUnion() {
         int start = lastParseIndex;
         Query command = parseSelectSub();
@@ -1940,7 +1955,14 @@ public class Parser {
             return command;
         }
         if (readIf("WITH")) {
-            Query query = parseWith();
+            Query query = null;
+            try {
+                query = (Query) parseWith();
+            }
+            catch(ClassCastException e){
+                throw DbException.get(ErrorCode.SYNTAX_ERROR_1,
+                        "WITH statment supports only SELECT (query) in this context");
+            }
             // recursive can not be lazy
             query.setNeverLazy(true);
             return query;
@@ -4890,15 +4912,52 @@ public class Parser {
         return command;
     }
 
-    private Query parseWith() {
+    private Prepared parseWith() {
         List<TableView> viewsCreated = new ArrayList<TableView>();
         readIf("RECURSIVE");
         do {
             viewsCreated.add(parseSingleCommonTableExpression());
         } while (readIf(","));
 
-        Query q = parseSelectUnion();
-        q.setPrepareAlways(true);
+        Prepared p = null;
+
+        if(isToken("SELECT")) {
+            Query query = parseSelectUnion();
+            query.setPrepareAlways(true);
+            query.setNeverLazy(true);
+            p = query;
+        }
+        else if(readIf("INSERT")) {
+            p = parseInsert();
+            p.setPrepareAlways(true);
+        }
+        else if(readIf("UPDATE")) {
+            p = parseUpdate();
+            p.setPrepareAlways(true);
+        }
+        else if(readIf("MERGE")) {
+            p = parseMerge();
+            p.setPrepareAlways(true);
+        }
+        else if(readIf("DELETE")) {
+            p = parseDelete();
+            p.setPrepareAlways(true);
+        }
+        else if(readIf("CREATE")) {
+            if (!isToken("TABLE")){
+                throw DbException.get(ErrorCode.SYNTAX_ERROR_1,
+                        "WITH statment supports only SELECT, CREATE TABLE, INSERT, UPDATE, MERGE or DELETE statements");
+
+            }
+            p = parseCreate();
+            p.setPrepareAlways(true);
+        }
+        else {
+            throw DbException.get(ErrorCode.SYNTAX_ERROR_1,
+                    "WITH statment supports only SELECT, CREATE TABLE, INSERT, UPDATE, MERGE or DELETE statements");
+
+        }
+
                 List<Runnable> cleanupCallbacks = new ArrayList<Runnable>();
 
         // clean up temp views starting with last to first (in case of dependencies)
@@ -4918,8 +4977,8 @@ public class Parser {
                 }
             });
         }
-        q.setCleanupCallbacks(cleanupCallbacks);
-        return q;
+        p.setCleanupCallbacks(cleanupCallbacks);
+        return p;
     }
 
     private TableView parseSingleCommonTableExpression() {

h2/src/test/org/h2/test/db/TestGeneralCommonTableQueries.java

@@ -29,14 +29,20 @@ public class TestGeneralCommonTableQueries extends TestBase {
 
     @Override
     public void test() throws Exception {
-        testSimple();
+        testSimpleSelect();
         testImpliedColumnNames();
         testChainedQuery();
         testParameterizedQuery();
         testNumberedParameterizedQuery();
+
+        testInsert();
+        testUpdate();
+        testDelete();
+        testMerge();
+        testCreateTable();
     }
 
-    private void testSimple() throws Exception {
+    private void testSimpleSelect() throws Exception {
         deleteDb("commonTableExpressionQueries");
         Connection conn = getConnection("commonTableExpressionQueries");
         Statement stat;
@@ -227,4 +233,157 @@ public class TestGeneralCommonTableQueries extends TestBase {
         conn.close();
         deleteDb("commonTableExpressionQueries");
     }
+
+    private void testInsert() throws Exception {
+        deleteDb("commonTableExpressionQueries");
+        Connection conn = getConnection("commonTableExpressionQueries");
+        Statement stat;
+        PreparedStatement prep;
+        ResultSet rs;
+        int rowCount;
+
+        stat = conn.createStatement();
+        stat.execute("CREATE TABLE T1 ( ID INT IDENTITY,  X INT NULL, Y VARCHAR(100) NULL )");
+
+        prep = conn.prepareStatement("WITH v1 AS ("
+                + "     SELECT R.X, 'X1' AS Y FROM SYSTEM_RANGE(?1,?2) R"
+                + ")"
+                + "INSERT INTO T1 (X,Y) SELECT v1.X, v1.Y FROM v1");
+        prep.setInt(1, 1);
+        prep.setInt(2, 2);
+        rowCount = prep.executeUpdate();
+
+        assertEquals(2, rowCount);
+
+        rs = stat.executeQuery("SELECT ID, X,Y FROM T1");
+
+        for (int n : new int[]{1, 2}) {
+            assertTrue(rs.next());
+            assertTrue(rs.getInt(1) != 0);
+            assertEquals(n, rs.getInt(2));
+            assertEquals("X1", rs.getString(3));
+        }
+        conn.close();
+        deleteDb("commonTableExpressionQueries");
+    }
+
+    private void testUpdate() throws Exception {
+        deleteDb("commonTableExpressionQueries");
+        Connection conn = getConnection("commonTableExpressionQueries");
+        Statement stat;
+        PreparedStatement prep;
+        ResultSet rs;
+        int rowCount;
+
+        stat = conn.createStatement();
+        stat.execute("CREATE TABLE IF NOT EXISTS T1 AS SELECT R.X AS ID, R.X, 'X1' AS Y FROM SYSTEM_RANGE(1,2) R");
+
+        prep = conn.prepareStatement("WITH v1 AS ("
+                +"     SELECT R.X, 'X1' AS Y FROM SYSTEM_RANGE(?1,?2) R"
+                +")"
+                +"UPDATE T1 SET Y = 'Y1' WHERE X IN ( SELECT v1.X FROM v1 )");
+        prep.setInt(1, 1);
+        prep.setInt(2, 2);
+        rowCount = prep.executeUpdate();
+
+        assertEquals(2,rowCount);
+
+        rs = stat.executeQuery("SELECT ID, X,Y FROM T1");
+
+        for (int n : new int[] { 1, 2 }) {
+            assertTrue(rs.next());
+            assertTrue(rs.getInt(1)!=0);
+            assertEquals(n, rs.getInt(2));
+            assertEquals("Y1", rs.getString(3));
+        }
+        conn.close();
+        deleteDb("commonTableExpressionQueries");
+    }
+
+    private void testDelete() throws Exception {
+        deleteDb("commonTableExpressionQueries");
+        Connection conn = getConnection("commonTableExpressionQueries");
+        Statement stat;
+        PreparedStatement prep;
+        ResultSet rs;
+        int rowCount;
+
+        stat = conn.createStatement();
+        stat.execute("CREATE TABLE IF NOT EXISTS T1 AS SELECT R.X AS ID, R.X, 'X1' AS Y FROM SYSTEM_RANGE(1,2) R");
+
+        prep = conn.prepareStatement("WITH v1 AS ("
+                +"     SELECT R.X, 'X1' AS Y FROM SYSTEM_RANGE(1,2) R"
+                +")"
+                +"DELETE FROM T1 WHERE X IN ( SELECT v1.X FROM v1 )");
+        rowCount = prep.executeUpdate();
+
+        assertEquals(2,rowCount);
+
+        rs = stat.executeQuery("SELECT ID, X,Y FROM T1");
+
+        assertFalse(rs.next());
+
+        conn.close();
+        deleteDb("commonTableExpressionQueries");
+    }
+
+    private void testMerge() throws Exception {
+        deleteDb("commonTableExpressionQueries");
+        Connection conn = getConnection("commonTableExpressionQueries");
+        Statement stat;
+        PreparedStatement prep;
+        ResultSet rs;
+        int rowCount;
+
+        stat = conn.createStatement();
+        stat.execute("CREATE TABLE IF NOT EXISTS T1 AS SELECT R.X AS ID, R.X, 'X1' AS Y FROM SYSTEM_RANGE(1,2) R");
+
+        prep = conn.prepareStatement("WITH v1 AS ("
+                +"     SELECT R.X, 'X1' AS Y FROM SYSTEM_RANGE(1,3) R"
+                +")"
+                +"MERGE INTO T1 KEY(ID) SELECT v1.X AS ID, v1.X, v1.Y FROM v1");
+        rowCount = prep.executeUpdate();
+
+        assertEquals(3,rowCount);
+
+        rs = stat.executeQuery("SELECT ID, X,Y FROM T1");
+
+        for (int n : new int[] { 1, 2, 3 }) {
+            assertTrue(rs.next());
+            assertTrue(rs.getInt(1)!=0);
+            assertEquals(n, rs.getInt(2));
+            assertEquals("X1", rs.getString(3));
+        }
+        conn.close();
+        deleteDb("commonTableExpressionQueries");
+    }
+
+    private void testCreateTable() throws Exception {
+        deleteDb("commonTableExpressionQueries");
+        Connection conn = getConnection("commonTableExpressionQueries");
+        Statement stat;
+        PreparedStatement prep;
+        ResultSet rs;
+        boolean success;
+
+        stat = conn.createStatement();
+        prep = conn.prepareStatement("WITH v1 AS ("
+                +"     SELECT R.X, 'X1' AS Y FROM SYSTEM_RANGE(1,3) R"
+                +")"
+                +"CREATE TABLE IF NOT EXISTS T1 AS SELECT v1.X AS ID, v1.X, v1.Y FROM v1");
+        success = prep.execute();
+
+        assertEquals(false,success);
+
+        rs = stat.executeQuery("SELECT ID, X,Y FROM T1");
+
+        for (int n : new int[] { 1, 2, 3 }) {
+            assertTrue(rs.next());
+            assertTrue(rs.getInt(1)!=0);
+            assertEquals(n, rs.getInt(2));
+            assertEquals("X1", rs.getString(3));
+        }
+        conn.close();
+        deleteDb("commonTableExpressionQueries");
+    }
 }
\ No newline at end of file