Skip to content

Commit e68056a

Browse files
yaojiejiayaojiejia
committed
Feat(postgres): Add Window Function Support For Query Generation
Implement comprehensive window function support for PostgreSQL testing, including: - Window function expressions with partition by, order by, ow_number, rank, and frame clauses - Named window definitions via window clause - Integration with existing expression generation framework The implementation allows SQLancer to generate and test queries containing window functions, expanding test coverage for PostgreSQL's analytical capabilities. Window functions are now included in the expression generation options and can be mutated during testing like other query components. Key changes: - Add PostgresWindowFunction and supporting classes - Enhance PostgresSelect to handle window clauses - Update expression generator for window function creation - Modify visitor for proper SQL generation
1 parent c50b87d commit e68056a

6 files changed

Lines changed: 388 additions & 3 deletions

‎src/sqlancer/postgres/PostgresToStringVisitor.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,4 +362,37 @@ public void visit(PostgresLikeOperation op) {
362362
super.visit((BinaryOperation<PostgresExpression>) op);
363363
}
364364

365+
@Override
366+
public void visit(PostgresWindowFunction windowFunction) {
367+
sb.append(windowFunction.getFunctionName());
368+
sb.append("(");
369+
visit(windowFunction.getArguments());
370+
sb.append(") OVER (");
371+
372+
WindowSpecification spec = windowFunction.getWindowSpec();
373+
if (!spec.getPartitionBy().isEmpty()) {
374+
sb.append("PARTITION BY ");
375+
visit(spec.getPartitionBy());
376+
}
377+
378+
if (!spec.getOrderBy().isEmpty()) {
379+
if (!spec.getPartitionBy().isEmpty()) {
380+
sb.append(" ");
381+
}
382+
sb.append("ORDER BY ");
383+
visit(spec.getOrderBy());
384+
}
385+
386+
if (spec.getFrame() != null) {
387+
sb.append(" ");
388+
WindowFrame frame = spec.getFrame();
389+
sb.append(frame.getType().getSQL());
390+
sb.append(" BETWEEN ");
391+
visit(frame.getStartExpr());
392+
sb.append(" AND ");
393+
visit(frame.getEndExpr());
394+
}
395+
396+
sb.append(")");
397+
}
365398
}

