Use generic TLP where oracle for Postgres

sqlancer · Apr 11, 2024 · 66a2fe6 · 66a2fe6
1 parent aaf67ee
commit 66a2fe6
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 31 deletions.
diff --git a/src/sqlancer/citus/oracle/tlp/CitusTLPWhereOracle.java b/src/sqlancer/citus/oracle/tlp/CitusTLPWhereOracle.java
@@ -1,14 +1,20 @@
 package sqlancer.citus.oracle.tlp;
 
 import java.sql.SQLException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 
+import sqlancer.ComparatorHelper;
+import sqlancer.Randomly;
 import sqlancer.citus.CitusGlobalState;
 import sqlancer.citus.gen.CitusCommon;
 import sqlancer.postgres.PostgresGlobalState;
-import sqlancer.postgres.oracle.tlp.PostgresTLPWhereOracle;
+import sqlancer.postgres.PostgresVisitor;
+import sqlancer.postgres.oracle.tlp.PostgresTLPBase;
 
-public class CitusTLPWhereOracle extends PostgresTLPWhereOracle {
+public class CitusTLPWhereOracle extends PostgresTLPBase {
 
     private final CitusTLPBase citusTLPBase;
 
@@ -32,4 +38,25 @@ public void check() throws SQLException {
         whereCheck();
         state.setDefaultAllowedFunctionTypes();
     }
+
+    void whereCheck() throws SQLException {
+        if (Randomly.getBooleanWithRatherLowProbability()) {
+            select.setOrderByClauses(gen.generateOrderBys());
+        }
+        String originalQueryString = PostgresVisitor.asString(select);
+        List<String> resultSet = ComparatorHelper.getResultSetFirstColumnAsString(originalQueryString, errors, state);
+
+        select.setOrderByClauses(Collections.emptyList());
+        select.setWhereClause(predicate);
+        String firstQueryString = PostgresVisitor.asString(select);
+        select.setWhereClause(negatedPredicate);
+        String secondQueryString = PostgresVisitor.asString(select);
+        select.setWhereClause(isNullPredicate);
+        String thirdQueryString = PostgresVisitor.asString(select);
+        List<String> combinedString = new ArrayList<>();
+        List<String> secondResultSet = ComparatorHelper.getCombinedResultSet(firstQueryString, secondQueryString,
+                thirdQueryString, combinedString, Randomly.getBoolean(), state, errors);
+        ComparatorHelper.assumeResultSetsAreEqual(resultSet, secondResultSet, originalQueryString, combinedString,
+                state);
+    }
 }
diff --git a/src/sqlancer/postgres/oracle/tlp/PostgresTLPWhereOracle.java b/src/sqlancer/postgres/oracle/tlp/PostgresTLPWhereOracle.java
@@ -1,45 +1,40 @@
 package sqlancer.postgres.oracle.tlp;
 
 import java.sql.SQLException;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
 
-import sqlancer.ComparatorHelper;
-import sqlancer.Randomly;
+import sqlancer.common.oracle.TLPWhereOracle;
+import sqlancer.common.oracle.TestOracle;
+import sqlancer.common.query.ExpectedErrors;
 import sqlancer.postgres.PostgresGlobalState;
-import sqlancer.postgres.PostgresVisitor;
+import sqlancer.postgres.PostgresSchema;
+import sqlancer.postgres.PostgresSchema.PostgresColumn;
+import sqlancer.postgres.PostgresSchema.PostgresTable;
+import sqlancer.postgres.ast.PostgresExpression;
+import sqlancer.postgres.ast.PostgresJoin;
+import sqlancer.postgres.ast.PostgresSelect;
+import sqlancer.postgres.gen.PostgresCommon;
+import sqlancer.postgres.gen.PostgresExpressionGenerator;
 
-public class PostgresTLPWhereOracle extends PostgresTLPBase {
+public class PostgresTLPWhereOracle implements TestOracle<PostgresGlobalState> {
+
+    private final TLPWhereOracle<PostgresSelect, PostgresJoin, PostgresExpression, PostgresSchema, PostgresTable, PostgresColumn, PostgresGlobalState> oracle;
 
     public PostgresTLPWhereOracle(PostgresGlobalState state) {
-        super(state);
+        PostgresExpressionGenerator gen = new PostgresExpressionGenerator(state);
+        ExpectedErrors expectedErrors = ExpectedErrors.newErrors().with(PostgresCommon.getCommonExpressionErrors())
+                .with(PostgresCommon.getCommonFetchErrors()).withRegex(PostgresCommon.getCommonExpressionRegexErrors())
+                .build();
+
+        this.oracle = new TLPWhereOracle<>(state, gen, expectedErrors);
     }
 
     @Override
     public void check() throws SQLException {
-        super.check();
-        whereCheck();
+        oracle.check();
     }
 
-    protected void whereCheck() throws SQLException {
-        if (Randomly.getBooleanWithRatherLowProbability()) {
-            select.setOrderByClauses(gen.generateOrderBy());
-        }
-        String originalQueryString = PostgresVisitor.asString(select);
-        List<String> resultSet = ComparatorHelper.getResultSetFirstColumnAsString(originalQueryString, errors, state);
-
-        select.setOrderByClauses(Collections.emptyList());
-        select.setWhereClause(predicate);
-        String firstQueryString = PostgresVisitor.asString(select);
-        select.setWhereClause(negatedPredicate);
-        String secondQueryString = PostgresVisitor.asString(select);
-        select.setWhereClause(isNullPredicate);
-        String thirdQueryString = PostgresVisitor.asString(select);
-        List<String> combinedString = new ArrayList<>();
-        List<String> secondResultSet = ComparatorHelper.getCombinedResultSet(firstQueryString, secondQueryString,
-                thirdQueryString, combinedString, Randomly.getBoolean(), state, errors);
-        ComparatorHelper.assumeResultSetsAreEqual(resultSet, secondResultSet, originalQueryString, combinedString,
-                state);
+    @Override
+    public String getLastQueryString() {
+        return oracle.getLastQueryString();
     }
 }