sqlancer · yaojiejia · Jul 2, 2025
diff --git a/src/sqlancer/postgres/PostgresExpectedValueVisitor.java b/src/sqlancer/postgres/PostgresExpectedValueVisitor.java
@@ -9,6 +9,7 @@
 import sqlancer.postgres.ast.PostgresColumnReference;
 import sqlancer.postgres.ast.PostgresColumnValue;
 import sqlancer.postgres.ast.PostgresConstant;
+import sqlancer.postgres.ast.PostgresCTE;
 import sqlancer.postgres.ast.PostgresExpression;
 import sqlancer.postgres.ast.PostgresFunction;
 import sqlancer.postgres.ast.PostgresInOperation;
@@ -24,6 +25,7 @@
 import sqlancer.postgres.ast.PostgresSimilarTo;
 import sqlancer.postgres.ast.PostgresTableReference;
 import sqlancer.postgres.ast.PostgresWindowFunction;
+import sqlancer.postgres.ast.PostgresWithClause;
 
 public final class PostgresExpectedValueVisitor implements PostgresVisitor {
 
@@ -195,4 +197,18 @@ public void visit(PostgresLikeOperation op) {
         visit(op.getRight());
     }
 
+    @Override
+    public void visit(PostgresCTE cte) {
+        print(cte);
+        visit(cte.getSubquery());
+    }
+
+    @Override
+    public void visit(PostgresWithClause withClause) {
+        print(withClause);
+        for (PostgresCTE cte : withClause.getCteList()) {
+            visit(cte);
+        }
+    }
+
 }
diff --git a/src/sqlancer/postgres/PostgresToStringVisitor.java b/src/sqlancer/postgres/PostgresToStringVisitor.java
@@ -15,6 +15,7 @@
 import sqlancer.postgres.ast.PostgresColumnReference;
 import sqlancer.postgres.ast.PostgresColumnValue;
 import sqlancer.postgres.ast.PostgresConstant;
+import sqlancer.postgres.ast.PostgresCTE;
 import sqlancer.postgres.ast.PostgresExpression;
 import sqlancer.postgres.ast.PostgresFunction;
 import sqlancer.postgres.ast.PostgresInOperation;
@@ -34,6 +35,7 @@
 import sqlancer.postgres.ast.PostgresWindowFunction;
 import sqlancer.postgres.ast.PostgresWindowFunction.WindowFrame;
 import sqlancer.postgres.ast.PostgresWindowFunction.WindowSpecification;
