From e8bea0b58fa05bab5dfce86d94472194ba65dda1 Mon Sep 17 00:00:00 2001 From: yaojiejia Date: Tue, 1 Jul 2025 23:00:36 -0400 Subject: [PATCH] Add CTE Function to Postgres System --- .../PostgresExpectedValueVisitor.java | 16 +++++ .../postgres/PostgresToStringVisitor.java | 26 ++++++++ src/sqlancer/postgres/PostgresVisitor.java | 10 +++ src/sqlancer/postgres/ast/PostgresCTE.java | 65 +++++++++++++++++++ src/sqlancer/postgres/ast/PostgresSelect.java | 9 +++ .../postgres/ast/PostgresWithClause.java | 64 ++++++++++++++++++ .../gen/PostgresExpressionGenerator.java | 54 +++++++++++++++ .../gen/PostgresRandomQueryGenerator.java | 6 ++ 8 files changed, 250 insertions(+) create mode 100644 src/sqlancer/postgres/ast/PostgresCTE.java create mode 100644 src/sqlancer/postgres/ast/PostgresWithClause.java diff --git a/src/sqlancer/postgres/PostgresExpectedValueVisitor.java b/src/sqlancer/postgres/PostgresExpectedValueVisitor.java index 8985caf0b..620dc892c 100644 --- 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 index 87bd3c429..a351d0864 100644 --- 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 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 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 index d5bf71e81..bd35c7fe9 100644 --- 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 new file mode 100644 index 000000000..295b2add1 --- /dev/null +++ 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); + } + } +} \ No newline at end of file diff --git a/src/sqlancer/postgres/ast/PostgresSelect.java b/src/sqlancer/postgres/ast/PostgresSelect.java index 1d7172e3e..2e22fe95d 100644 --- a/src/sqlancer/postgres/ast/PostgresSelect.java +++ b/src/sqlancer/postgres/ast/PostgresSelect.java @@ -24,6 +24,7 @@ public class PostgresSelect extends SelectBase private ForClause forClause; private List windowFunctions = new ArrayList<>(); private final Map 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 new file mode 100644 index 000000000..1b92ffceb --- /dev/null +++ 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 cteList; + + public PostgresWithClause(List cteList) { + if (cteList == null || cteList.isEmpty()) { + throw new IllegalArgumentException("CTE list cannot be null or empty"); + } + this.cteList = new ArrayList<>(cteList); + } + + public List 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 cteList = new ArrayList<>(); + + public Builder addCTE(PostgresCTE cte) { + this.cteList.add(cte); + return this; + } + + public Builder cteList(List cteList) { + this.cteList = new ArrayList<>(cteList); + return this; + } + + public PostgresWithClause build() { + return new PostgresWithClause(cteList); + } + } +} \ No newline at end of file diff --git a/src/sqlancer/postgres/gen/PostgresExpressionGenerator.java b/src/sqlancer/postgres/gen/PostgresExpressionGenerator.java index bad87affc..c6054c8a4 100644 --- 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, NoRECGenerator, @@ -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 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 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 index 64d6278de..706b2bcdd 100644 --- 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));