Skip to content

Commit 74c599f

Browse files
qvadqvad
committed
Add YSQL generators for COPY, cursor operations, DO blocks, domains, EXPLAIN, functions, GRANT/REVOKE, table locks, and row-level policies with enhanced error handling
1 parent 0b336ba commit 74c599f

13 files changed

Lines changed: 1263 additions & 0 deletions

‎src/sqlancer/yugabyte/ysql/gen/YSQLCopyGenerator.java‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package sqlancer.yugabyte.ysql.gen;
2+
3+
import java.util.List;
4+
import java.util.stream.Collectors;
5+
6+
import sqlancer.IgnoreMeException;
7+
import sqlancer.Randomly;
8+
import sqlancer.common.query.ExpectedErrors;
9+
import sqlancer.common.query.SQLQueryAdapter;
10+
import sqlancer.yugabyte.ysql.YSQLErrors;
11+
import sqlancer.yugabyte.ysql.YSQLGlobalState;
12+
import sqlancer.yugabyte.ysql.YSQLSchema.YSQLColumn;
13+
import sqlancer.yugabyte.ysql.YSQLSchema.YSQLTable;
14+
15+
public final class YSQLCopyGenerator {
16+
17+
private YSQLCopyGenerator() {
18+
}
19+
20+
public static SQLQueryAdapter generate(YSQLGlobalState globalState) {
21+
ExpectedErrors errors = new ExpectedErrors();
22+
YSQLErrors.addTransactionErrors(errors);
23+
errors.add("does not exist");
24+
errors.add("is not a table");
25+
errors.add("COPY");
26+
errors.add("cannot copy");
27+
errors.add("must be superuser");
28+
errors.add("relative path not allowed");
29+
errors.add("could not open file");
30+
errors.add("invalid input syntax");
31+
errors.add("violates");
32+
errors.add("duplicate key");
33+
errors.add("No such file or directory");
34+
errors.add("Permission denied");
35+
errors.add("extra data after");
36+
errors.add("missing data for column");
37+
errors.add("cannot specify HEADER in BINARY mode");
38+
errors.add("not supported yet");
39+
errors.add("BINARY is not supported yet");
40+
41+
// Only use COPY TO STDOUT (safe - no file system access needed)
42+
return generateCopyToStdout(globalState, errors);
43+
}
44+
45+
private static SQLQueryAdapter generateCopyToStdout(YSQLGlobalState globalState, ExpectedErrors errors) {
46+
if (globalState.getSchema().getDatabaseTables().isEmpty()) {
47+
throw new IgnoreMeException();
48+
}
49+
YSQLTable table = globalState.getSchema().getRandomTable(t -> !t.isView());
50+
if (table == null) {
51+
throw new IgnoreMeException();
52+
}
53+
StringBuilder sb = new StringBuilder();
54+
sb.append("COPY ");
55+
56+
if (Randomly.getBoolean()) {
57+
// COPY table TO
58+
sb.append(table.getName());
59+
if (Randomly.getBoolean()) {
60+
List<String> cols = Randomly.nonEmptySubset(table.getColumns()).stream().map(YSQLColumn::getName)
61+
.collect(Collectors.toList());
62+
sb.append(" (").append(String.join(", ", cols)).append(")");
63+
}
64+
} else {
65+
// COPY (query) TO
66+
YSQLColumn col = table.getRandomColumn();
67+
sb.append("(SELECT ").append(col.getName()).append(" FROM ").append(table.getName());
68+
if (Randomly.getBoolean()) {
69+
sb.append(" LIMIT ").append(Randomly.smallNumber());
70+
}
71+
sb.append(")");
72+
}
73+
sb.append(" TO STDOUT");
74+
75+
if (Randomly.getBoolean()) {
76+
sb.append(" WITH (");
77+
sb.append("FORMAT ").append(Randomly.fromOptions("text", "csv", "binary"));
78+
if (Randomly.getBoolean()) {
79+
sb.append(", HEADER TRUE");
80+
}
81+
sb.append(")");
82+
}
83+
return new SQLQueryAdapter(sb.toString(), errors);
84+
}
85+
86+
}

