[SPARK-14525][SQL] Make DataFrameWrite.save work for jdbc

sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala

@@ -244,7 +244,11 @@ final class DataFrameWriter private[sql](df: DataFrame) {
       bucketSpec = getBucketSpec,
       options = extraOptions.toMap)
 
-    dataSource.write(mode, df)
+    dataSource.providingClass.newInstance() match {
+      case jdbc: execution.datasources.jdbc.DefaultSource =>
+        jdbc.write(mode, df, extraOptions.toMap)
+      case _ => dataSource.write(mode, df)
+    }
   }
 
   /**
@@ -489,46 +493,12 @@ final class DataFrameWriter private[sql](df: DataFrame) {
    * @since 1.4.0
    */
   def jdbc(url: String, table: String, connectionProperties: Properties): Unit = {
-    val props = new Properties()
-    extraOptions.foreach { case (key, value) =>
-      props.put(key, value)
-    }
+    import scala.collection.JavaConverters._
     // connectionProperties should override settings in extraOptions
-    props.putAll(connectionProperties)
-    val conn = JdbcUtils.createConnectionFactory(url, props)()
-
-    try {
-      var tableExists = JdbcUtils.tableExists(conn, url, table)
-
-      if (mode == SaveMode.Ignore && tableExists) {
-        return
-      }
-
-      if (mode == SaveMode.ErrorIfExists && tableExists) {
-        sys.error(s"Table $table already exists.")
-      }
-
-      if (mode == SaveMode.Overwrite && tableExists) {
-        JdbcUtils.dropTable(conn, table)
-        tableExists = false
-      }
-
-      // Create the table if the table didn't exist.
-      if (!tableExists) {
-        val schema = JdbcUtils.schemaString(df, url)
-        val sql = s"CREATE TABLE $table ($schema)"
-        val statement = conn.createStatement
-        try {
-          statement.executeUpdate(sql)
-        } finally {
-          statement.close()
-        }
-      }
-    } finally {
-      conn.close()
-    }
-
-    JdbcUtils.saveTable(df, url, table, props)
+    this.extraOptions = this.extraOptions ++ (connectionProperties.asScala)
+    // explicit url and dbtable should override all
+    this.extraOptions += ("url" -> url, "dbtable" -> table)
+    format("jdbc").save
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/DefaultSource.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.datasources.jdbc
 
 import java.util.Properties
 
-import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.{DataFrame, SaveMode, SQLContext}
 import org.apache.spark.sql.sources.{BaseRelation, DataSourceRegister, RelationProvider}
 
 class DefaultSource extends RelationProvider with DataSourceRegister {
@@ -56,4 +56,52 @@ class DefaultSource extends RelationProvider with DataSourceRegister {
     parameters.foreach(kv => properties.setProperty(kv._1, kv._2))
     JDBCRelation(url, table, parts, properties)(sqlContext)
   }
+
+  def write(mode: SaveMode, data: DataFrame, options: Map[String, String]): Unit = {
+    val url = options.getOrElse("url",
+      sys.error("Saving jdbc source requires url to be set." +
+        " (ie. df.option(\"url\", \"ACTUAL_URL\")"))
+    val table = options.getOrElse("dbtable", options.getOrElse("table",
+      sys.error("Saving jdbc source requires dbtable to be set." +
+        " (ie. df.option(\"dbtable\", \"ACTUAL_DB_TABLE\")")))
+
+    import collection.JavaConverters._
+    val props = new Properties()
+    props.putAll(options.asJava)
+
+    val conn = JdbcUtils.createConnectionFactory(url, props)()
+
+    try {
+      var tableExists = JdbcUtils.tableExists(conn, url, table)
+
+      if (mode == SaveMode.Ignore && tableExists) {
+        return
+      }
+
+      if (mode == SaveMode.ErrorIfExists && tableExists) {
+        sys.error(s"Table $table already exists.")
+      }
+
+      if (mode == SaveMode.Overwrite && tableExists) {
+        JdbcUtils.dropTable(conn, table)
+        tableExists = false
+      }
+
+      // Create the table if the table didn't exist.
+      if (!tableExists) {
+        val schema = JdbcUtils.schemaString(data, url)
+        val sql = s"CREATE TABLE $table ($schema)"
+        val statement = conn.createStatement
+        try {
+          statement.executeUpdate(sql)
+        } finally {
+          statement.close()
+        }
+      }
+    } finally {
+      conn.close()
+    }
+
+    JdbcUtils.saveTable(data, url, table, props)
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCWriteSuite.scala

@@ -151,4 +151,16 @@ class JDBCWriteSuite extends SharedSQLContext with BeforeAndAfter {
     assert(2 === sqlContext.read.jdbc(url1, "TEST.PEOPLE1", properties).count)
     assert(2 === sqlContext.read.jdbc(url1, "TEST.PEOPLE1", properties).collect()(0).length)
   }
+
+  test("save works for format(\"jdbc\") if url and dbtable are set") {
+    val df = sqlContext.createDataFrame(sparkContext.parallelize(arr2x2), schema2)
+
+    df.write.format("jdbc")
+    .options(Map("url" -> url, "dbtable" -> "TEST.BASICCREATETEST"))
+    .save
+
+    assert(2 === sqlContext.read.jdbc(url, "TEST.BASICCREATETEST", new Properties).count)
+    assert(
+      2 === sqlContext.read.jdbc(url, "TEST.BASICCREATETEST", new Properties).collect()(0).length)
+  }
 }