Skip to content

Commit 495ad71

Browse files
yaojiejiayaojiejia
committed
Refactored the module
1 parent 7e0b0f4 commit 495ad71

3 files changed

Lines changed: 51 additions & 9 deletions

File tree

‎src/sqlancer/postgres/PostgresOptions.java‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
public class PostgresOptions implements DBMSSpecificOptions<PostgresOracleFactory> {
1515
public static final String DEFAULT_HOST = "localhost";
1616
public static final int DEFAULT_PORT = 5432;
17-
private static final boolean DEFAULT_TEST_TABLESPACES = determineDefaultTablespaceSupport();
17+
private static Boolean defaultTestTablespaces = null;
1818

1919
@Parameter(names = "--bulk-insert", description = "Specifies whether INSERT statements should be issued in bulk", arity = 1)
2020
public boolean allowBulkInsert;
@@ -26,7 +26,7 @@ public class PostgresOptions implements DBMSSpecificOptions<PostgresOracleFactor
2626
public boolean testCollations = true;
2727

2828
@Parameter(names = "--test-tablespaces", description = "Specifies whether to test tablespace creation (default is OS-dependent)", arity = 1)
29-
public boolean testTablespaces = DEFAULT_TEST_TABLESPACES;
29+
public boolean testTablespaces = false;
3030

3131
@Parameter(names = "--tablespace-path", description = "Base path for tablespace directories (default is OS-dependent)", arity = 1)
3232
public String tablespacePath = getDefaultTablespacePath();
@@ -74,4 +74,46 @@ public static String getDefaultTablespacePath() {
7474
public List<PostgresOracleFactory> getTestOracleFactory() {
7575
return oracle;
7676
}
77+
78+
public String getTablespacePath() {
79+
if (tablespacePath == null || tablespacePath.trim().isEmpty()) {
80+
throw new AssertionError("Tablespace path is null or empty. Please configure --tablespace-path");
81+
}
82+
83+
File path = new File(tablespacePath);
84+
85+
// Check if the directory exists or can be created
86+
if (!path.exists()) {
87+
if (!path.mkdirs()) {
88+
throw new AssertionError("Cannot create tablespace directory: " + tablespacePath +
89+
". Please ensure the parent directory exists and you have write permissions.");
90+
}
91+
}
92+
93+
// Check if it's actually a directory
94+
if (!path.isDirectory()) {
95+
throw new AssertionError("Tablespace path is not a directory: " + tablespacePath);
96+
}
97+
98+
// Check write permissions
99+
if (!path.canWrite()) {
100+
throw new AssertionError("No write permissions for tablespace directory: " + tablespacePath +
101+
". Please ensure you have write permissions to this directory.");
102+
}
103+
104+
return tablespacePath;
105+
}
106+
107+
public boolean isTestTablespaces() {
108+
// If the user explicitly set the value via command line, use that
109+
// Otherwise, use the OS-dependent default
110+
return testTablespaces || getDefaultTablespaceSupport();
111+
}
112+
113+
private static boolean getDefaultTablespaceSupport() {
114+
if (defaultTestTablespaces == null) {
115+
defaultTestTablespaces = determineDefaultTablespaceSupport();
116+
}
117+
return defaultTestTablespaces;
118+
}
77119
}

‎src/sqlancer/postgres/PostgresProvider.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,8 @@ public enum Action implements AbstractAction<PostgresGlobalState> {
126126
LISTEN((g) -> PostgresNotifyGenerator.createListen()), //
127127
UNLISTEN((g) -> PostgresNotifyGenerator.createUnlisten()), //
128128
CREATE_SEQUENCE(PostgresSequenceGenerator::createSequence), //
129-
CREATE_VIEW(PostgresViewGenerator::create), CREATE_TABLESPACE(PostgresTableSpaceGenerator::generate);
129+
CREATE_VIEW(PostgresViewGenerator::create),
130+
CREATE_TABLESPACE(PostgresTableSpaceGenerator::generate);
130131

131132
private final SQLQueryProvider<PostgresGlobalState> sqlQueryProvider;
132133

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,7 @@ public class PostgresTableSpaceGenerator {
1212

1313
public PostgresTableSpaceGenerator(PostgresGlobalState globalState) {
1414
this.globalState = globalState;
15-
errors.addRegexString("ERROR: (?:tablespace )?directory \".*[\\\\/]tablespace[1-5]\" does not exist");
16-
errors.add("ERROR: must be a directory");
15+
errors.addRegexString("ERROR: (?:tablespace )?directory \".*[\\\\/]tablespace\\d+\" does not exist");
1716
errors.add("ERROR: permission denied");
1817
errors.add("ERROR: already exists");
1918
errors.add("ERROR: is not empty");
@@ -23,25 +22,25 @@ public PostgresTableSpaceGenerator(PostgresGlobalState globalState) {
2322
public static SQLQueryAdapter generate(PostgresGlobalState globalState) {
2423
// Skip tablespace generation if the option is disabled
2524
PostgresOptions options = globalState.getDbmsSpecificOptions();
26-
if (!options.testTablespaces) {
25+
if (!options.isTestTablespaces()) {
2726
return null;
2827
}
2928
return new PostgresTableSpaceGenerator(globalState).generateTableSpace();
3029
}
3130

3231
private SQLQueryAdapter generateTableSpace() {
3332
StringBuilder sb = new StringBuilder();
34-
int tableSpaceNum = globalState.getRandomly().getInteger(1, 5);
33+
int tableSpaceNum = globalState.getRandomly().getInteger(1, Integer.MAX_VALUE);
3534

3635
// CREATE TABLESPACE syntax
3736
sb.append("CREATE TABLESPACE ");
3837
sb.append("tablespace");
3938
sb.append(tableSpaceNum);
4039
sb.append(" LOCATION '");
4140

42-
// Get the base path from options and append the tablespace number
41+
// Get the validated base path from options and append the tablespace number
4342
PostgresOptions options = globalState.getDbmsSpecificOptions();
44-
String path = options.tablespacePath + tableSpaceNum;
43+
String path = options.getTablespacePath() + tableSpaceNum;
4544

4645
// Convert backslashes to forward slashes for PostgreSQL
4746
path = path.replace('\\', '/');

0 commit comments

Comments
 (0)