‎src/sqlancer/yugabyte/ysql/gen/YSQLCursorGenerator.java‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
package sqlancer.yugabyte.ysql.gen;
2+
3+
import sqlancer.IgnoreMeException;
4+
import sqlancer.Randomly;
5+
import sqlancer.common.query.ExpectedErrors;
6+
import sqlancer.common.query.SQLQueryAdapter;
7+
import sqlancer.yugabyte.ysql.YSQLErrors;
8+
import sqlancer.yugabyte.ysql.YSQLGlobalState;
9+
import sqlancer.yugabyte.ysql.YSQLSchema.YSQLColumn;
10+
import sqlancer.yugabyte.ysql.YSQLSchema.YSQLDataType;
11+
import sqlancer.yugabyte.ysql.YSQLSchema.YSQLTable;
12+
import sqlancer.yugabyte.ysql.YSQLVisitor;
13+
14+
public final class YSQLCursorGenerator {
15+
16+
private static final int MAX_CURSORS = 5;
17+
18+
private YSQLCursorGenerator() {
19+
}
20+
21+
public static SQLQueryAdapter generate(YSQLGlobalState globalState) {
22+
ExpectedErrors errors = new ExpectedErrors();
23+
YSQLErrors.addTransactionErrors(errors);
24+
YSQLErrors.addCommonExpressionErrors(errors);
25+
errors.add("does not exist");
26+
errors.add("already in use");
27+
errors.add("is not open");
28+
errors.add("no such cursor");
29+
errors.add("cursor can only scan forward");
30+
errors.add("DECLARE CURSOR can only be used in transaction blocks");
31+
errors.add("is not a simply updatable relation");
32+
errors.add("no result set");
33+
errors.add("portal");
34+
35+
switch (Randomly.fromOptions(0, 1, 2, 3)) {
36+
case 0:
37+
return generateDeclare(globalState, errors);
38+
case 1:
39+
return generateFetch(errors);
40+
case 2:
41+
return generateMove(errors);
42+
case 3:
43+
return generateClose(errors);
44+
default:
45+
throw new AssertionError();
46+
}
47+
}
48+
49+
private static SQLQueryAdapter generateDeclare(YSQLGlobalState globalState, ExpectedErrors errors) {
50+
if (globalState.getSchema().getDatabaseTables().isEmpty()) {
51+
throw new IgnoreMeException();
52+
}
53+
YSQLTable table = globalState.getSchema().getRandomTable();
54+
StringBuilder sb = new StringBuilder();
55+
String cursorName = "cur" + globalState.getRandomly().getInteger(0, MAX_CURSORS);
56+
sb.append("DECLARE ").append(cursorName);
57+
if (Randomly.getBoolean()) {
58+
sb.append(" BINARY");
59+
}
60+
if (Randomly.getBoolean()) {
61+
sb.append(Randomly.fromOptions(" SCROLL", " NO SCROLL"));
62+
}
63+
sb.append(" CURSOR");
64+
if (Randomly.getBoolean()) {
65+
sb.append(Randomly.fromOptions(" WITH HOLD", " WITHOUT HOLD"));
66+
}
67+
sb.append(" FOR SELECT ");
68+
YSQLColumn col = table.getRandomColumn();
69+
sb.append(col.getName());
70+
sb.append(" FROM ").append(table.getName());
71+
if (Randomly.getBoolean()) {
72+
sb.append(" WHERE ");
73+
sb.append(YSQLVisitor.asString(
74+
YSQLExpressionGenerator.generateExpression(globalState, table.getColumns(), YSQLDataType.BOOLEAN)));
75+
}
76+
if (Randomly.getBoolean()) {
77+
sb.append(" ORDER BY ").append(col.getName());
78+
}
79+
return new SQLQueryAdapter(sb.toString(), errors);
80+
}
81+
82+
private static SQLQueryAdapter generateFetch(ExpectedErrors errors) {
83+
StringBuilder sb = new StringBuilder();
84+
sb.append("FETCH ");
85+
appendDirection(sb);
86+
sb.append(" FROM cur").append(Randomly.smallNumber() % MAX_CURSORS);
87+
return new SQLQueryAdapter(sb.toString(), errors);
88+
}
89+
90+
private static SQLQueryAdapter generateMove(ExpectedErrors errors) {
91+
StringBuilder sb = new StringBuilder();
92+
sb.append("MOVE ");
93+
appendDirection(sb);
94+
sb.append(" FROM cur").append(Randomly.smallNumber() % MAX_CURSORS);
95+
return new SQLQueryAdapter(sb.toString(), errors);
96+
}
97+
98+
private static SQLQueryAdapter generateClose(ExpectedErrors errors) {
99+
StringBuilder sb = new StringBuilder();
100+
sb.append("CLOSE ");
101+
if (Randomly.getBoolean()) {
102+
sb.append("ALL");
103+
} else {
104+
sb.append("cur").append(Randomly.smallNumber() % MAX_CURSORS);
105+
}
106+
return new SQLQueryAdapter(sb.toString(), errors);
107+
}
108+
109+
private static void appendDirection(StringBuilder sb) {
110+
switch (Randomly.fromOptions(0, 1, 2, 3, 4, 5, 6, 7)) {
111+
case 0:
112+
sb.append("NEXT");
113+
break;
114+
case 1:
115+
sb.append("PRIOR");
116+
break;
117+
case 2:
118+
sb.append("FIRST");
119+
break;
120+
case 3:
121+
sb.append("LAST");
122+
break;
123+
case 4:
124+
sb.append("FORWARD");
125+
break;
126+
case 5:
127+
sb.append("BACKWARD");
128+
break;
129+
case 6:
130+
sb.append("FORWARD ").append(Randomly.smallNumber());
131+
break;
132+
case 7:
133+
sb.append("ABSOLUTE ").append(Randomly.smallNumber());
134+
break;
135+
default:
136+
throw new AssertionError();
137+
}
138+
}
139+
140+
}

