use reflection to check for mysql transient exception type (#12205)

* use reflection to check for mysql transient exception type

* better

* oops

extensions-core/mysql-metadata-storage/src/main/java/org/apache/druid/metadata/storage/mysql/MySQLConnector.java

@@ -19,11 +19,11 @@
 
 package org.apache.druid.metadata.storage.mysql;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
-import com.mysql.jdbc.exceptions.MySQLTransientException;
 import org.apache.commons.dbcp2.BasicDataSource;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
@@ -35,6 +35,7 @@ import org.skife.jdbi.v2.DBI;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.util.StringMapper;
 
+import javax.annotation.Nullable;
 import java.io.File;
 import java.sql.SQLException;
 
@@ -45,7 +46,10 @@ public class MySQLConnector extends SQLMetadataConnector
   private static final String SERIAL_TYPE = "BIGINT(20) AUTO_INCREMENT";
   private static final String QUOTE_STRING = "`";
   private static final String COLLATION = "CHARACTER SET utf8mb4 COLLATE utf8mb4_bin";
+  private static final String MYSQL_TRANSIENT_EXCEPTION_CLASS_NAME = "com.mysql.jdbc.exceptions.MySQLTransientException";
 
+  @Nullable
+  private final Class<?> myTransientExceptionClass;
   private final DBI dbi;
 
   @Inject
@@ -57,17 +61,13 @@ public class MySQLConnector extends SQLMetadataConnector
   )
   {
     super(config, dbTables);
-    try {
-      log.info("Loading \"MySQL\" metadata connector driver %s", driverConfig.getDriverClassName());
-      Class.forName(driverConfig.getDriverClassName(), false, getClass().getClassLoader());
-    }
-    catch (ClassNotFoundException e) {
-      throw new ISE(e, "Could not find %s on the classpath. The MySQL Connector library is not included in the Druid "
-                   + "distribution but is required to use MySQL. Please download a compatible library (for example "
-                   + "'mysql-connector-java-5.1.48.jar') and place it under 'extensions/mysql-metadata-storage/'. See "
-                   + "https://druid.apache.org/downloads for more details.",
-                    driverConfig.getDriverClassName()
-      );
+    log.info("Loading \"MySQL\" metadata connector driver %s", driverConfig.getDriverClassName());
+    tryLoadDriverClass(driverConfig.getDriverClassName());
+
+    if (driverConfig.getDriverClassName().contains("mysql")) {
+      myTransientExceptionClass = tryLoadDriverClass(MYSQL_TRANSIENT_EXCEPTION_CLASS_NAME);
+    } else {
+      myTransientExceptionClass = null;
     }
 
     final BasicDataSource datasource = getDatasource();
@@ -205,8 +205,7 @@ public class MySQLConnector extends SQLMetadataConnector
   @Override
   protected boolean connectorIsTransientException(Throwable e)
   {
-    return e instanceof MySQLTransientException
-           || (e instanceof SQLException && ((SQLException) e).getErrorCode() == 1317 /* ER_QUERY_INTERRUPTED */);
+    return isTransientException(myTransientExceptionClass, e);
   }
 
   @Override
@@ -241,4 +240,29 @@ public class MySQLConnector extends SQLMetadataConnector
   {
     return dbi;
   }
+
+  private Class<?> tryLoadDriverClass(String className)
+  {
+    try {
+      return Class.forName(className, false, getClass().getClassLoader());
+    }
+    catch (ClassNotFoundException e) {
+      throw new ISE(e, "Could not find %s on the classpath. The MySQL Connector library is not included in the Druid "
+                       + "distribution but is required to use MySQL. Please download a compatible library (for example "
+                       + "'mysql-connector-java-5.1.48.jar') and place it under 'extensions/mysql-metadata-storage/'. See "
+                       + "https://druid.apache.org/downloads for more details.",
+                    className
+      );
+    }
+  }
+
+  @VisibleForTesting
+  static boolean isTransientException(@Nullable Class<?> driverClazz, Throwable e)
+  {
+    if (driverClazz != null) {
+      return driverClazz.isAssignableFrom(e.getClass())
+             || e instanceof SQLException && ((SQLException) e).getErrorCode() == 1317 /* ER_QUERY_INTERRUPTED */;
+    }
+    return false;
+  }
 }

extensions-core/mysql-metadata-storage/src/test/java/org/apache/druid/metadata/storage/mysql/MySQLConnectorTest.java

@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.metadata.storage.mysql;
+
+import com.mysql.jdbc.exceptions.MySQLTransactionRollbackException;
+import com.mysql.jdbc.exceptions.MySQLTransientException;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.sql.SQLException;
+import java.sql.SQLTransientConnectionException;
+
+public class MySQLConnectorTest
+{
+  @Test
+  public void testIsExceptionTransient()
+  {
+    final Class<?> clazz = MySQLTransientException.class;
+    Assert.assertTrue(MySQLConnector.isTransientException(clazz, new MySQLTransientException()));
+    Assert.assertTrue(MySQLConnector.isTransientException(clazz, new MySQLTransactionRollbackException()));
+    Assert.assertTrue(
+        MySQLConnector.isTransientException(clazz, new SQLException("some transient failure", "wtf", 1317))
+    );
+    Assert.assertFalse(
+        MySQLConnector.isTransientException(clazz, new SQLException("totally realistic test data", "wtf", 1337))
+    );
+    // this method does not specially handle normal transient exceptions either, since it is not vendor specific
+    Assert.assertFalse(
+        MySQLConnector.isTransientException(clazz, new SQLTransientConnectionException("transient"))
+    );
+  }
+
+  @Test
+  public void testIsExceptionTransientNoMySqlClazz()
+  {
+    // no vendor specific for MariaDb, so should always be false
+    Assert.assertFalse(MySQLConnector.isTransientException(null, new MySQLTransientException()));
+    Assert.assertFalse(
+        MySQLConnector.isTransientException(null, new SQLException("some transient failure", "wtf", 1317))
+    );
+    Assert.assertFalse(
+        MySQLConnector.isTransientException(null, new SQLException("totally realistic test data", "wtf", 1337))
+    );
+    Assert.assertFalse(
+        MySQLConnector.isTransientException(null, new SQLTransientConnectionException("transient"))
+    );
+  }
+}

server/src/main/java/org/apache/druid/metadata/SQLMetadataConnector.java

@@ -171,6 +171,9 @@ public abstract class SQLMetadataConnector implements MetadataStorageConnector
                          || (e instanceof DBIException && isTransientException(e.getCause())));
   }
 
+  /**
+   * Vendor specific errors that are not covered by {@link #isTransientException(Throwable)}
+   */
   protected boolean connectorIsTransientException(Throwable e)
   {
     return false;