Allow more flexible construction of PrestoExceptionClassifier

prestodb · Jan 11, 2020 · ca6a2d3 · ca6a2d3
1 parent a5bb640
commit ca6a2d3
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 23 deletions.
diff --git a/presto-verifier/src/main/java/com/facebook/presto/verifier/PrestoVerifyCommand.java b/presto-verifier/src/main/java/com/facebook/presto/verifier/PrestoVerifyCommand.java
@@ -69,7 +69,7 @@ public List<Class<? extends Predicate<SourceQuery>>> getCustomQueryFilterClasses
     @Override
     public SqlExceptionClassifier getSqlExceptionClassifier()
     {
-        return new PrestoExceptionClassifier(ImmutableSet.of(), ImmutableSet.of());
+        return PrestoExceptionClassifier.createDefault();
     }
 
     @Override

diff --git a/...er/src/main/java/com/facebook/presto/verifier/prestoaction/PrestoExceptionClassifier.java b/...er/src/main/java/com/facebook/presto/verifier/prestoaction/PrestoExceptionClassifier.java
@@ -63,15 +63,15 @@
 public class PrestoExceptionClassifier
         implements SqlExceptionClassifier
 {
-    private static final Set<ErrorCodeSupplier> DEFAULT_ERRORS = ImmutableSet.<ErrorCodeSupplier>builder()
+    public static final Set<ErrorCodeSupplier> DEFAULT_ERRORS = ImmutableSet.<ErrorCodeSupplier>builder()
             .addAll(asList(StandardErrorCode.values()))
             .addAll(asList(MetastoreErrorCode.values()))
             .addAll(asList(HiveErrorCode.values()))
             .addAll(asList(JdbcErrorCode.values()))
             .addAll(asList(ThriftErrorCode.values()))
             .build();
 
-    private static final Set<ErrorCodeSupplier> DEFAULT_RETRYABLE_ERRORS = ImmutableSet.of(
+    public static final Set<ErrorCodeSupplier> DEFAULT_RETRYABLE_ERRORS = ImmutableSet.of(
             // From StandardErrorCode
             NO_NODES_AVAILABLE,
             REMOTE_TASK_ERROR,
@@ -104,20 +104,21 @@ public class PrestoExceptionClassifier
     private final Map<Integer, ErrorCodeSupplier> errorByCode;
     private final Set<ErrorCodeSupplier> retryableErrors;
 
-    public PrestoExceptionClassifier(
-            Set<ErrorCodeSupplier> additionalErrors,
-            Set<ErrorCodeSupplier> additionalRetryableErrors)
+    private PrestoExceptionClassifier(Set<ErrorCodeSupplier> recognizedErrors, Set<ErrorCodeSupplier> retryableErrors)
     {
-        this.errorByCode = ImmutableSet.<ErrorCodeSupplier>builder()
-                .addAll(DEFAULT_ERRORS)
-                .addAll(additionalErrors)
-                .build()
-                .stream()
+        this.errorByCode = recognizedErrors.stream()
                 .collect(toImmutableMap(errorCode -> errorCode.toErrorCode().getCode(), identity()));
-        this.retryableErrors = ImmutableSet.<ErrorCodeSupplier>builder()
-                .addAll(DEFAULT_RETRYABLE_ERRORS)
-                .addAll(additionalRetryableErrors)
-                .build();
+        this.retryableErrors = ImmutableSet.copyOf(retryableErrors);
+    }
+
+    public static PrestoExceptionClassifier createDefault()
+    {
+        return create(DEFAULT_ERRORS, DEFAULT_RETRYABLE_ERRORS);
+    }
+
+    public static PrestoExceptionClassifier create(Set<ErrorCodeSupplier> recognizedErrors, Set<ErrorCodeSupplier> retryableErrors)
+    {
+        return new PrestoExceptionClassifier(recognizedErrors, retryableErrors);
     }
 
     public QueryException createException(QueryStage queryStage, Optional<QueryStats> queryStats, SQLException cause)

diff --git a/...o-verifier/src/test/java/com/facebook/presto/verifier/framework/TestDataVerification.java b/...o-verifier/src/test/java/com/facebook/presto/verifier/framework/TestDataVerification.java
@@ -35,7 +35,6 @@
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
@@ -97,7 +96,7 @@ private DataVerification createVerification(String controlQuery, String testQuer
         VerificationContext verificationContext = new VerificationContext();
         RetryConfig retryConfig = new RetryConfig();
         PrestoAction prestoAction = new JdbcPrestoAction(
-                new PrestoExceptionClassifier(ImmutableSet.of(), ImmutableSet.of()),
+                PrestoExceptionClassifier.createDefault(),
                 configuration,
                 verificationContext,
                 new PrestoClusterConfig()

diff --git a/...erifier/src/test/java/com/facebook/presto/verifier/prestoaction/TestJdbcPrestoAction.java b/...erifier/src/test/java/com/facebook/presto/verifier/prestoaction/TestJdbcPrestoAction.java
@@ -25,7 +25,6 @@
 import com.facebook.presto.verifier.framework.VerificationContext;
 import com.facebook.presto.verifier.retry.RetryConfig;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -75,7 +74,7 @@ public void setup()
     {
         verificationContext = new VerificationContext();
         prestoAction = new JdbcPrestoAction(
-                new PrestoExceptionClassifier(ImmutableSet.of(), ImmutableSet.of()),
+                PrestoExceptionClassifier.createDefault(),
                 CONFIGURATION,
                 verificationContext,
                 new PrestoClusterConfig()

diff --git a/...rc/test/java/com/facebook/presto/verifier/prestoaction/TestPrestoExceptionClassifier.java b/...rc/test/java/com/facebook/presto/verifier/prestoaction/TestPrestoExceptionClassifier.java
@@ -18,7 +18,6 @@
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.verifier.framework.QueryException;
 import com.facebook.presto.verifier.framework.QueryStage;
-import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.Test;
 
 import java.io.EOFException;
@@ -47,7 +46,7 @@ public class TestPrestoExceptionClassifier
     private static final QueryStage QUERY_STAGE = CONTROL_MAIN;
     private static final QueryStats QUERY_STATS = new QueryStats("id", "", false, false, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, Optional.empty());
 
-    private final SqlExceptionClassifier classifier = new PrestoExceptionClassifier(ImmutableSet.of(), ImmutableSet.of());
+    private final SqlExceptionClassifier classifier = PrestoExceptionClassifier.createDefault();
 
     @Test
     public void testNetworkException()

diff --git a/presto-verifier/src/test/java/com/facebook/presto/verifier/rewrite/TestQueryRewriter.java b/presto-verifier/src/test/java/com/facebook/presto/verifier/rewrite/TestQueryRewriter.java
@@ -31,7 +31,6 @@
 import com.facebook.presto.verifier.retry.RetryConfig;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
 import org.intellij.lang.annotations.Language;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -203,7 +202,7 @@ private QueryRewriter getQueryRewriter(QualifiedName prefix)
         return new QueryRewriter(
                 sqlParser,
                 new JdbcPrestoAction(
-                        new PrestoExceptionClassifier(ImmutableSet.of(), ImmutableSet.of()),
+                        PrestoExceptionClassifier.createDefault(),
                         CONFIGURATION,
                         new VerificationContext(),
                         new PrestoClusterConfig()