[HSQLDB] Implement TLP-where generator interface

malwaregarry · malwaregarry · commit fab83b92b14d · 2024-09-16T17:13:03.000+08:00
diff --git a/src/sqlancer/hsqldb/gen/HSQLDBExpressionGenerator.java b/src/sqlancer/hsqldb/gen/HSQLDBExpressionGenerator.java
@@ -8,6 +8,7 @@
 import sqlancer.Randomly;
 import sqlancer.common.ast.BinaryOperatorNode;
 import sqlancer.common.gen.NoRECGenerator;
+import sqlancer.common.gen.TLPWhereGenerator;
 import sqlancer.common.gen.TypedExpressionGenerator;
 import sqlancer.common.schema.AbstractTables;
 import sqlancer.hsqldb.HSQLDBProvider;
@@ -26,7 +27,8 @@
 
 public final class HSQLDBExpressionGenerator extends
         TypedExpressionGenerator<HSQLDBExpression, HSQLDBSchema.HSQLDBColumn, HSQLDBSchema.HSQLDBCompositeDataType>
-        implements NoRECGenerator<HSQLDBSelect, HSQLDBJoin, HSQLDBExpression, HSQLDBTable, HSQLDBColumn> {
+        implements NoRECGenerator<HSQLDBSelect, HSQLDBJoin, HSQLDBExpression, HSQLDBTable, HSQLDBColumn>,
+        TLPWhereGenerator<HSQLDBSelect, HSQLDBJoin, HSQLDBExpression, HSQLDBTable, HSQLDBColumn> {
 
     List<HSQLDBTable> tables;
 
@@ -205,8 +207,7 @@ public List<HSQLDBExpression> generateOrderBys() {
     }
 
     @Override
-    public NoRECGenerator<HSQLDBSelect, HSQLDBJoin, HSQLDBExpression, HSQLDBTable, HSQLDBColumn> setTablesAndColumns(
-            AbstractTables<HSQLDBTable, HSQLDBColumn> tables) {
+    public HSQLDBExpressionGenerator setTablesAndColumns(AbstractTables<HSQLDBTable, HSQLDBColumn> tables) {
         this.columns = tables.getColumns();
         this.tables = tables.getTables();
 
@@ -289,4 +290,13 @@ public String generateUnoptimizedQueryString(HSQLDBSelect select, HSQLDBExpressi
         select.setWhereClause(null);
         return "SELECT SUM(count) FROM (" + select.asString() + ") as res";
     }
+
+    @Override
+    public List<HSQLDBExpression> generateFetchColumns(boolean shouldCreateDummy) {
+        if (shouldCreateDummy) {
+            return List.of(new HSQLDBColumnReference(new HSQLDBSchema.HSQLDBColumn("*", null, null)));
+        }
+        return Randomly
+                .nonEmptySubset(columns.stream().map(c -> new HSQLDBColumnReference(c)).collect(Collectors.toList()));
+    }
 }
