implement support for Common Table Expressions

nukoyluoglu · nukoyluoglu · commit 0d916dd4e592 · 2020-07-23T15:45:43.000+03:00
diff --git a/src/sqlancer/postgres/PostgresExpectedValueVisitor.java b/src/sqlancer/postgres/PostgresExpectedValueVisitor.java
@@ -16,6 +16,7 @@
 import sqlancer.postgres.ast.PostgresPrefixOperation;
 import sqlancer.postgres.ast.PostgresSelect;
 import sqlancer.postgres.ast.PostgresSelect.PostgresFromTable;
+import sqlancer.postgres.ast.PostgresSelect.PostgresCTE;
 import sqlancer.postgres.ast.PostgresSimilarTo;
 
 public final class PostgresExpectedValueVisitor implements PostgresVisitor {
@@ -155,4 +156,9 @@ public void visit(PostgresFromTable from) {
         print(from);
     }
 
+    @Override
+    public void visit(PostgresCTE cte) {
+        print(cte);
+    }
+
 }
diff --git a/src/sqlancer/postgres/PostgresToStringVisitor.java b/src/sqlancer/postgres/PostgresToStringVisitor.java
@@ -22,6 +22,7 @@
 import sqlancer.postgres.ast.PostgresPrefixOperation;
 import sqlancer.postgres.ast.PostgresSelect;
 import sqlancer.postgres.ast.PostgresSelect.PostgresFromTable;
+import sqlancer.postgres.ast.PostgresSelect.PostgresCTE;
 import sqlancer.postgres.ast.PostgresSimilarTo;
 import sqlancer.visitor.ToStringVisitor;
 
@@ -75,6 +76,14 @@ public void visit(PostgresFromTable from) {
         }
     }
 
+    @Override
+    public void visit(PostgresCTE cte) {
+        sb.append("(");
+        visit(cte.getSelect());
+        sb.append(") AS ");
+        sb.append(cte.getName());
+    }
+
     @Override
     public void visit(PostgresSelect s) {
         sb.append("SELECT ");
@@ -126,7 +135,11 @@ public void visit(PostgresSelect s) {
                 throw new AssertionError(j.getType());
             }
             sb.append(" ");
-            sb.append(j.getTable().getName());
+            if (j.joinCTE()) {
+                visit(j.getCTE());
+            } else {
+                sb.append(j.getTable().getName());
+            }
             if (j.getType() != PostgresJoinType.CROSS) {
                 sb.append(" ON ");
                 visit(j.getOnClause());
diff --git a/src/sqlancer/postgres/PostgresVisitor.java b/src/sqlancer/postgres/PostgresVisitor.java
@@ -20,6 +20,7 @@
 import sqlancer.postgres.ast.PostgresPrefixOperation;
 import sqlancer.postgres.ast.PostgresSelect;
 import sqlancer.postgres.ast.PostgresSelect.PostgresFromTable;
+import sqlancer.postgres.ast.PostgresSelect.PostgresCTE;
 import sqlancer.postgres.ast.PostgresSimilarTo;
 import sqlancer.postgres.gen.PostgresExpressionGenerator;
 
@@ -57,6 +58,8 @@ public interface PostgresVisitor {
 
     void visit(PostgresFromTable from);
 
+    void visit(PostgresCTE cte);
+
     default void visit(PostgresExpression expression) {
         if (expression instanceof PostgresConstant) {
             visit((PostgresConstant) expression);
@@ -90,6 +93,8 @@ default void visit(PostgresExpression expression) {
             visit((PostgresCollate) expression);
         } else if (expression instanceof PostgresFromTable) {
             visit((PostgresFromTable) expression);
+        } else if (expression instanceof PostgresCTE) {
+            visit((PostgresCTE) expression);
         } else {
             throw new AssertionError(expression);
         }
diff --git a/src/sqlancer/postgres/ast/PostgresJoin.java b/src/sqlancer/postgres/ast/PostgresJoin.java
@@ -3,6 +3,7 @@
 import sqlancer.Randomly;
 import sqlancer.postgres.PostgresSchema.PostgresDataType;
 import sqlancer.postgres.PostgresSchema.PostgresTable;
+import sqlancer.postgres.ast.PostgresSelect.PostgresCTE;
 
 public class PostgresJoin implements PostgresExpression {
 
@@ -18,17 +19,33 @@ public static PostgresJoinType getRandom() {
     private final PostgresTable table;
     private final PostgresExpression onClause;
     private final PostgresJoinType type;
+    private PostgresCTE CTE = null;
 
     public PostgresJoin(PostgresTable table, PostgresExpression onClause, PostgresJoinType type) {
         this.table = table;
         this.onClause = onClause;
         this.type = type;
     }
 
+    public PostgresJoin(PostgresCTE CTE, PostgresExpression onClause, PostgresJoinType type) {
+        this.CTE = CTE;
+        this.onClause = onClause;
+        this.type = type;
+        this.table = null;
+    }
+
     public PostgresTable getTable() {
         return table;
     }
 
+    public PostgresCTE getCTE() {
+        return CTE;
+    }
+
+    public boolean joinCTE() {
+        return (CTE != null);
+    }
+
     public PostgresExpression getOnClause() {
         return onClause;
     }
diff --git a/src/sqlancer/postgres/ast/PostgresSelect.java b/src/sqlancer/postgres/ast/PostgresSelect.java
@@ -56,6 +56,29 @@ public PostgresDataType getExpressionType() {
         }
     }
 
+    public static class PostgresCTE implements PostgresExpression {
+        private final PostgresSelect s;
+        private final String name;
+
+        public PostgresCTE(PostgresSelect s, String name) {
+            this.s = s;
+            this.name = name;
+        }
+
+        public PostgresSelect getSelect() {
+            return s;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        @Override
+        public PostgresDataType getExpressionType() {
+            return null;
+        }
+    }
+
     public enum SelectType {
         DISTINCT, ALL;
 
