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

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

@@ -587,47 +587,12 @@ final class DataFrameWriter private[sql](df: DataFrame) {
    */
   def jdbc(url: String, table: String, connectionProperties: Properties): Unit = {
     assertNotStreaming("jdbc() can only be called on non-continuous queries")
-
-    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/JDBCRelation.scala

@@ -87,7 +87,8 @@ private[sql] case class JDBCRelation(
     url: String,
     table: String,
     parts: Array[Partition],
-    properties: Properties = new Properties())(@transient val sparkSession: SparkSession)
+    properties: Properties = new Properties(),
+    providedSchemaOption: Option[StructType] = None)(@transient val sparkSession: SparkSession)
   extends BaseRelation
   with PrunedFilteredScan
   with InsertableRelation {
@@ -96,7 +97,16 @@ private[sql] case class JDBCRelation(
 
   override val needConversion: Boolean = false
 
-  override val schema: StructType = JDBCRDD.resolveTable(url, table, properties)
+  override val schema: StructType = {
+    val resolvedSchema = JDBCRDD.resolveTable(url, table, properties)
+    providedSchemaOption match {
+      case Some(providedSchema) =>
+        if (providedSchema.sql.toLowerCase == resolvedSchema.sql.toLowerCase) resolvedSchema
+        else sys.error(s"User specified schema, $providedSchema, " +
+          s"does not match the actual schema, $resolvedSchema.")
+      case None => resolvedSchema
+    }
+  }
 
   // Check if JDBCRDD.compileFilter can accept input filters
   override def unhandledFilters(filters: Array[Filter]): Array[Filter] = {

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

@@ -19,37 +19,105 @@ package org.apache.spark.sql.execution.datasources.jdbc
 
 import java.util.Properties
 
-import org.apache.spark.sql.SQLContext
-import org.apache.spark.sql.sources.{BaseRelation, DataSourceRegister, RelationProvider}
+import org.apache.spark.sql.{DataFrame, SaveMode, SQLContext}
+import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider, DataSourceRegister, RelationProvider, SchemaRelationProvider}
+import org.apache.spark.sql.types.StructType
 
-class JdbcRelationProvider extends RelationProvider with DataSourceRegister {
+class JdbcRelationProvider extends CreatableRelationProvider
+  with SchemaRelationProvider with RelationProvider with DataSourceRegister {
 
   override def shortName(): String = "jdbc"
 
-  /** Returns a new base relation with the given parameters. */
   override def createRelation(
       sqlContext: SQLContext,
       parameters: Map[String, String]): BaseRelation = {
-    val jdbcOptions = new JDBCOptions(parameters)
-    if (jdbcOptions.partitionColumn != null
-      && (jdbcOptions.lowerBound == null
-        || jdbcOptions.upperBound == null
-        || jdbcOptions.numPartitions == null)) {
+    createRelation(sqlContext, parameters, null)
+  }
+
+  /** Returns a new base relation with the given parameters. */
+  override def createRelation(
+      sqlContext: SQLContext,
+      parameters: Map[String, String],
+      schema: StructType): BaseRelation = {
+    val url = parameters.getOrElse("url", sys.error("Option 'url' not specified"))
+    val table = parameters.getOrElse("dbtable", sys.error("Option 'dbtable' not specified"))
+    val partitionColumn = parameters.getOrElse("partitionColumn", null)
+    val lowerBound = parameters.getOrElse("lowerBound", null)
+    val upperBound = parameters.getOrElse("upperBound", null)
+    val numPartitions = parameters.getOrElse("numPartitions", null)
+
+    if (partitionColumn != null
+      && (lowerBound == null || upperBound == null || numPartitions == null)) {
       sys.error("Partitioning incompletely specified")
     }
 
-    val partitionInfo = if (jdbcOptions.partitionColumn == null) {
-      null
-    } else {
+    val partitionInfo = if (partitionColumn == null) null
+    else {
       JDBCPartitioningInfo(
-        jdbcOptions.partitionColumn,
-        jdbcOptions.lowerBound.toLong,
-        jdbcOptions.upperBound.toLong,
-        jdbcOptions.numPartitions.toInt)
+        partitionColumn, lowerBound.toLong, upperBound.toLong, numPartitions.toInt)
     }
     val parts = JDBCRelation.columnPartition(partitionInfo)
     val properties = new Properties() // Additional properties that we will pass to getConnection
     parameters.foreach(kv => properties.setProperty(kv._1, kv._2))
-    JDBCRelation(jdbcOptions.url, jdbcOptions.table, parts, properties)(sqlContext.sparkSession)
+    JDBCRelation(url, table, parts, properties, Option(schema))(sqlContext.sparkSession)
+  }
+
+  /*
+   * The following structure applies to this code:
+   *                 |    tableExists            |          !tableExists
+   *------------------------------------------------------------------------------------
+   * Ignore          | BaseRelation              | CreateTable, saveTable, BaseRelation
+   * ErrorIfExists   | ERROR                     | CreateTable, saveTable, BaseRelation
+   * Overwrite       | DropTable, CreateTable,   | CreateTable, saveTable, BaseRelation
+   *                 | saveTable, BaseRelation   |
+   * Append          | saveTable, BaseRelation   | CreateTable, saveTable, BaseRelation
+   */
+  override def createRelation(
+      sqlContext: SQLContext,
+      mode: SaveMode,
+      parameters: Map[String, String],
+      data: DataFrame): BaseRelation = {
+    val url = parameters.getOrElse("url",
+      sys.error("Saving jdbc source requires url to be set." +
+        " (ie. df.option(\"url\", \"ACTUAL_URL\")"))
+    val table = parameters.getOrElse("dbtable", parameters.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(parameters.asJava)
+    val conn = JdbcUtils.createConnectionFactory(url, props)()
+
+    try {
+      val tableExists = JdbcUtils.tableExists(conn, url, table)
+
+      val (doCreate, doSave) = (mode, tableExists) match {
+        case (SaveMode.Ignore, true) => (false, false)
+        case (SaveMode.ErrorIfExists, true) => sys.error(s"Table $table already exists.")
+        case (SaveMode.Overwrite, true) =>
+          JdbcUtils.dropTable(conn, table)
+          (true, true)
+        case (SaveMode.Append, true) => (false, true)
+        case (_, true) => sys.error(s"Unexpected SaveMode, '$mode', for handling existing tables.")
+        case (_, false) => (true, true)
+      }
+
+      if(doCreate) {
+        val schema = JdbcUtils.schemaString(data, url)
+        val sql = s"CREATE TABLE $table ($schema)"
+        val statement = conn.createStatement
+        try {
+          statement.executeUpdate(sql)
+        } finally {
+          statement.close()
+        }
+      }
+      if(doSave) JdbcUtils.saveTable(data, url, table, props)
+    } finally {
+      conn.close()
+    }
+
+    createRelation(sqlContext, parameters, data.schema)
   }
 }

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

@@ -151,4 +151,16 @@ class JDBCWriteSuite extends SharedSQLContext with BeforeAndAfter {
     assert(2 === spark.read.jdbc(url1, "TEST.PEOPLE1", properties).count)
     assert(2 === spark.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)
+  }
 }