Skip to content

Commit 875e656

Browse files
mriggermrigger
committed
[Postgres] Improve the query partitioning oracle
1 parent 346bbfc commit 875e656

8 files changed

Lines changed: 74 additions & 30 deletions

‎src/sqlancer/postgres/PostgresOptions.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,14 @@ public TestOracle create(PostgresGlobalState globalState) throws SQLException {
3434
return new PostgresPivotedQuerySynthesisGenerator(globalState);
3535
}
3636
},
37+
HAVING {
38+
39+
@Override
40+
public TestOracle create(PostgresGlobalState globalState) throws SQLException {
41+
return new PostgresQueryPartitioningHavingTester(globalState);
42+
}
43+
44+
},
3745
QUERY_PARTITIONING {
3846
@Override
3947
public TestOracle create(PostgresGlobalState globalState) throws SQLException {

‎src/sqlancer/postgres/PostgresProvider.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ private static int mapActions(PostgresGlobalState globalState, Action a) {
140140
int nrPerformed;
141141
switch (a) {
142142
case CREATE_INDEX:
143-
nrPerformed = r.getInteger(0, 10);
143+
nrPerformed = r.getInteger(0, 3);
144144
break;
145145
case CREATE_STATISTICS:
146146
nrPerformed = r.getInteger(0, 5);
@@ -158,7 +158,7 @@ private static int mapActions(PostgresGlobalState globalState, Action a) {
158158
break;
159159
case REINDEX:
160160
case RESET:
161-
nrPerformed = r.getInteger(0, 5);
161+
nrPerformed = r.getInteger(0, 3);
162162
break;
163163
case DELETE:
164164
case RESET_ROLE:
@@ -167,7 +167,7 @@ private static int mapActions(PostgresGlobalState globalState, Action a) {
167167
nrPerformed = r.getInteger(0, 5);
168168
break;
169169
case ANALYZE:
170-
nrPerformed = r.getInteger(0, 10);
170+
nrPerformed = r.getInteger(0, 3);
171171
break;
172172
case VACUUM:
173173
case SET_CONSTRAINTS:

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ public static void addCommonFetchErrors(Set<String> errors) {
2626
errors.add("FULL JOIN is only supported with merge-joinable or hash-joinable join conditions");
2727
errors.add("but it cannot be referenced from this part of the query");
2828
errors.add("missing FROM-clause entry for table");
29+
30+
errors.add("canceling statement due to statement timeout");
2931
}
3032

3133
public static void addCommonTableErrors(Set<String> errors) {
@@ -405,6 +407,7 @@ public static void addGroupingErrors(Set<String> errors) {
405407
errors.add("non-integer constant in GROUP BY"); // TODO
406408
errors.add("must appear in the GROUP BY clause or be used in an aggregate function");
407409
errors.add("is not in select list");
410+
errors.add("aggregate functions are not allowed in GROUP BY");
408411
}
409412

410413
}

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,8 +158,13 @@ private PostgresExpression generateBooleanExpression(int depth) {
158158
return new PostgresPrefixOperation(generateExpression(depth + 1, PostgresDataType.BOOLEAN),
159159
PrefixOperator.NOT);
160160
case BINARY_LOGICAL_OPERATOR:
161-
return new PostgresBinaryLogicalOperation(generateExpression(depth + 1, PostgresDataType.BOOLEAN),
162-
generateExpression(depth + 1, PostgresDataType.BOOLEAN), BinaryLogicalOperator.getRandom());
161+
PostgresExpression first = generateExpression(depth + 1, PostgresDataType.BOOLEAN);
162+
int nr = Randomly.smallNumber() + 1;
163+
for (int i = 0; i < nr; i++) {
164+
first = new PostgresBinaryLogicalOperation(first,
165+
generateExpression(depth + 1, PostgresDataType.BOOLEAN), BinaryLogicalOperator.getRandom());
166+
}
167+
return first;
163168
case BINARY_COMPARISON:
164169
PostgresDataType dataType = getMeaningfulType();
165170
return generateComparison(depth, dataType);
@@ -562,4 +567,9 @@ public PostgresAggregate generateArgsForAggregate(PostgresDataType dataType, Pos
562567
return new PostgresAggregate(args, agg);
563568
}
564569

570+
public PostgresExpressionGenerator allowAggregates(boolean value) {
571+
allowAggregateFunctions = value;
572+
return this;
573+
}
574+
565575
}

‎src/sqlancer/postgres/test/PostgresNoRECOracle.java‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -76,15 +76,7 @@ public void check() throws SQLException {
7676
}
7777
List<PostgresTable> tables = randomTables.getTables();
7878

79-
List<PostgresJoin> joinStatements = new ArrayList<>();
80-
for (int i = 1; i < tables.size(); i++) {
81-
PostgresExpression joinClause = getRandomWhereCondition(columns);
82-
PostgresTable table = Randomly.fromList(tables);
83-
tables.remove(table);
84-
PostgresJoinType options = PostgresJoinType.getRandom();
85-
PostgresJoin j = new PostgresJoin(table, joinClause, options);
86-
joinStatements.add(j);
87-
}
79+
List<PostgresJoin> joinStatements = getJoinStatements(globalState, columns, tables);
8880
List<PostgresExpression> fromTables = tables.stream().map(t -> new PostgresFromTable(t, Randomly.getBoolean()))
8981
.collect(Collectors.toList());
9082
int secondCount = getSecondQuery(fromTables, randomWhereCondition, groupBys, joinStatements);
@@ -99,6 +91,20 @@ public void check() throws SQLException {
9991
}
10092
}
10193

94+
public static List<PostgresJoin> getJoinStatements(PostgresGlobalState globalState, List<PostgresColumn> columns, List<PostgresTable> tables) {
95+
List<PostgresJoin> joinStatements = new ArrayList<>();
96+
PostgresExpressionGenerator gen = new PostgresExpressionGenerator(globalState).setColumns(columns);
97+
for (int i = 1; i < tables.size(); i++) {
98+
PostgresExpression joinClause = gen.generateExpression(PostgresDataType.BOOLEAN);
99+
PostgresTable table = Randomly.fromList(tables);
100+
tables.remove(table);
101+
PostgresJoinType options = PostgresJoinType.getRandom();
102+
PostgresJoin j = new PostgresJoin(table, joinClause, options);
103+
joinStatements.add(j);
104+
}
105+
return joinStatements;
106+
}
107+
102108
private List<PostgresExpression> getRandomExpressions(List<PostgresColumn> columns) {
103109
List<PostgresExpression> randomExpressions = columns.stream().map(c -> new PostgresColumnValue(c, null))
104110
.collect(Collectors.toList());

‎src/sqlancer/postgres/test/PostgresQueryPartitioningAggregateTester.java‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import sqlancer.postgres.ast.PostgresAggregate.PostgresAggregateFunction;
2121
import sqlancer.postgres.ast.PostgresAlias;
2222
import sqlancer.postgres.ast.PostgresExpression;
23+
import sqlancer.postgres.ast.PostgresJoin;
2324
import sqlancer.postgres.ast.PostgresPostfixOperation;
2425
import sqlancer.postgres.ast.PostgresPostfixOperation.PostfixOperator;
2526
import sqlancer.postgres.ast.PostgresPrefixOperation;
@@ -38,10 +39,13 @@ public PostgresQueryPartitioningAggregateTester(PostgresGlobalState state) {
3839
super(state);
3940
PostgresCommon.addGroupingErrors(errors);
4041
}
41-
42+
4243
public void check() throws SQLException {
4344
super.check();
44-
PostgresAggregateFunction aggregateFunction = Randomly.fromOptions(PostgresAggregateFunction.MAX, PostgresAggregateFunction.MIN, PostgresAggregateFunction.SUM, PostgresAggregateFunction.BIT_AND, PostgresAggregateFunction.BIT_OR, PostgresAggregateFunction.BOOL_AND, PostgresAggregateFunction.BOOL_OR);
45+
PostgresAggregateFunction aggregateFunction = Randomly.fromOptions(PostgresAggregateFunction.MAX,
46+
PostgresAggregateFunction.MIN, PostgresAggregateFunction.SUM, PostgresAggregateFunction.BIT_AND,
47+
PostgresAggregateFunction.BIT_OR, PostgresAggregateFunction.BOOL_AND, PostgresAggregateFunction.BOOL_OR,
48+
PostgresAggregateFunction.COUNT);
4549
PostgresAggregate aggregate = gen.generateArgsForAggregate(aggregateFunction.getRandomReturnType(),
4650
aggregateFunction);
4751
List<PostgresExpression> fetchColumns = new ArrayList<>();
@@ -50,9 +54,6 @@ public void check() throws SQLException {
5054
fetchColumns.add(gen.generateAggregate());
5155
}
5256
select.setFetchColumns(Arrays.asList(aggregate));
53-
// if (Randomly.getBooleanWithRatherLowProbability()) {
54-
// select.setJoinList(PostgresNoRECTester.getJoins(from, state));
55-
// }
5657
if (Randomly.getBooleanWithRatherLowProbability()) {
5758
select.setOrderByExpressions(gen.generateOrderBy());
5859
}
@@ -85,9 +86,9 @@ private String createMetamorphicUnionQuery(PostgresSelect select, PostgresAggreg
8586
PostgresExpression negatedClause = new PostgresPrefixOperation(whereClause, PrefixOperator.NOT);
8687
PostgresExpression notNullClause = new PostgresPostfixOperation(whereClause, PostfixOperator.IS_NULL);
8788
List<PostgresExpression> mappedAggregate = mapped(aggregate);
88-
PostgresSelect leftSelect = getSelect(mappedAggregate, from, whereClause/* , select.getJoinList() */);
89-
PostgresSelect middleSelect = getSelect(mappedAggregate, from, negatedClause/* , select.getJoinList() */);
90-
PostgresSelect rightSelect = getSelect(mappedAggregate, from, notNullClause/* , select.getJoinList() */);
89+
PostgresSelect leftSelect = getSelect(mappedAggregate, from, whereClause, select.getJoinClauses());
90+
PostgresSelect middleSelect = getSelect(mappedAggregate, from, negatedClause, select.getJoinClauses());
91+
PostgresSelect rightSelect = getSelect(mappedAggregate, from, notNullClause, select.getJoinClauses());
9192
metamorphicQuery = "SELECT " + getOuterAggregateFunction(aggregate).toString() + " FROM (";
9293
metamorphicQuery += PostgresVisitor.asString(leftSelect) + " UNION ALL "
9394
+ PostgresVisitor.asString(middleSelect) + " UNION ALL " + PostgresVisitor.asString(rightSelect);
@@ -116,7 +117,7 @@ private String getAggregateResult(String queryString) throws SQLException {
116117
private List<PostgresExpression> mapped(PostgresAggregate aggregate) {
117118
switch (aggregate.getFunction()) {
118119
case SUM:
119-
// case COUNT:
120+
case COUNT:
120121
case BIT_AND:
121122
case BIT_OR:
122123
case BOOL_AND:
@@ -150,21 +151,20 @@ private String getOuterAggregateFunction(PostgresAggregate aggregate) {
150151
switch (aggregate.getFunction()) {
151152
// case AVG:
152153
// return "SUM(agg0::DECIMAL)/SUM(agg1)::DECIMAL";
153-
// case COUNT:
154-
// case COUNT_ROWS:
155-
// return PostgresAggregateFunction.SUM.toString() + "(agg0)";
154+
case COUNT:
155+
return PostgresAggregateFunction.SUM.toString() + "(agg0)";
156156
default:
157157
return aggregate.getFunction().toString() + "(agg0)";
158158
}
159159
}
160160

161161
private PostgresSelect getSelect(List<PostgresExpression> aggregates, List<PostgresExpression> from,
162-
PostgresExpression whereClause/* , List<PostgresJoin> joinList */) {
162+
PostgresExpression whereClause, List<PostgresJoin> joinList) {
163163
PostgresSelect leftSelect = new PostgresSelect();
164164
leftSelect.setFetchColumns(aggregates);
165165
leftSelect.setFromList(from);
166166
leftSelect.setWhereClause(whereClause);
167-
// leftSelect.setJoinList(joinList);
167+
leftSelect.setJoinClauses(joinList);
168168
if (Randomly.getBooleanWithSmallProbability()) {
169169
leftSelect.setGroupByExpressions(gen.generateExpressions(Randomly.smallNumber() + 1));
170170
}

‎src/sqlancer/postgres/test/PostgresQueryPartitioningBase.java‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,11 @@
1212
import sqlancer.postgres.PostgresGlobalState;
1313
import sqlancer.postgres.PostgresSchema;
1414
import sqlancer.postgres.PostgresSchema.PostgresDataType;
15+
import sqlancer.postgres.PostgresSchema.PostgresTable;
1516
import sqlancer.postgres.PostgresSchema.PostgresTables;
1617
import sqlancer.postgres.ast.PostgresColumnValue;
1718
import sqlancer.postgres.ast.PostgresExpression;
19+
import sqlancer.postgres.ast.PostgresJoin;
1820
import sqlancer.postgres.ast.PostgresPostfixOperation;
1921
import sqlancer.postgres.ast.PostgresPostfixOperation.PostfixOperator;
2022
import sqlancer.postgres.ast.PostgresPrefixOperation;
@@ -49,12 +51,15 @@ public void check() throws SQLException {
4951
targetTables = s.getRandomTableNonEmptyTables();
5052
gen = new PostgresExpressionGenerator(state).setColumns(targetTables.getColumns());
5153
select = new PostgresSelect();
52-
select.setFetchColumns(Arrays.asList(new PostgresColumnValue(targetTables.getColumns().get(0), null)));
53-
List<PostgresExpression> tableList = targetTables.getTables().stream()
54+
select.setFetchColumns(generateFetchColumns());
55+
List<PostgresTable> tables = targetTables.getTables();
56+
List<PostgresJoin> joins = PostgresNoRECOracle.getJoinStatements(state, targetTables.getColumns(), tables);
57+
List<PostgresExpression> tableList = tables.stream()
5458
.map(t -> new PostgresFromTable(t, Randomly.getBoolean())).collect(Collectors.toList());
5559
// TODO joins
5660
select.setFromList(tableList);
5761
select.setWhereClause(null);
62+
select.setJoinClauses(joins);
5863
predicate = generatePredicate();
5964
negatedPredicate = new PostgresPrefixOperation(predicate, PostgresPrefixOperation.PrefixOperator.NOT);
6065
isNullPredicate = new PostgresPostfixOperation(predicate, PostfixOperator.IS_NULL);
@@ -63,6 +68,10 @@ public void check() throws SQLException {
6368
}
6469
}
6570

71+
List<PostgresExpression> generateFetchColumns() {
72+
return Arrays.asList(new PostgresColumnValue(targetTables.getColumns().get(0), null));
73+
}
74+
6675
PostgresExpression generatePredicate() {
6776
return gen.generateExpression(PostgresDataType.BOOLEAN);
6877
}

‎src/sqlancer/postgres/test/PostgresQueryPartitioningHavingTester.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,4 +73,12 @@ PostgresExpression generatePredicate() {
7373
return gen.generateHavingClause();
7474
}
7575

76+
@Override
77+
List<PostgresExpression> generateFetchColumns() {
78+
List<PostgresExpression> expressions = gen.allowAggregates(true)
79+
.generateExpressions(Randomly.smallNumber() + 1);
80+
gen.allowAggregates(false);
81+
return expressions;
82+
}
83+
7684
}

0 commit comments

Comments
 (0)