‎src/sqlancer/postgres/PostgresVisitor.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ public interface PostgresVisitor {
7272

7373
void visit(PostgresLikeOperation op);
7474

75+
void visit(PostgresWindowFunction windowFunction);
76+
7577
default void visit(PostgresExpression expression) {
7678
if (expression instanceof PostgresConstant) {
7779
visit((PostgresConstant) expression);
@@ -113,7 +115,10 @@ default void visit(PostgresExpression expression) {
113115
visit((PostgresColumnReference) expression);
114116
} else if (expression instanceof PostgresTableReference) {
115117
visit((PostgresTableReference) expression);
116-
} else {
118+
} else if (expression instanceof PostgresWindowFunction) {
119+
visit((PostgresWindowFunction) expression);
120+
}
121+
else {
117122
throw new AssertionError(expression);
118123
}
119124
}

‎src/sqlancer/postgres/ast/PostgresSelect.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,55 @@ public static ForClause getRandom() {
3737
}
3838
}
3939

40+
public static class WindowDefinition {
41+
private final List<PostgresExpression> partitionBy;
42+
private final List<PostgresOrderByTerm> orderBy;
43+
private final WindowFrame frame;
44+
45+
public WindowDefinition(List<PostgresExpression> partitionBy,
46+
List<PostgresOrderByTerm> orderBy,
47+
WindowFrame frame) {
48+
this.partitionBy = partitionBy;
49+
this.orderBy = orderBy;
50+
this.frame = frame;
51+
}
52+
53+
public List<PostgresExpression> getPartitionBy() {
54+
return partitionBy;
55+
}
56+
57+
public List<PostgresOrderByTerm> getOrderBy() {
58+
return orderBy;
59+
}
60+
61+
public WindowFrame getFrame() {
62+
return frame;
63+
}
64+
}
65+
66+
// Getters setters for windowfunctions
67+
public List<PostgresExpression> getWindowFunctions() {
68+
return windowFunctions;
69+
}
70+
71+
public void setWindowFunctions(List<PostgresExpression> windowFunctions) {
72+
this.windowFunctions = windowFunctions != null ? windowFunctions : new ArrayList<>();
73+
}
74+
75+
// Add methods for window definitions
76+
public void addWindowDefinition(String name, WindowDefinition definition) {
77+
windowDefinitions.put(name, definition);
78+
}
79+
80+
public WindowDefinition getWindowDefinition(String name) {
81+
return windowDefinitions.get(name);
82+
}
83+
84+
public Map<String, WindowDefinition> getWindowDefinitions() {
85+
return windowDefinitions;
86+
}
87+
88+
4089
public static class PostgresFromTable implements PostgresExpression {
4190
private final PostgresTable t;
4291
private final boolean only;

‎src/sqlancer/postgres/ast/PostgresWindowFunction.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
package sqlancer.postgres.ast;
2+
3+
import java.util.List;
4+
5+
import sqlancer.Randomly;
6+
import sqlancer.postgres.PostgresSchema.PostgresDataType;
7+
8+
public class PostgresWindowFunction implements PostgresExpression {
9+
10+
private final String functionName;
11+
private final List<PostgresExpression> arguments;
12+
private final WindowSpecification windowSpec;
13+
private final PostgresDataType returnType;
14+
15+
public static class WindowSpecification {
16+
private final List<PostgresExpression> partitionBy;
17+
private final List<PostgresOrderByTerm> orderBy;
18+
private final WindowFrame frame;
19+
20+
public WindowSpecification(List<PostgresExpression> partitionBy,
21+
List<PostgresOrderByTerm> orderBy,
22+
WindowFrame frame) {
23+
this.partitionBy = partitionBy;
24+
this.orderBy = orderBy;
25+
this.frame = frame;
26+
}
27+
28+
public List<PostgresExpression> getPartitionBy() {
29+
return partitionBy;
30+
}
31+
32+
public List<PostgresOrderByTerm> getOrderBy() {
33+
return orderBy;
34+
}
35+
36+
public WindowFrame getFrame() {
37+
return frame;
38+
}
39+
}
40+
41+
public static class WindowFrame {
42+
public enum FrameType {
43+
ROWS("ROWS"),
44+
RANGE("RANGE");
45+
46+
private final String sql;
47+
48+
FrameType(String sql) {
49+
this.sql = sql;
50+
}
51+
52+
public String getSQL() {
53+
return sql;
54+
}
55+
}
56+
57+
private final FrameType type;
58+
private final PostgresExpression startExpr;
59+
private final PostgresExpression endExpr;
60+
61+
public WindowFrame(FrameType type, PostgresExpression startExpr, PostgresExpression endExpr) {
62+
this.type = type;
63+
this.startExpr = startExpr;
64+
this.endExpr = endExpr;
65+
}
66+
67+
public FrameType getType() {
68+
return type;
69+
}
70+
71+
public PostgresExpression getStartExpr() {
72+
return startExpr;
73+
}
74+
75+
public PostgresExpression getEndExpr() {
76+
return endExpr;
77+
}
78+
}
79+
80+
public PostgresWindowFunction(String functionName, List<PostgresExpression> arguments,
81+
WindowSpecification windowSpec, PostgresDataType returnType) {
82+
this.functionName = functionName;
83+
this.arguments = arguments;
84+
this.windowSpec = windowSpec;
85+
this.returnType = returnType;
86+
}
87+
88+
public String getFunctionName() {
89+
return functionName;
90+
}
91+
92+
public List<PostgresExpression> getArguments() {
93+
return arguments;
94+
}
95+
96+
public WindowSpecification getWindowSpec() {
97+
return windowSpec;
98+
}
99+
100+
@Override
101+
public PostgresDataType getExpressionType() {
102+
return returnType;
103+
}
104+
}

‎src/sqlancer/postgres/gen/PostgresExpressionGenerator.java‎

Lines changed: 100 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,53 @@ private PostgresExpression generateTextExpression(int depth) {
422422
throw new AssertionError();
423423
}
424424
}
425+
426+
private PostgresExpression generateWindowFunction(int depth, PostgresDataType returnType) {
427+
// Generate window function arguments
428+
List<PostgresExpression> arguments = new ArrayList<>();
429+
if (Randomly.getBoolean()) {
430+
arguments.add(generateExpression(depth + 1));
431+
}
432+
433+
// Generate PARTITION BY expressions
434+
List<PostgresExpression> partitionBy = new ArrayList<>();
435+
if (Randomly.getBoolean()) {
436+
int count = Randomly.smallNumber();
437+
for (int i = 0; i < count; i++) {
438+
partitionBy.add(generateExpression(depth + 1));
439+
}
440+
}
441+
442+
// Generate ORDER BY expressions
443+
List<PostgresOrderByTerm> orderBy = new ArrayList<>();
444+
if (Randomly.getBoolean()) {
445+
int count = Randomly.smallNumber();
446+
for (int i = 0; i < count; i++) {
447+
PostgresExpression expr = generateExpression(depth + 1);
448+
orderBy.add(new PostgresOrderByTerm(expr, Randomly.getBoolean()));
449+
}
450+
}
451+
452+
// Generate window frame
453+
WindowFrame frame = null;
454+
if (Randomly.getBoolean()) {
455+
WindowFrame.FrameType frameType = Randomly.fromOptions(WindowFrame.FrameType.values());
456+
PostgresExpression startExpr = generateConstant(globalState.getRandomly(), PostgresDataType.INT);
457+
PostgresExpression endExpr = generateConstant(globalState.getRandomly(), PostgresDataType.INT);
458+
frame = new WindowFrame(frameType, startExpr, endExpr);
459+
}
460+
461+
WindowSpecification windowSpec = new WindowSpecification(partitionBy, orderBy, frame);
462+
463+
// Select a window function
464+
String functionName = Randomly.fromList(Arrays.asList(
465+
"row_number", "rank", "dense_rank", "percent_rank",
466+
"cume_dist", "ntile", "lag", "lead", "first_value",
467+
"last_value", "nth_value"
468+
));
469+
470+
return new PostgresWindowFunction(functionName, arguments, windowSpec, returnType);
471+
}
425472

426473
private PostgresExpression generateConcat(int depth) {
427474
PostgresExpression left = generateExpression(depth + 1, PostgresDataType.TEXT);
@@ -447,7 +494,7 @@ private PostgresExpression generateBitExpression(int depth) {
447494
}
448495

449496
private enum IntExpression {
450-
UNARY_OPERATION, FUNCTION, CAST, BINARY_ARITHMETIC_EXPRESSION
497+
UNARY_OPERATION, FUNCTION, CAST, BINARY_ARITHMETIC_EXPRESSION, WINDOW_FUNCTION
451498
}
452499

453500
private PostgresExpression generateIntExpression(int depth) {
@@ -667,7 +714,21 @@ public PostgresExpression generateBooleanExpression() {
667714

668715
@Override
669716
public PostgresSelect generateSelect() {
670-
return new PostgresSelect();
717+
PostgresSelect select = new PostgresSelect();
718+
719+
// Add window functions to fetch columns if appropriate
720+
if (Randomly.getBoolean()) {
721+
List<PostgresExpression> windowFunctions = new ArrayList<>();
722+
int numWindowFunctions = Randomly.smallNumber();
723+
for (int i = 0; i < numWindowFunctions; i++) {
724+
windowFunctions.add(generateWindowFunction(0, Randomly.fromList(Arrays.asList(
725+
PostgresDataType.INT, PostgresDataType.FLOAT
726+
))));
727+
}
728+
select.setWindowFunctions(windowFunctions);
729+
}
730+
731+
return select;
671732
}
672733

673734
@Override
@@ -768,6 +829,43 @@ public boolean mutate(PostgresSelect select) {
768829
return Randomly.fromList(mutators).apply(select);
769830
}
770831

832+
833+
@Override
834+
public boolean mutate(PostgresSelect select) {
835+
List<Function<PostgresSelect, Boolean>> mutators = new ArrayList<>();
836+
837+
mutators.add(this::mutateJoin);
838+
mutators.add(this::mutateWhere);
839+
mutators.add(this::mutateGroupBy);
840+
mutators.add(this::mutateHaving);
841+
mutators.add(this::mutateWindowFunction); // Add window function mutation
842+
if (!PostgresBugs.bug18643) {
843+
mutators.add(this::mutateAnd);
844+
mutators.add(this::mutateOr);
845+
}
846+
mutators.add(this::mutateDistinct);
847+
848+
return Randomly.fromList(mutators).apply(select);
849+
}
850+
851+
private boolean mutateWindowFunction(PostgresSelect select) {
852+
List<PostgresExpression> windowFunctions = select.getWindowFunctions();
853+
if (windowFunctions == null || windowFunctions.isEmpty()) {
854+
// Add a new window function
855+
windowFunctions = new ArrayList<>();
856+
windowFunctions.add(generateWindowFunction(0, PostgresDataType.INT));
857+
select.setWindowFunctions(windowFunctions);
858+
return false;
859+
} else {
860+
// Remove a random window function
861+
windowFunctions.remove(Randomly.fromList(windowFunctions));
862+
if (windowFunctions.isEmpty()) {
863+
select.setWindowFunctions(null);
864+
}
865+
return true;
866+
}
867+
}
868+
771869
boolean mutateJoin(PostgresSelect select) {
772870
if (select.getJoinList().isEmpty()) {
773871
return false;

0 commit comments

Comments
 (0)