[verifier] Adding ability to run checksum queries on test cluster.

presto-verifier/src/main/java/com/facebook/presto/verifier/annotation/ForChecksum.java

@@ -0,0 +1,31 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.verifier.annotation;
+
+import javax.inject.Qualifier;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@Qualifier
+public @interface ForChecksum
+{
+}

presto-verifier/src/main/java/com/facebook/presto/verifier/framework/AbstractVerification.java

@@ -166,6 +166,11 @@ protected PrestoAction getHelperAction()
         return queryActions.getHelperAction();
     }
 
+    protected PrestoAction getChecksumAction()
+    {
+        return queryActions.getChecksumAction();
+    }
+
     protected boolean isControlEnabled()
     {
         return !skipControl || saveSnapshot;

presto-verifier/src/main/java/com/facebook/presto/verifier/framework/DataVerification.java

@@ -107,20 +107,20 @@ public DataMatchResult verify(
             ChecksumQueryContext controlChecksumQueryContext,
             ChecksumQueryContext testChecksumQueryContext)
     {
-        List<Column> testColumns = getColumns(getHelperAction(), typeManager, test.getObjectName());
+        List<Column> testColumns = getColumns(getChecksumAction(), typeManager, test.getObjectName());
         Query testChecksumQuery = checksumValidator.generateChecksumQuery(test.getObjectName(), testColumns, test.getPartitionsPredicate());
         testChecksumQueryContext.setChecksumQuery(formatSql(testChecksumQuery));
 
         List<Column> controlColumns = null;
         ChecksumResult controlChecksumResult = null;
 
         if (isControlEnabled()) {
-            controlColumns = getColumns(getHelperAction(), typeManager, control.getObjectName());
+            controlColumns = getColumns(getChecksumAction(), typeManager, control.getObjectName());
             Query controlChecksumQuery = checksumValidator.generateChecksumQuery(control.getObjectName(), controlColumns, control.getPartitionsPredicate());
             controlChecksumQueryContext.setChecksumQuery(formatSql(controlChecksumQuery));
 
             QueryResult<ChecksumResult> controlChecksum = callAndConsume(
-                    () -> getHelperAction().execute(controlChecksumQuery, CONTROL_CHECKSUM, ChecksumResult::fromResultSet),
+                    () -> getChecksumAction().execute(controlChecksumQuery, CONTROL_CHECKSUM, ChecksumResult::fromResultSet),
                     stats -> stats.getQueryStats().map(QueryStats::getQueryId).ifPresent(controlChecksumQueryContext::setChecksumQueryId));
             controlChecksumResult = getOnlyElement(controlChecksum.getResults());
 
@@ -151,7 +151,7 @@ else if (QUERY_BANK_MODE.equals(runningMode)) {
         }
 
         QueryResult<ChecksumResult> testChecksum = callAndConsume(
-                () -> getHelperAction().execute(testChecksumQuery, TEST_CHECKSUM, ChecksumResult::fromResultSet),
+                () -> getChecksumAction().execute(testChecksumQuery, TEST_CHECKSUM, ChecksumResult::fromResultSet),
                 stats -> stats.getQueryStats().map(QueryStats::getQueryId).ifPresent(testChecksumQueryContext::setChecksumQueryId));
         ChecksumResult testChecksumResult = getOnlyElement(testChecksum.getResults());
 

presto-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/QueryActions.java

@@ -23,12 +23,14 @@
 public class QueryActions
 {
     private final PrestoAction helperAction;
+    private final PrestoAction checksumAction;
     private final QueryAction controlAction;
     private final QueryAction testAction;
 
-    public QueryActions(PrestoAction helperAction, QueryAction controlAction, QueryAction testAction)
+    public QueryActions(PrestoAction helperAction, PrestoAction checksumAction, QueryAction controlAction, QueryAction testAction)
     {
         this.helperAction = requireNonNull(helperAction, "helperAction is null");
+        this.checksumAction = requireNonNull(helperAction, "checksumAction is null");
         this.controlAction = requireNonNull(controlAction, "controlAction is null");
         this.testAction = requireNonNull(testAction, "testAction is null");
     }
@@ -38,6 +40,11 @@ public PrestoAction getHelperAction()
         return helperAction;
     }
 
+    public PrestoAction getChecksumAction()
+    {
+        return checksumAction;
+    }
+
     public QueryAction getControlAction()
     {
         return controlAction;

presto-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/QueryActionsConfig.java

@@ -26,6 +26,7 @@ public class QueryActionsConfig
     private String controlQueryActionType = JdbcPrestoAction.QUERY_ACTION_TYPE;
     private String testQueryActionType = JdbcPrestoAction.QUERY_ACTION_TYPE;
     private boolean runHelperQueriesOnControl = true;
+    private boolean runChecksumQueriesOnControl;
 
     private Duration metadataTimeout = new Duration(3, MINUTES);
     private Duration checksumTimeout = new Duration(30, MINUTES);
@@ -68,6 +69,18 @@ public QueryActionsConfig setRunHelperQueriesOnControl(boolean runHelperQueriesO
         return this;
     }
 
+    public boolean isRunChecksumQueriesOnControl()
+    {
+        return runChecksumQueriesOnControl;
+    }
+
+    @Config("run-checksum-queries-on-control")
+    public QueryActionsConfig setRunChecksumQueriesOnControl(boolean runChecksumQueriesOnControl)
+    {
+        this.runChecksumQueriesOnControl = runChecksumQueriesOnControl;
+        return this;
+    }
+
     @MinDuration("1s")
     public Duration getMetadataTimeout()
     {

presto-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/QueryActionsModule.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.verifier.prestoaction;
 
 import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
+import com.facebook.presto.verifier.annotation.ForChecksum;
 import com.facebook.presto.verifier.annotation.ForControl;
 import com.facebook.presto.verifier.annotation.ForHelper;
 import com.facebook.presto.verifier.annotation.ForTest;
@@ -84,6 +85,14 @@ protected void setup(Binder binder)
             configBinder(binder).bindConfig(PrestoActionConfig.class, ForHelper.class, "helper");
             binder.bind(PrestoActionFactory.class).annotatedWith(ForHelper.class).toProvider(new JdbcPrestoActionFactoryProvider(ForHelper.class)).in(SINGLETON);
         }
+        if (config.isRunChecksumQueriesOnControl()) {
+            checkArgument(controlQueryActionType.equals(JdbcPrestoAction.QUERY_ACTION_TYPE), "Cannot run checksum queries on control cluster because it is not a presto-jdbc action");
+            binder.bind(PrestoActionFactory.class).annotatedWith(ForChecksum.class).toProvider(new JdbcPrestoActionFactoryProvider(ForControl.class)).in(SINGLETON);
+        }
+        else {
+            checkArgument(testQueryActionType.equals(JdbcPrestoAction.QUERY_ACTION_TYPE), "Cannot run checksum queries on test cluster because it is not a presto-jdbc action");
+            binder.bind(PrestoActionFactory.class).annotatedWith(ForChecksum.class).toProvider(new JdbcPrestoActionFactoryProvider(ForTest.class)).in(SINGLETON);
+        }
     }
 
     public static class JdbcPrestoActionFactoryProvider

presto-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/QueryActionsProvider.java

@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.verifier.prestoaction;
 
+import com.facebook.presto.verifier.annotation.ForChecksum;
 import com.facebook.presto.verifier.annotation.ForControl;
 import com.facebook.presto.verifier.annotation.ForHelper;
 import com.facebook.presto.verifier.annotation.ForTest;
@@ -26,16 +27,19 @@ public class QueryActionsProvider
         implements QueryActionsFactory
 {
     private final PrestoActionFactory helpActionFactory;
+    private final PrestoActionFactory checksumActionFactory;
     private final QueryActionFactory controlActionFactory;
     private final QueryActionFactory testActionFactory;
 
     @Inject
     public QueryActionsProvider(
             @ForHelper PrestoActionFactory helpActionFactory,
+            @ForChecksum PrestoActionFactory checksumActionFactory,
             @ForControl QueryActionFactory controlActionFactory,
             @ForTest QueryActionFactory testActionFactory)
     {
         this.helpActionFactory = requireNonNull(helpActionFactory, "helpActionFactory is null");
+        this.checksumActionFactory = requireNonNull(checksumActionFactory, "checksumActionFactory is null");
         this.controlActionFactory = requireNonNull(controlActionFactory, "controlActionFactory is null");
         this.testActionFactory = requireNonNull(testActionFactory, "testActionFactory is null");
     }
@@ -44,6 +48,7 @@ public QueryActions create(SourceQuery sourceQuery, VerificationContext verifica
     {
         return new QueryActions(
                 helpActionFactory.create(sourceQuery.getControlConfiguration(), verificationContext),
+                checksumActionFactory.create(sourceQuery.getTestConfiguration(), verificationContext),
                 controlActionFactory.create(sourceQuery.getControlConfiguration(), verificationContext),
                 testActionFactory.create(sourceQuery.getTestConfiguration(), verificationContext));
     }