‎src/sqlancer/yugabyte/ysql/gen/YSQLDoBlockGenerator.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package sqlancer.yugabyte.ysql.gen;
2+
3+
import sqlancer.Randomly;
4+
import sqlancer.common.query.ExpectedErrors;
5+
import sqlancer.common.query.SQLQueryAdapter;
6+
import sqlancer.yugabyte.ysql.YSQLErrors;
7+
import sqlancer.yugabyte.ysql.YSQLGlobalState;
8+
import sqlancer.yugabyte.ysql.YSQLSchema.YSQLTable;
9+
10+
public final class YSQLDoBlockGenerator {
11+
12+
private YSQLDoBlockGenerator() {
13+
}
14+
15+
public static SQLQueryAdapter generate(YSQLGlobalState globalState) {
16+
ExpectedErrors errors = new ExpectedErrors();
17+
YSQLErrors.addTransactionErrors(errors);
18+
YSQLErrors.addCommonExpressionErrors(errors);
19+
errors.add("does not exist");
20+
errors.add("division by zero");
21+
errors.add("out of range");
22+
errors.add("syntax error");
23+
errors.add("query has no destination for result data");
24+
errors.add("query returned no rows");
25+
errors.add("duplicate key");
26+
errors.add("violates");
27+
errors.add("cannot cast");
28+
errors.add("invalid input syntax");
29+
30+
StringBuilder sb = new StringBuilder();
31+
sb.append("DO $$ ");
32+
33+
switch (Randomly.fromOptions(0, 1, 2, 3)) {
34+
case 0:
35+
appendRaiseBlock(sb);
36+
break;
37+
case 1:
38+
appendPerformBlock(sb, globalState);
39+
break;
40+
case 2:
41+
appendDeclareBlock(sb);
42+
break;
43+
case 3:
44+
appendIfBlock(sb, globalState);
45+
break;
46+
default:
47+
throw new AssertionError();
48+
}
49+
50+
sb.append(" $$");
51+
return new SQLQueryAdapter(sb.toString(), errors);
52+
}
53+
54+
private static void appendRaiseBlock(StringBuilder sb) {
55+
sb.append("BEGIN RAISE NOTICE 'test notice %', ");
56+
sb.append(Randomly.smallNumber());
57+
sb.append("; END;");
58+
}
59+
60+
private static void appendPerformBlock(StringBuilder sb, YSQLGlobalState globalState) {
61+
sb.append("BEGIN PERFORM ");
62+
if (globalState.getSchema().getDatabaseTables().isEmpty()) {
63+
sb.append("1");
64+
} else {
65+
YSQLTable table = globalState.getSchema().getRandomTable();
66+
sb.append("count(*) FROM ").append(table.getName());
67+
}
68+
sb.append("; END;");
69+
}
70+
71+
private static void appendDeclareBlock(StringBuilder sb) {
72+
sb.append("DECLARE v_val INTEGER; ");
73+
sb.append("BEGIN v_val := ").append(Randomly.smallNumber()).append("; ");
74+
sb.append("IF v_val > ").append(Randomly.smallNumber()).append(" THEN ");
75+
sb.append("RAISE NOTICE 'high: %', v_val; ");
76+
sb.append("END IF; END;");
77+
}
78+
79+
private static void appendIfBlock(StringBuilder sb, YSQLGlobalState globalState) {
80+
sb.append("DECLARE v_cnt BIGINT; ");
81+
sb.append("BEGIN ");
82+
if (globalState.getSchema().getDatabaseTables().isEmpty()) {
83+
sb.append("v_cnt := 0; ");
84+
} else {
85+
YSQLTable table = globalState.getSchema().getRandomTable();
86+
sb.append("SELECT count(*) INTO v_cnt FROM ").append(table.getName()).append("; ");
87+
}
88+
sb.append("IF v_cnt > ").append(Randomly.smallNumber()).append(" THEN ");
89+
sb.append("RAISE NOTICE 'count: %', v_cnt; ");
90+
sb.append("END IF; END;");
91+
}
92+
93+
}

0 commit comments

Comments
 (0)