+import sqlancer.postgres.ast.PostgresWithClause;
 
 public final class PostgresToStringVisitor extends ToStringVisitor<PostgresExpression> implements PostgresVisitor {
 
@@ -105,6 +107,10 @@ public void visit(PostgresTableReference ref) {
 
     @Override
     public void visit(PostgresSelect s) {
+        if (s.getWithClause() != null) {
+            visit(s.getWithClause());
+            sb.append(" ");
+        }
         sb.append("SELECT ");
         switch (s.getSelectOption()) {
         case DISTINCT:
@@ -399,4 +405,24 @@ public void visit(PostgresWindowFunction windowFunction) {
 
         sb.append(")");
     }
+
+    @Override
+    public void visit(PostgresCTE cte) {
+        sb.append(cte.getName());
+        sb.append(" AS (");
+        visit(cte.getSubquery());
+        sb.append(")");
+    }
+
+    @Override
+    public void visit(PostgresWithClause withClause) {
+        sb.append("WITH ");
+        List<PostgresCTE> cteList = withClause.getCteList();
+        for (int i = 0; i < cteList.size(); i++) {
+            if (i > 0) {
+                sb.append(", ");
+            }
+            visit(cteList.get(i));
+        }
+    }
 }
diff --git a/src/sqlancer/postgres/PostgresVisitor.java b/src/sqlancer/postgres/PostgresVisitor.java
@@ -12,6 +12,7 @@
 import sqlancer.postgres.ast.PostgresColumnReference;
 import sqlancer.postgres.ast.PostgresColumnValue;
 import sqlancer.postgres.ast.PostgresConstant;
+import sqlancer.postgres.ast.PostgresCTE;
 import sqlancer.postgres.ast.PostgresExpression;
 import sqlancer.postgres.ast.PostgresFunction;
 import sqlancer.postgres.ast.PostgresInOperation;
@@ -27,6 +28,7 @@
 import sqlancer.postgres.ast.PostgresSimilarTo;
 import sqlancer.postgres.ast.PostgresTableReference;
 import sqlancer.postgres.ast.PostgresWindowFunction;
+import sqlancer.postgres.ast.PostgresWithClause;
 import sqlancer.postgres.gen.PostgresExpressionGenerator;
 
 public interface PostgresVisitor {
@@ -75,6 +77,10 @@ public interface PostgresVisitor {
 
     void visit(PostgresWindowFunction windowFunction);
 
+    void visit(PostgresCTE cte);
+
+    void visit(PostgresWithClause withClause);
+
     default void visit(PostgresExpression expression) {
         if (expression instanceof PostgresConstant) {
             visit((PostgresConstant) expression);
@@ -118,6 +124,10 @@ default void visit(PostgresExpression expression) {
             visit((PostgresTableReference) expression);
         } else if (expression instanceof PostgresWindowFunction) {
             visit((PostgresWindowFunction) expression);
+        } else if (expression instanceof PostgresCTE) {
+            visit((PostgresCTE) expression);
+        } else if (expression instanceof PostgresWithClause) {
+            visit((PostgresWithClause) expression);
         } else {
             throw new AssertionError(expression);
         }

diff --git a/src/sqlancer/postgres/ast/PostgresCTE.java b/src/sqlancer/postgres/ast/PostgresCTE.java
@@ -0,0 +1,65 @@
+package sqlancer.postgres.ast;
+
+import sqlancer.postgres.PostgresSchema.PostgresDataType;
+
+public class PostgresCTE implements PostgresExpression {
+    private final String name;
+    private final PostgresSelect subquery;
+
+    public PostgresCTE(String name, PostgresSelect subquery) {
+        if (name == null || name.trim().isEmpty()) {
+            throw new IllegalArgumentException("CTE name cannot be null or empty");
+        }
+        if (subquery == null) {
+            throw new IllegalArgumentException("CTE subquery cannot be null");
+        }
+        this.name = name.trim();
+        this.subquery = subquery;
+    }
+
+    public String getName() {
+        return name;
+    }
+
+    public PostgresSelect getSubquery() {
+        return subquery;
+    }
+
+    @Override
+    public PostgresDataType getExpressionType() {
+        return null;
+    }
+
+    public String asString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append(name);
+        sb.append(" AS (");
+        sb.append(subquery.asString());
+        sb.append(")");
+        return sb.toString();
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+
+    public static class Builder {
+        private String name;
+        private PostgresSelect subquery;
+
+        public Builder name(String name) {
+            this.name = name;
+            return this;
+        }
+
+        public Builder subquery(PostgresSelect subquery) {
+            this.subquery = subquery;
+            return this;
+        }
+
+        public PostgresCTE build() {
+            return new PostgresCTE(name, subquery);
+        }
+    }
+} 
diff --git a/src/sqlancer/postgres/ast/PostgresSelect.java b/src/sqlancer/postgres/ast/PostgresSelect.java
@@ -24,6 +24,7 @@ public class PostgresSelect extends SelectBase<PostgresExpression>
     private ForClause forClause;
     private List<PostgresExpression> windowFunctions = new ArrayList<>();
     private final Map<String, WindowDefinition> windowDefinitions = new HashMap<>();
+    private PostgresWithClause withClause;
 
     public enum ForClause {
         UPDATE("UPDATE"), NO_KEY_UPDATE("NO KEY UPDATE"), SHARE("SHARE"), KEY_SHARE("KEY SHARE");
@@ -191,6 +192,14 @@ public ForClause getForClause() {
         return forClause;
     }
 
+    public PostgresWithClause getWithClause() {
+        return withClause;
+    }
+
+    public void setWithClause(PostgresWithClause withClause) {
+        this.withClause = withClause;
+    }
+
     @Override
     public String asString() {
         return PostgresVisitor.asString(this);

diff --git a/src/sqlancer/postgres/ast/PostgresWithClause.java b/src/sqlancer/postgres/ast/PostgresWithClause.java
@@ -0,0 +1,64 @@
+package sqlancer.postgres.ast;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import sqlancer.postgres.PostgresSchema.PostgresDataType;
+
+
+public class PostgresWithClause implements PostgresExpression {
+    private final List<PostgresCTE> cteList;
+
+    public PostgresWithClause(List<PostgresCTE> cteList) {
+        if (cteList == null || cteList.isEmpty()) {
+            throw new IllegalArgumentException("CTE list cannot be null or empty");
+        }
+        this.cteList = new ArrayList<>(cteList);
+    }
+
+    public List<PostgresCTE> getCteList() {
+        return new ArrayList<>(cteList);
+    }
+
+    @Override
+    public PostgresDataType getExpressionType() {
+        return null;
+    }
+
+
+    public String asString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append("WITH ");
+
+        sb.append(cteList.stream()
+                 .map(PostgresCTE::asString)
+                 .collect(Collectors.joining(", ")));
+
+        return sb.toString();
+    }
+
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+
+    public static class Builder {
+        private List<PostgresCTE> cteList = new ArrayList<>();
+
+        public Builder addCTE(PostgresCTE cte) {
+            this.cteList.add(cte);
+            return this;
+        }
+
+        public Builder cteList(List<PostgresCTE> cteList) {
+            this.cteList = new ArrayList<>(cteList);
+            return this;
+        }
+
+        public PostgresWithClause build() {
+            return new PostgresWithClause(cteList);
+        }
+    }
+} 
diff --git a/src/sqlancer/postgres/gen/PostgresExpressionGenerator.java b/src/sqlancer/postgres/gen/PostgresExpressionGenerator.java
@@ -42,6 +42,7 @@
 import sqlancer.postgres.ast.PostgresColumnValue;
 import sqlancer.postgres.ast.PostgresConcatOperation;
 import sqlancer.postgres.ast.PostgresConstant;
+import sqlancer.postgres.ast.PostgresCTE;
 import sqlancer.postgres.ast.PostgresExpression;
 import sqlancer.postgres.ast.PostgresFunction;
 import sqlancer.postgres.ast.PostgresFunction.PostgresFunctionWithResult;
@@ -68,6 +69,7 @@
 import sqlancer.postgres.ast.PostgresWindowFunction;
 import sqlancer.postgres.ast.PostgresWindowFunction.WindowFrame;
 import sqlancer.postgres.ast.PostgresWindowFunction.WindowSpecification;
+import sqlancer.postgres.ast.PostgresWithClause;
 
 public class PostgresExpressionGenerator implements ExpressionGenerator<PostgresExpression>,
         NoRECGenerator<PostgresSelect, PostgresJoin, PostgresExpression, PostgresTable, PostgresColumn>,
@@ -972,4 +974,56 @@ boolean mutateLimit(PostgresSelect select) {
         }
         return increase;
     }
+
+    public static PostgresCTE generateCTE(PostgresGlobalState globalState, String cteName) {
+        PostgresTables tables = globalState.getSchema().getRandomTableNonEmptyTables();
+        PostgresSelect subquery = new PostgresSelect();
+
+        // Generate a simple subquery for the CTE
+        List<PostgresExpression> columns = new ArrayList<>();
+        PostgresExpressionGenerator gen = new PostgresExpressionGenerator(globalState).setColumns(tables.getColumns());
+
+        // Select 1-3 columns
+        int numColumns = Randomly.smallNumber() + 1;
+        for (int i = 0; i < numColumns; i++) {
+            columns.add(gen.generateExpression(0));
+        }
+
+        subquery.setSelectType(SelectType.getRandom());
+        subquery.setFromList(tables.getTables().stream()
+                .map(t -> new PostgresFromTable(t, Randomly.getBoolean()))
+                .collect(Collectors.toList()));
+        subquery.setFetchColumns(columns);
+
+        // Optionally add WHERE clause
+        if (Randomly.getBoolean()) {
+            subquery.setWhereClause(gen.generateExpression(0, PostgresDataType.BOOLEAN));
+        }
+
+        // Optionally add LIMIT
+        if (Randomly.getBoolean()) {
+            subquery.setLimitClause(PostgresConstant.createIntConstant(Randomly.getPositiveOrZeroNonCachedInteger()));
+        }
+
+        return new PostgresCTE(cteName, subquery);
+    }
+
+    /**
+     * Generates a WITH clause containing 1-3 CTEs.
+     * 
+     * @param globalState the global state
+     * @return a PostgresWithClause instance
+     */
+    public static PostgresWithClause generateWithClause(PostgresGlobalState globalState) {
+        List<PostgresCTE> cteList = new ArrayList<>();
+        int numCTEs = Randomly.smallNumber() + 1; // 1-3 CTEs
+
+        for (int i = 0; i < numCTEs; i++) {
+            String cteName = "cte_" + i;
+            PostgresCTE cte = generateCTE(globalState, cteName);
+            cteList.add(cte);
+        }
+
+        return new PostgresWithClause(cteList);
+    }
 }
diff --git a/src/sqlancer/postgres/gen/PostgresRandomQueryGenerator.java b/src/sqlancer/postgres/gen/PostgresRandomQueryGenerator.java
@@ -28,6 +28,12 @@ public static PostgresSelect createRandomQuery(int nrColumns, PostgresGlobalStat
             columns.add(gen.generateExpression(0));
         }
         PostgresSelect select = new PostgresSelect();
+
+        // Optionally add CTEs (WITH clause)
+        if (Randomly.getBooleanWithRatherLowProbability()) {
+            select.setWithClause(PostgresExpressionGenerator.generateWithClause(globalState));
+        }
+
         select.setSelectType(SelectType.getRandom());
         if (select.getSelectOption() == SelectType.DISTINCT && Randomly.getBoolean()) {
             select.setDistinctOnClause(gen.generateExpression(0));