Extract query rewrite from AbstractVerification

prestodb · Sep 10, 2020 · b050fc1 · b050fc1
1 parent aae0157
commit b050fc1
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 23 deletions.
diff --git a/presto-verifier/src/main/java/com/facebook/presto/verifier/event/VerifierQueryEvent.java b/presto-verifier/src/main/java/com/facebook/presto/verifier/event/VerifierQueryEvent.java
@@ -26,6 +26,8 @@
 import java.util.Optional;
 
 import static com.facebook.presto.verifier.event.VerifierQueryEvent.EventStatus.SKIPPED;
+import static com.facebook.presto.verifier.framework.ClusterType.CONTROL;
+import static com.facebook.presto.verifier.framework.ClusterType.TEST;
 import static java.util.Objects.requireNonNull;
 
 @Immutable
@@ -119,11 +121,11 @@ public static VerifierQueryEvent skipped(
                         Optional.of(new QueryInfo(
                                 sourceQuery.getControlConfiguration().getCatalog(),
                                 sourceQuery.getControlConfiguration().getSchema(),
-                                sourceQuery.getControlQuery())),
+                                sourceQuery.getQuery(CONTROL))),
                 new QueryInfo(
                         sourceQuery.getTestConfiguration().getCatalog(),
                         sourceQuery.getTestConfiguration().getSchema(),
-                        sourceQuery.getTestQuery()),
+                        sourceQuery.getQuery(TEST)),
                 Optional.empty(),
                 Optional.empty(),
                 Optional.empty(),

diff --git a/...o-verifier/src/main/java/com/facebook/presto/verifier/framework/AbstractVerification.java b/...o-verifier/src/main/java/com/facebook/presto/verifier/framework/AbstractVerification.java
@@ -70,7 +70,6 @@ public abstract class AbstractVerification
 
     private final QueryActions queryActions;
     private final SourceQuery sourceQuery;
-    private final QueryRewriter queryRewriter;
     private final DeterminismAnalyzer determinismAnalyzer;
     private final FailureResolverManager failureResolverManager;
     private final SqlExceptionClassifier exceptionClassifier;
