Change setFromTables to setFromList for SQLite and MariaDB

malwaregarry · malwaregarry · commit ab066addf704 · 2024-03-25T10:40:00.000+08:00
diff --git a/src/sqlancer/common/ast/newast/Select.java b/src/sqlancer/common/ast/newast/Select.java
@@ -8,8 +8,6 @@
 public interface Select<J extends Join<E, T, C>, E extends Expression<C>, T extends AbstractTable<C, ?, ?>, C extends AbstractTableColumn<?, ?>>
         extends Expression<C> {
 
-    void setFromTables(List<E> fromTables);
-
     List<E> getFromList();
 
     void setFromList(List<E> fromList);
diff --git a/src/sqlancer/mariadb/ast/MariaDBSelectStatement.java b/src/sqlancer/mariadb/ast/MariaDBSelectStatement.java
@@ -4,7 +4,6 @@
 import java.util.List;
 
 import sqlancer.common.ast.SelectBase;
-import sqlancer.mariadb.MariaDBSchema.MariaDBTable;
 
 public class MariaDBSelectStatement extends SelectBase<MariaDBExpression> implements MariaDBExpression {
 
@@ -14,7 +13,6 @@ public enum MariaDBSelectType {
 
     private List<MariaDBExpression> groupBys = new ArrayList<>();
     private List<MariaDBExpression> columns = new ArrayList<>();
-    private List<MariaDBTable> tables = new ArrayList<>();
     private MariaDBSelectType selectType = MariaDBSelectType.ALL;
     private MariaDBExpression whereCondition;
 
@@ -28,10 +26,6 @@ public void setFetchColumns(List<MariaDBExpression> columns) {
 
     }
 
-    public void setFromTables(List<MariaDBTable> tables) {
-        this.tables = tables;
-    }
-
     public void setSelectType(MariaDBSelectType selectType) {
         this.selectType = selectType;
     }
@@ -53,10 +47,6 @@ public MariaDBSelectType getSelectType() {
         return selectType;
     }
 
-    public List<MariaDBTable> getTables() {
-        return tables;
-    }
-
     public MariaDBExpression getWhereCondition() {
         return whereCondition;
     }
diff --git a/src/sqlancer/mariadb/ast/MariaDBStringVisitor.java b/src/sqlancer/mariadb/ast/MariaDBStringVisitor.java
@@ -1,7 +1,6 @@
 package sqlancer.mariadb.ast;
 
 import java.util.List;
-import java.util.stream.Collectors;
 
 public class MariaDBStringVisitor extends MariaDBVisitor {
 
@@ -41,8 +40,13 @@ public void visit(MariaDBSelectStatement s) {
             visit(column);
         }
         sb.append(" FROM ");
-        sb.append(s.getTables().stream().map(t -> t.getName()).collect(Collectors.joining(", ")));
 
+        for (int j = 0; j < s.getFromList().size(); j++) {
+            if (j != 0) {
+                sb.append(", ");
+            }
+            visit(s.getFromList().get(j));
+        }
         for (MariaDBExpression j : s.getJoinList()) {
             visit(j);
         }
@@ -167,4 +171,9 @@ public void visit(MariaDBJoin join) {
             visit(join.getOnClause());
         }
     }
+
+    @Override
+    public void visit(MariaDBTableReference ref) {
+        sb.append(ref.getTable().getName());
+    }
 }
diff --git a/src/sqlancer/mariadb/ast/MariaDBTableReference.java b/src/sqlancer/mariadb/ast/MariaDBTableReference.java
@@ -0,0 +1,16 @@
+package sqlancer.mariadb.ast;
+
+import sqlancer.mariadb.MariaDBSchema.MariaDBTable;
+
+public class MariaDBTableReference implements MariaDBExpression {
+
+    private final MariaDBTable table;
+
+    public MariaDBTableReference(MariaDBTable table) {
+        this.table = table;
+    }
+
+    public MariaDBTable getTable() {
+        return table;
+    }
+}
diff --git a/src/sqlancer/mariadb/ast/MariaDBVisitor.java b/src/sqlancer/mariadb/ast/MariaDBVisitor.java
@@ -24,6 +24,8 @@ public abstract class MariaDBVisitor {
 
     public abstract void visit(MariaDBJoin join);
 
+    public abstract void visit(MariaDBTableReference join);
+
     public void visit(MariaDBExpression expr) {
         if (expr instanceof MariaDBConstant) {
             visit((MariaDBConstant) expr);
@@ -47,6 +49,8 @@ public void visit(MariaDBExpression expr) {
             visit((MariaDBInOperation) expr);
         } else if (expr instanceof MariaDBJoin) {
             visit((MariaDBJoin) expr);
+        } else if (expr instanceof MariaDBTableReference) {
+            visit((MariaDBTableReference) expr);
         } else {
             throw new AssertionError(expr.getClass());
         }
diff --git a/src/sqlancer/mariadb/oracle/MariaDBDQPOracle.java b/src/sqlancer/mariadb/oracle/MariaDBDQPOracle.java
@@ -17,6 +17,7 @@
 import sqlancer.mariadb.ast.MariaDBExpression;
 import sqlancer.mariadb.ast.MariaDBJoin;
 import sqlancer.mariadb.ast.MariaDBSelectStatement;
+import sqlancer.mariadb.ast.MariaDBTableReference;
 import sqlancer.mariadb.ast.MariaDBVisitor;
 import sqlancer.mariadb.gen.MariaDBExpressionGenerator;
 import sqlancer.mariadb.gen.MariaDBSetGenerator;
@@ -60,7 +61,8 @@ public void check() throws Exception {
         select.setJoinList(joinExpressions.stream().map(j -> (MariaDBExpression) j).collect(Collectors.toList()));
 
         // Set the from clause from the tables that are not used in the join.
-        select.setFromTables(tables.getTables());
+        select.setFromList(
+                tables.getTables().stream().map(t -> new MariaDBTableReference(t)).collect(Collectors.toList()));
 
         // Get the result of the first query
         String originalQueryString = MariaDBVisitor.asString(select);
diff --git a/src/sqlancer/mariadb/oracle/MariaDBNoRECOracle.java b/src/sqlancer/mariadb/oracle/MariaDBNoRECOracle.java
@@ -23,6 +23,7 @@
 import sqlancer.mariadb.ast.MariaDBPostfixUnaryOperation.MariaDBPostfixUnaryOperator;
 import sqlancer.mariadb.ast.MariaDBSelectStatement;
 import sqlancer.mariadb.ast.MariaDBSelectStatement.MariaDBSelectType;
+import sqlancer.mariadb.ast.MariaDBTableReference;
 import sqlancer.mariadb.ast.MariaDBText;
 import sqlancer.mariadb.ast.MariaDBVisitor;
 import sqlancer.mariadb.gen.MariaDBExpressionGenerator;
@@ -76,7 +77,7 @@ private int getUnoptimizedQuery(MariaDBTable randomTable, MariaDBExpression rand
                 randomWhereCondition);
         MariaDBText asText = new MariaDBText(isTrue, " as count", false);
         select.setFetchColumns(Arrays.asList(asText));
-        select.setFromTables(Arrays.asList(randomTable));
+        select.setFromList(Arrays.asList(new MariaDBTableReference(randomTable)));
         select.setSelectType(MariaDBSelectType.ALL);
         int secondCount = 0;
 
@@ -103,7 +104,7 @@ private int getOptimizedQuery(MariaDBTable randomTable, MariaDBExpression random
                 new MariaDBColumnName(new MariaDBColumn("*", MariaDBDataType.INT, false, 0)),
                 MariaDBAggregateFunction.COUNT);
         select.setFetchColumns(Arrays.asList(aggr));
-        select.setFromTables(Arrays.asList(randomTable));
+        select.setFromList(Arrays.asList(new MariaDBTableReference(randomTable)));
         select.setWhereClause(randomWhereCondition);
         select.setSelectType(MariaDBSelectType.ALL);
         int firstCount;
diff --git a/src/sqlancer/sqlite3/ast/SQLite3Select.java b/src/sqlancer/sqlite3/ast/SQLite3Select.java
@@ -46,10 +46,6 @@ public void setSelectType(SelectType fromOptions) {
         this.setFromOptions(fromOptions);
     }
 
-    public void setFromTables(List<SQLite3Expression> fromTables) {
-        this.setFromList(fromTables);
-    }
-
     public SelectType getFromOptions() {
         return fromOptions;
     }
diff --git a/src/sqlancer/sqlite3/oracle/SQLite3NoRECOracle.java b/src/sqlancer/sqlite3/oracle/SQLite3NoRECOracle.java
@@ -80,7 +80,7 @@ public void check() throws SQLException {
         List<Join> joinStatements = gen.getRandomJoinClauses(tables);
         List<SQLite3Expression> tableRefs = SQLite3Common.getTableRefs(tables, s);
         SQLite3Select select = new SQLite3Select();
-        select.setFromTables(tableRefs);
+        select.setFromList(tableRefs);
         select.setJoinClauses(joinStatements);
 
         Function<SQLite3GlobalState, Integer> optimizedQuery = getOptimizedQuery(select, randomWhereCondition);
diff --git a/src/sqlancer/sqlite3/oracle/SQLite3PivotedQuerySynthesisOracle.java b/src/sqlancer/sqlite3/oracle/SQLite3PivotedQuerySynthesisOracle.java
@@ -75,7 +75,7 @@ public SQLite3Select getQuery() throws SQLException {
                 .filter(c -> !SQLite3Schema.ROWID_STRINGS.contains(c.getName())).collect(Collectors.toList());
         List<Join> joinStatements = getJoinStatements(globalState, tables, columnsWithoutRowid);
         selectStatement.setJoinClauses(joinStatements);
-        selectStatement.setFromTables(SQLite3Common.getTableRefs(tables, globalState.getSchema()));
+        selectStatement.setFromList(SQLite3Common.getTableRefs(tables, globalState.getSchema()));
 
         fetchColumns = Randomly.nonEmptySubset(columnsWithoutRowid);
         List<SQLite3Table> allTables = new ArrayList<>();
diff --git a/src/sqlancer/sqlite3/oracle/tlp/SQLite3TLPBase.java b/src/sqlancer/sqlite3/oracle/tlp/SQLite3TLPBase.java
@@ -48,7 +48,7 @@ public void check() throws SQLException {
         List<Join> joinStatements = gen.getRandomJoinClauses(tables);
         List<SQLite3Expression> tableRefs = SQLite3Common.getTableRefs(tables, s);
         select.setJoinClauses(joinStatements.stream().collect(Collectors.toList()));
-        select.setFromTables(tableRefs);
+        select.setFromList(tableRefs);
         select.setWhereClause(null);
     }
 
diff --git a/src/sqlancer/sqlite3/oracle/tlp/SQLite3TLPHavingOracle.java b/src/sqlancer/sqlite3/oracle/tlp/SQLite3TLPHavingOracle.java
@@ -57,7 +57,7 @@ public void check() throws SQLException {
         List<SQLite3Expression> from = SQLite3Common.getTableRefs(tables, state.getSchema());
         select.setJoinClauses(joinStatements);
         select.setSelectType(SelectType.ALL);
-        select.setFromTables(from);
+        select.setFromList(from);
         // TODO order by?
         select.setGroupByClause(groupByColumns);
         select.setHavingClause(null);
diff --git a/test/sqlancer/dbms/TestMariaDB.java b/test/sqlancer/dbms/TestMariaDB.java
@@ -3,12 +3,40 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
+import java.util.List;
+import java.util.stream.Collectors;
+
 import org.junit.jupiter.api.Test;
 
 import sqlancer.Main;
+import sqlancer.mariadb.MariaDBSchema;
+import sqlancer.mariadb.ast.MariaDBColumnName;
+import sqlancer.mariadb.ast.MariaDBSelectStatement;
+import sqlancer.mariadb.ast.MariaDBTableReference;
+import sqlancer.mariadb.ast.MariaDBVisitor;
 
 public class TestMariaDB {
 
+    @Test
+    public void testSelectAsString() {
+        MariaDBSchema.MariaDBColumn c0 = new MariaDBSchema.MariaDBColumn("c0", MariaDBSchema.MariaDBDataType.INT, true,
+                0);
+        MariaDBSchema.MariaDBColumn c1 = new MariaDBSchema.MariaDBColumn("c1", MariaDBSchema.MariaDBDataType.INT, false,
+                0);
+        List<MariaDBSchema.MariaDBColumn> columns = List.of(c0, c1);
+        List<MariaDBSchema.MariaDBIndex> indices = List.of();
+        MariaDBSchema.MariaDBTable t1 = new MariaDBSchema.MariaDBTable("t1", columns, indices,
+                MariaDBSchema.MariaDBTable.MariaDBEngine.INNO_DB);
+        MariaDBSchema.MariaDBTables tables = new MariaDBSchema.MariaDBTables(List.of(t1));
+
+        MariaDBSelectStatement select = new MariaDBSelectStatement();
+        select.setFetchColumns(tables.getColumns().stream().map(MariaDBColumnName::new).collect(Collectors.toList()));
+        select.setFromList(tables.getTables().stream().map(MariaDBTableReference::new).collect(Collectors.toList()));
+
+        String selectString = MariaDBVisitor.asString(select);
+        assertEquals("SELECT c0, c1 FROM t1", selectString);
+    }
+
     @Test
     public void testMariaDB() {
         String mariaDBAvailable = System.getenv("MARIADB_AVAILABLE");

Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,6 @@`
`1`	`1`	`package sqlancer.mariadb.ast;`
`2`	`2`
`3`	`3`	`import java.util.List;`
`4`		`-import java.util.stream.Collectors;`
`5`	`4`
`6`	`5`	`public class MariaDBStringVisitor extends MariaDBVisitor {`
`7`	`6`
`@@ -41,8 +40,13 @@ public void visit(MariaDBSelectStatement s) {`
`41`	`40`	`visit(column);`
`42`	`41`	`}`
`43`	`42`	`sb.append(" FROM ");`
`44`		`- sb.append(s.getTables().stream().map(t -> t.getName()).collect(Collectors.joining(", ")));`
`45`	`43`
	`44`	`+ for (int j = 0; j < s.getFromList().size(); j++) {`
	`45`	`+ if (j != 0) {`
	`46`	`+ sb.append(", ");`
	`47`	`+ }`
	`48`	`+ visit(s.getFromList().get(j));`
	`49`	`+ }`
`46`	`50`	`for (MariaDBExpression j : s.getJoinList()) {`
`47`	`51`	`visit(j);`
`48`	`52`	`}`
`@@ -167,4 +171,9 @@ public void visit(MariaDBJoin join) {`
`167`	`171`	`visit(join.getOnClause());`
`168`	`172`	`}`
`169`	`173`	`}`
	`174`	`+`
	`175`	`+ @Override`
	`176`	`+ public void visit(MariaDBTableReference ref) {`
	`177`	`+ sb.append(ref.getTable().getName());`
	`178`	`+ }`
`170`	`179`	`}`
Original file line number	Diff line number	Diff line change
`@@ -46,10 +46,6 @@ public void setSelectType(SelectType fromOptions) {`
`46`	`46`	`this.setFromOptions(fromOptions);`
`47`	`47`	`}`
`48`	`48`
`49`		`- public void setFromTables(List<SQLite3Expression> fromTables) {`
`50`		`- this.setFromList(fromTables);`
`51`		`- }`
`52`		`-`
`53`	`49`	`public SelectType getFromOptions() {`
`54`	`50`	`return fromOptions;`
`55`	`51`	`}`