Handle DriverWrapper when scanning registered drivers in JDBWrapper.g…

…etConnector() This is a followup to #143 which fixes a corner-case bug in our scanning of registered JDBC drivers: we need to properly handle Spark's `DriverWrapper` drivers, which are used to wrap JDBC drivers in order to make them accessible from the root classloader so that the `DriverManager` can find them. A simpler, reflection-free version of this change was incorporated into apache/spark#10519 Author: Josh Rosen <[email protected]> Closes #147 from JoshRosen/jdbc-driver-precedence-round-2.
databricks · Jan 6, 2016 · 038d061 · 038d061
1 parent 61c04da
commit 038d061
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/src/main/scala/com/databricks/spark/redshift/RedshiftJDBCWrapper.scala b/src/main/scala/com/databricks/spark/redshift/RedshiftJDBCWrapper.scala
@@ -201,12 +201,23 @@ private[redshift] class JDBCWrapper {
     val subprotocol = url.stripPrefix("jdbc:").split(":")(0)
     val driverClass: String = getDriverClass(subprotocol, userProvidedDriverClass)
     registerDriver(driverClass)
+    val driverWrapperClass: Class[_] = if (SPARK_VERSION.startsWith("1.4")) {
+      Utils.classForName("org.apache.spark.sql.jdbc.package$DriverWrapper")
+    } else { // Spark 1.5.0+
+      Utils.classForName("org.apache.spark.sql.execution.datasources.jdbc.DriverWrapper")
+    }
+    def getWrapped(d: Driver): Driver = {
+      require(driverWrapperClass.isAssignableFrom(d.getClass))
+      driverWrapperClass.getDeclaredMethod("wrapped").invoke(d).asInstanceOf[Driver]
+    }
     // Note that we purposely don't call DriverManager.getConnection() here: we want to ensure
     // that an explicitly-specified user-provided driver class can take precedence over the default
     // class, but DriverManager.getConnection() might return a according to a different precedence.
     // At the same time, we don't want to create a driver-per-connection, so we use the
     // DriverManager's driver instances to handle that singleton logic for us.
     val driver: Driver = DriverManager.getDrivers.asScala.collectFirst {
+      case d if driverWrapperClass.isAssignableFrom(d.getClass)
+        && getWrapped(d).getClass.getCanonicalName == driverClass => d
       case d if d.getClass.getCanonicalName == driverClass => d
     }.getOrElse {
       throw new IllegalArgumentException(s"Did not find registered driver with class $driverClass")