@@ -87,7 +86,6 @@ public abstract class AbstractVerification
     public AbstractVerification(
             QueryActions queryActions,
             SourceQuery sourceQuery,
-            QueryRewriter queryRewriter,
             DeterminismAnalyzer determinismAnalyzer,
             FailureResolverManager failureResolverManager,
             SqlExceptionClassifier exceptionClassifier,
@@ -96,7 +94,6 @@ public AbstractVerification(
     {
         this.queryActions = requireNonNull(queryActions, "queryActions is null");
         this.sourceQuery = requireNonNull(sourceQuery, "sourceQuery is null");
-        this.queryRewriter = requireNonNull(queryRewriter, "queryRewriter is null");
         this.determinismAnalyzer = requireNonNull(determinismAnalyzer, "determinismAnalyzer is null");
         this.failureResolverManager = requireNonNull(failureResolverManager, "failureResolverManager is null");
         this.exceptionClassifier = requireNonNull(exceptionClassifier, "exceptionClassifier is null");
@@ -110,6 +107,8 @@ public AbstractVerification(
         this.skipControl = verifierConfig.isSkipControl();
     }
 
+    protected abstract QueryBundle getQueryRewrite(ClusterType clusterType);
+
     protected abstract MatchResult verify(QueryBundle control, QueryBundle test, ChecksumQueryContext controlContext, ChecksumQueryContext testContext);
 
     protected PrestoAction getHelperAction()
@@ -148,9 +147,9 @@ public VerificationResult run()
         try {
             // Rewrite queries
             if (!skipControl) {
-                control = Optional.of(queryRewriter.rewriteQuery(sourceQuery.getControlQuery(), CONTROL));
+                control = Optional.of(getQueryRewrite(CONTROL));
             }
-            test = Optional.of(queryRewriter.rewriteQuery(sourceQuery.getTestQuery(), TEST));
+            test = Optional.of(getQueryRewrite(TEST));
 
             // Run control queries
             if (skipControl) {
@@ -318,13 +317,13 @@ private VerificationResult concludeVerification(
                         Optional.empty() :
                         Optional.of(buildQueryInfo(
                                 sourceQuery.getControlConfiguration(),
-                                sourceQuery.getControlQuery(),
+                                sourceQuery.getQuery(CONTROL),
                                 controlChecksumQueryContext,
                                 control,
                                 controlQueryContext)),
                 buildQueryInfo(
                         sourceQuery.getTestConfiguration(),
-                        sourceQuery.getTestQuery(),
+                        sourceQuery.getQuery(TEST),
                         testChecksumQueryContext,
                         test,
                         testQueryContext),

diff --git a/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/DataVerification.java b/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/DataVerification.java
@@ -36,6 +36,7 @@
 public class DataVerification
         extends AbstractVerification
 {
+    private final QueryRewriter queryRewriter;
     private final TypeManager typeManager;
     private final ChecksumValidator checksumValidator;
 
@@ -51,11 +52,18 @@ public DataVerification(
             TypeManager typeManager,
             ChecksumValidator checksumValidator)
     {
-        super(queryActions, sourceQuery, queryRewriter, determinismAnalyzer, failureResolverManager, exceptionClassifier, verificationContext, verifierConfig);
+        super(queryActions, sourceQuery, determinismAnalyzer, failureResolverManager, exceptionClassifier, verificationContext, verifierConfig);
+        this.queryRewriter = requireNonNull(queryRewriter, "queryRewriter is null");
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
         this.checksumValidator = requireNonNull(checksumValidator, "checksumValidator is null");
     }
 
+    @Override
+    protected QueryBundle getQueryRewrite(ClusterType clusterType)
+    {
+        return queryRewriter.rewriteQuery(getSourceQuery().getQuery(clusterType), clusterType);
+    }
+
     @Override
     public MatchResult verify(QueryBundle control, QueryBundle test, ChecksumQueryContext controlContext, ChecksumQueryContext testContext)
     {

diff --git a/...to-verifier/src/main/java/com/facebook/presto/verifier/framework/DeterminismAnalyzer.java b/...to-verifier/src/main/java/com/facebook/presto/verifier/framework/DeterminismAnalyzer.java
@@ -125,7 +125,7 @@ protected DeterminismAnalysis analyze(QueryBundle control, ChecksumResult contro
         Map<QueryBundle, DeterminismAnalysisRun.Builder> queryRuns = new HashMap<>();
         try {
             for (int i = 0; i < maxAnalysisRuns; i++) {
-                QueryBundle queryBundle = queryRewriter.rewriteQuery(sourceQuery.getControlQuery(), CONTROL);
+                QueryBundle queryBundle = queryRewriter.rewriteQuery(sourceQuery.getQuery(CONTROL), CONTROL);
                 DeterminismAnalysisRun.Builder run = determinismAnalysisDetails.addRun().setTableName(queryBundle.getTableName().toString());
                 queryRuns.put(queryBundle, run);
 

diff --git a/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/SourceQuery.java b/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/SourceQuery.java
@@ -19,7 +19,10 @@
 
 import java.util.Objects;
 
+import static com.facebook.presto.verifier.framework.ClusterType.CONTROL;
+import static com.facebook.presto.verifier.framework.ClusterType.TEST;
 import static com.google.common.base.MoreObjects.toStringHelper;
+import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 
 public class SourceQuery
@@ -58,14 +61,10 @@ public String getName()
         return name;
     }
 
-    public String getControlQuery()
+    public String getQuery(ClusterType clusterType)
     {
-        return controlQuery;
-    }
-
-    public String getTestQuery()
-    {
-        return testQuery;
+        checkArgument(clusterType == CONTROL || clusterType == TEST, "Invalid ClusterType: %s", clusterType);
+        return clusterType == CONTROL ? controlQuery : testQuery;
     }
 
     public QueryConfiguration getControlConfiguration()

diff --git a/...to-verifier/src/main/java/com/facebook/presto/verifier/framework/VerificationFactory.java b/...to-verifier/src/main/java/com/facebook/presto/verifier/framework/VerificationFactory.java
@@ -29,6 +29,7 @@
 
 import java.util.Optional;
 
+import static com.facebook.presto.verifier.framework.ClusterType.CONTROL;
 import static com.facebook.presto.verifier.framework.VerifierUtil.PARSING_OPTIONS;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
@@ -70,7 +71,7 @@ public VerificationFactory(
 
     public Verification get(SourceQuery sourceQuery, Optional<VerificationContext> existingContext)
     {
-        QueryType queryType = QueryType.of(sqlParser.createStatement(sourceQuery.getControlQuery(), PARSING_OPTIONS));
+        QueryType queryType = QueryType.of(sqlParser.createStatement(sourceQuery.getQuery(CONTROL), PARSING_OPTIONS));
         switch (queryType.getCategory()) {
             case DATA_PRODUCING:
                 VerificationContext verificationContext = existingContext.map(VerificationContext::createForResubmission)

diff --git a/...to-verifier/src/main/java/com/facebook/presto/verifier/framework/VerificationManager.java b/...to-verifier/src/main/java/com/facebook/presto/verifier/framework/VerificationManager.java
@@ -50,6 +50,8 @@
 import static com.facebook.presto.verifier.event.VerifierQueryEvent.EventStatus.FAILED_RESOLVED;
 import static com.facebook.presto.verifier.event.VerifierQueryEvent.EventStatus.SKIPPED;
 import static com.facebook.presto.verifier.event.VerifierQueryEvent.EventStatus.SUCCEEDED;
+import static com.facebook.presto.verifier.framework.ClusterType.CONTROL;
+import static com.facebook.presto.verifier.framework.ClusterType.TEST;
 import static com.facebook.presto.verifier.framework.QueryType.Category.DATA_PRODUCING;
 import static com.facebook.presto.verifier.framework.SkippedReason.CUSTOM_FILTER;
 import static com.facebook.presto.verifier.framework.SkippedReason.MISMATCHED_QUERY_TYPE;
@@ -175,8 +177,8 @@ private List<SourceQuery> applyOverrides(List<SourceQuery> sourceQueries)
                 .map(sourceQuery -> new SourceQuery(
                         sourceQuery.getSuite(),
                         sourceQuery.getName(),
-                        sourceQuery.getControlQuery(),
-                        sourceQuery.getTestQuery(),
+                        sourceQuery.getQuery(CONTROL),
+                        sourceQuery.getQuery(TEST),
                         sourceQuery.getControlConfiguration().applyOverrides(controlOverrides),
                         sourceQuery.getTestConfiguration().applyOverrides(testOverrides)))
                 .collect(toImmutableList());
@@ -211,8 +213,8 @@ private List<SourceQuery> filterQueryType(List<SourceQuery> sourceQueries)
         ImmutableList.Builder<SourceQuery> selected = ImmutableList.builder();
         for (SourceQuery sourceQuery : sourceQueries) {
             try {
-                QueryType controlQueryType = QueryType.of(sqlParser.createStatement(sourceQuery.getControlQuery(), PARSING_OPTIONS));
-                QueryType testQueryType = QueryType.of(sqlParser.createStatement(sourceQuery.getTestQuery(), PARSING_OPTIONS));
+                QueryType controlQueryType = QueryType.of(sqlParser.createStatement(sourceQuery.getQuery(CONTROL), PARSING_OPTIONS));
+                QueryType testQueryType = QueryType.of(sqlParser.createStatement(sourceQuery.getQuery(TEST), PARSING_OPTIONS));
 
                 if (controlQueryType != testQueryType) {
                     postEvent(VerifierQueryEvent.skipped(sourceQuery.getSuite(), testId, sourceQuery, MISMATCHED_QUERY_TYPE, skipControl));