Adding support for pre and post queries for verifier

presto-docs/src/main/sphinx/release/release-0.124.rst

@@ -6,3 +6,17 @@ General Changes
 ---------------
 
 * The :func:`approx_percentile` aggregation now also accepts an array of percentages
+
+Verifier
+--------
+
+* Add support for pre and post control and test queries
+
+If you're upgrading from 0.123, you need to alter your verifier_queries table
+
+.. code-block:: sql
+
+    ALTER TABLE verifier_queries add column test_postqueries text null;
+    ALTER TABLE verifier_queries add column test_prequeries text null;
+    ALTER TABLE verifier_queries add column control_postqueries text null;
+    ALTER TABLE verifier_queries add column control_prequeries text null;

presto-verifier/src/main/java/com/facebook/presto/verifier/PrestoVerifier.java

@@ -177,15 +177,19 @@ private static List<QueryPair> applyOverrides(final VerifierConfig config, List<
                     Query test = new Query(
                             Optional.ofNullable(config.getTestCatalogOverride()).orElse(input.getTest().getCatalog()),
                             Optional.ofNullable(config.getTestSchemaOverride()).orElse(input.getTest().getSchema()),
+                            input.getTest().getPreQueries(),
                             input.getTest().getQuery(),
+                            input.getTest().getPostQueries(),
                             Optional.ofNullable(config.getTestUsernameOverride()).orElse(input.getTest().getUsername()),
                             Optional.ofNullable(config.getTestPasswordOverride()).orElse(
                                     Optional.ofNullable(input.getTest().getPassword()).orElse(null)),
                             input.getTest().getSessionProperties());
                     Query control = new Query(
                             Optional.ofNullable(config.getControlCatalogOverride()).orElse(input.getControl().getCatalog()),
                             Optional.ofNullable(config.getControlSchemaOverride()).orElse(input.getControl().getSchema()),
+                            input.getControl().getPreQueries(),
                             input.getControl().getQuery(),
+                            input.getControl().getPostQueries(),
                             Optional.ofNullable(config.getControlUsernameOverride()).orElse(input.getControl().getUsername()),
                             Optional.ofNullable(config.getControlPasswordOverride()).orElse(
                                     Optional.ofNullable(input.getControl().getPassword()).orElse(null)),

presto-verifier/src/main/java/com/facebook/presto/verifier/Query.java

@@ -15,22 +15,29 @@
 
 import com.google.common.collect.ImmutableMap;
 
+import java.util.List;
 import java.util.Map;
 
+import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
+
 public class Query
 {
     private final String catalog;
     private final String schema;
+    private final List<String> preQueries;
     private final String query;
+    private final List<String> postQueries;
     private final String username;
     private final String password;
     private final Map<String, String> sessionProperties;
 
-    public Query(String catalog, String schema, String query, String username, String password, Map<String, String> sessionProperties)
+    public Query(String catalog, String schema, List<String> preQueries, String query, List<String> postQueries, String username, String password, Map<String, String> sessionProperties)
     {
         this.catalog = catalog;
         this.schema = schema;
+        this.preQueries = preQueries.stream().map(Query::clean).collect(toImmutableList());
         this.query = clean(query);
+        this.postQueries = postQueries.stream().map(Query::clean).collect(toImmutableList());
         this.username = username;
         this.password = password;
         this.sessionProperties = ImmutableMap.copyOf(sessionProperties);
@@ -46,11 +53,21 @@ public String getSchema()
         return schema;
     }
 
+    public List<String> getPreQueries()
+    {
+        return preQueries;
+    }
+
     public String getQuery()
     {
         return query;
     }
 
+    public List<String> getPostQueries()
+    {
+        return postQueries;
+    }
+
     public String getUsername()
     {
         return username;

presto-verifier/src/main/java/com/facebook/presto/verifier/QueryPairMapper.java

@@ -13,19 +13,22 @@
  */
 package com.facebook.presto.verifier;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import io.airlift.json.JsonCodec;
 import org.skife.jdbi.v2.StatementContext;
 import org.skife.jdbi.v2.tweak.ResultSetMapper;
 
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.util.List;
 import java.util.Map;
 
 public class QueryPairMapper
         implements ResultSetMapper<QueryPair>
 {
-    private static final JsonCodec<Map<String, String>> jsonCodec = JsonCodec.mapJsonCodec(String.class, String.class);
+    private static final JsonCodec<Map<String, String>> propertiesJsonCodec = JsonCodec.mapJsonCodec(String.class, String.class);
+    private static final JsonCodec<List<String>> queriesJsonCodec = JsonCodec.listJsonCodec(String.class);
 
     @Override
     public QueryPair map(int index, ResultSet resultSet, StatementContext statementContext)
@@ -34,15 +37,23 @@ public QueryPair map(int index, ResultSet resultSet, StatementContext statementC
         Map<String, String> sessionProperties = ImmutableMap.of();
         String json = resultSet.getString("session_properties_json");
         if (json != null) {
-            sessionProperties = jsonCodec.fromJson(json);
+            sessionProperties = propertiesJsonCodec.fromJson(json);
         }
 
         return new QueryPair(
                 resultSet.getString("suite"),
                 resultSet.getString("name"),
-                new Query(resultSet.getString("test_catalog"), resultSet.getString("test_schema"), resultSet.getString("test_query"),
+                new Query(resultSet.getString("test_catalog"), resultSet.getString("test_schema"), fromJsonString(resultSet.getString("test_prequeries")), resultSet.getString("test_query"),
+                        fromJsonString(resultSet.getString("test_postqueries")),
                         resultSet.getString("test_username"), resultSet.getString("test_password"), sessionProperties),
-                new Query(resultSet.getString("control_catalog"), resultSet.getString("control_schema"), resultSet.getString("control_query"),
+                new Query(resultSet.getString("control_catalog"), resultSet.getString("control_schema"), fromJsonString(resultSet.getString("control_prequeries")), resultSet.getString("control_query"),
+                        fromJsonString(resultSet.getString("control_postqueries")),
                         resultSet.getString("control_username"), resultSet.getString("control_password"), sessionProperties));
     }
+
+    private static List<String> fromJsonString(String jsonString)
+            throws SQLException
+    {
+        return jsonString == null ? ImmutableList.of() : queriesJsonCodec.fromJson(jsonString);
+    }
 }

presto-verifier/src/main/java/com/facebook/presto/verifier/QueryResult.java

@@ -24,7 +24,7 @@ public class QueryResult
 {
     public enum State
     {
-        INVALID, FAILED, SUCCESS, TOO_MANY_ROWS, TIMEOUT
+        INVALID, FAILED, SUCCESS, TOO_MANY_ROWS, TIMEOUT, FAILED_TO_SETUP, FAILED_TO_TEARDOWN
     }
 
     private final State state;

presto-verifier/src/main/java/com/facebook/presto/verifier/Validator.java

@@ -47,6 +47,7 @@
 import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.regex.Pattern;
 
 import static com.facebook.presto.verifier.QueryResult.State;
@@ -182,36 +183,61 @@ public boolean isDeterministic()
 
     private boolean validate()
     {
-        controlResult = executeQuery(controlGateway, controlUsername, controlPassword, queryPair.getControl(), controlTimeout, sessionProperties);
+        controlResult = executeQueryControl();
 
         // query has too many rows. Consider blacklisting.
         if (controlResult.getState() == State.TOO_MANY_ROWS) {
             testResult = new QueryResult(State.INVALID, null, null, ImmutableList.<List<Object>>of());
             return false;
         }
         // query failed in the control
-        else if (controlResult.getState() != State.SUCCESS) {
+        if (controlResult.getState() != State.SUCCESS) {
             testResult = new QueryResult(State.INVALID, null, null, ImmutableList.<List<Object>>of());
             return true;
         }
 
-        testResult = executeQuery(testGateway, testUsername, testPassword, queryPair.getTest(), testTimeout, sessionProperties);
+        testResult = executeQueryTest();
 
-        if ((controlResult.getState() != State.SUCCESS) || (testResult.getState() != State.SUCCESS)) {
+        if (controlResult.getState() != State.SUCCESS || testResult.getState() != State.SUCCESS) {
             return false;
         }
 
         if (!checkCorrectness) {
             return true;
         }
 
-        if (resultsMatch(controlResult, testResult, precision)) {
-            return true;
+        return resultsMatch(controlResult, testResult, precision) || checkForDeterministicAndRerunTestQueriesIfNeeded();
+    }
+
+    private QueryResult tearDown(Query query, Function<String, QueryResult> executor)
+    {
+        for (String postqueryString : query.getPostQueries()) {
+            QueryResult queryResult = executor.apply(postqueryString);
+            if (queryResult.getState() != State.SUCCESS) {
+                return new QueryResult(State.FAILED_TO_TEARDOWN, queryResult.getException(), queryResult.getDuration(), ImmutableList.<List<Object>>of());
+            }
         }
 
+        return new QueryResult(State.SUCCESS, null, null, ImmutableList.of());
+    }
+
+    private QueryResult setup(Query query, Function<String, QueryResult> executor)
+    {
+        for (String prequeryString : query.getPreQueries()) {
+            QueryResult queryResult = executor.apply(prequeryString);
+            if (queryResult.getState() != State.SUCCESS) {
+                return new QueryResult(State.FAILED_TO_SETUP, queryResult.getException(), queryResult.getDuration(), ImmutableList.<List<Object>>of());
+            }
+        }
+
+        return new QueryResult(State.SUCCESS, null, null, ImmutableList.of());
+    }
+
+    private boolean checkForDeterministicAndRerunTestQueriesIfNeeded()
+    {
         // check if the control query is deterministic
         for (int i = 0; i < 3; i++) {
-            QueryResult results = executeQuery(controlGateway, controlUsername, controlPassword, queryPair.getControl(), controlTimeout, sessionProperties);
+            QueryResult results = executeQueryControl();
             if (results.getState() != State.SUCCESS) {
                 return false;
             }
@@ -225,7 +251,7 @@ else if (controlResult.getState() != State.SUCCESS) {
         // Re-run the test query to confirm that the results don't match, in case there was caching on the test tier,
         // but require that it matches 3 times in a row to rule out a non-deterministic correctness bug.
         for (int i = 0; i < 3; i++) {
-            testResult = executeQuery(testGateway, testUsername, testPassword, queryPair.getTest(), testTimeout, sessionProperties);
+            testResult = executeQueryTest();
             if (testResult.getState() != State.SUCCESS) {
                 return false;
             }
@@ -238,6 +264,52 @@ else if (controlResult.getState() != State.SUCCESS) {
         return true;
     }
 
+    private QueryResult executeQueryTest()
+    {
+        Query query = queryPair.getTest();
+        QueryResult queryResult = new QueryResult(State.INVALID, null, null, ImmutableList.<List<Object>>of());
+        try {
+            // startup
+            queryResult = setup(query, testPrequery -> executeQuery(testGateway, testUsername, testPassword, queryPair.getTest(), testPrequery, testTimeout, sessionProperties));
+
+            // if startup is successful -> execute query
+            if (queryResult.getState() == State.SUCCESS) {
+                queryResult = executeQuery(testGateway, testUsername, testPassword, queryPair.getTest(), query.getQuery(), testTimeout, sessionProperties);
+            }
+        }
+        finally {
+            // teardown no matter what
+            QueryResult tearDownResult = tearDown(query, testPostquery -> executeQuery(testGateway, testUsername, testPassword, queryPair.getTest(), testPostquery, testTimeout, sessionProperties));
+
+            // if teardown is not successful the query fails
+            queryResult = tearDownResult.getState() == State.SUCCESS ? queryResult : tearDownResult;
+        }
+        return queryResult;
+    }
+
+    private QueryResult executeQueryControl()
+    {
+        Query query = queryPair.getControl();
+        QueryResult queryResult = new QueryResult(State.INVALID, null, null, ImmutableList.<List<Object>>of());
+        try {
+            // startup
+            queryResult = setup(query, controlPrequery -> executeQuery(controlGateway, controlUsername, controlPassword, queryPair.getControl(), controlPrequery, controlTimeout, sessionProperties));
+
+            // if startup is successful -> execute query
+            if (queryResult.getState() == State.SUCCESS) {
+                queryResult = executeQuery(controlGateway, controlUsername, controlPassword, queryPair.getControl(), query.getQuery(), controlTimeout, sessionProperties);
+            }
+        }
+        finally {
+            // teardown no matter what
+            QueryResult tearDownResult = tearDown(query, controlPostquery -> executeQuery(controlGateway, controlUsername, controlPassword, queryPair.getControl(), controlPostquery, controlTimeout, sessionProperties));
+
+            // if teardown is not successful the query fails
+            queryResult = tearDownResult.getState() == State.SUCCESS ? queryResult : tearDownResult;
+        }
+        return queryResult;
+    }
+
     public QueryPair getQueryPair()
     {
         return queryPair;
@@ -253,7 +325,7 @@ public QueryResult getTestResult()
         return testResult;
     }
 
-    private QueryResult executeQuery(String url, String username, String password, Query query, Duration timeout, Map<String, String> sessionProperties)
+    private QueryResult executeQuery(String url, String username, String password, Query query, String sql, Duration timeout, Map<String, String> sessionProperties)
     {
         try (Connection connection = DriverManager.getConnection(url, username, password)) {
             trySetConnectionProperties(query, connection);
@@ -266,7 +338,6 @@ private QueryResult executeQuery(String url, String username, String password, Q
                 TimeLimiter limiter = new SimpleTimeLimiter();
                 Stopwatch stopwatch = Stopwatch.createStarted();
                 Statement limitedStatement = limiter.newProxy(statement, Statement.class, timeout.toMillis(), TimeUnit.MILLISECONDS);
-                String sql = query.getQuery();
                 if (explainOnly) {
                     sql = "EXPLAIN " + sql;
                 }
@@ -326,7 +397,8 @@ private void trySetConnectionProperties(Query query, Connection connection)
 
     private Callable<List<List<Object>>> getResultSetConverter(final ResultSet resultSet)
     {
-        return new Callable<List<List<Object>>>() {
+        return new Callable<List<List<Object>>>()
+        {
             @Override
             public List<List<Object>> call()
                     throws Exception

presto-verifier/src/main/java/com/facebook/presto/verifier/VerifierDao.java

@@ -27,12 +27,16 @@ public interface VerifierDao
             ", name\n" +
             ", test_catalog\n" +
             ", test_schema\n" +
+            ", test_prequeries\n" +
             ", test_query\n" +
+            ", test_postqueries\n" +
             ", test_username\n" +
             ", test_password\n" +
             ", control_catalog\n" +
             ", control_schema\n" +
+            ", control_prequeries\n" +
             ", control_query\n" +
+            ", control_postqueries\n" +
             ", control_username\n" +
             ", control_password\n" +
             ", session_properties_json\n" +