Include Column List option

README.md

@@ -571,6 +571,16 @@ must also set a distribution key with the <tt>distkey</tt> option.
     <p>Since setting <tt>usestagingtable=false</tt> operation risks data loss / unavailability, we have chosen to deprecate it in favor of requiring users to manually drop the destination table themselves.</p>
     </td>
  </tr>
+  <tr>
+    <td><tt>include_column_list</tt></td>
+    <td>No</td>
+    <td>false</td>
+    <td>
+        If <tt>true</tt> then this library will automatically extract the columns from the schema
+        and add them to the COPY command according to the <a href="http://docs.aws.amazon.com/redshift/latest/dg/copy-parameters-column-mapping.html">Column List docs</a>.
+        (e.g. `COPY "PUBLIC"."tablename" ("column1" [,"column2", ...])`).
+    </td>
+  </tr>
  <tr>
     <td><tt>description</tt></td>
     <td>No</td>

src/main/scala/com/databricks/spark/redshift/Parameters.scala

@@ -38,7 +38,8 @@ private[redshift] object Parameters {
     "diststyle" -> "EVEN",
     "usestagingtable" -> "true",
     "preactions" -> ";",
-    "postactions" -> ";"
+    "postactions" -> ";",
+    "include_column_list" -> "false"
   )
 
   val VALID_TEMP_FORMATS = Set("AVRO", "CSV", "CSV GZIP")
@@ -285,5 +286,11 @@ private[redshift] object Parameters {
           new BasicSessionCredentials(accessKey, secretAccessKey, sessionToken))
       }
     }
+
+    /**
+     * If true then this library will extract the column list from the schema to
+     * include in the COPY command (e.g. `COPY "PUBLIC"."tablename" ("column1" [,"column2", ...])`)
+     */
+    def includeColumnList: Boolean = parameters("include_column_list").toBoolean
   }
 }

src/main/scala/com/databricks/spark/redshift/RedshiftWriter.scala

@@ -88,6 +88,7 @@ private[redshift] class RedshiftWriter(
    */
   private def copySql(
       sqlContext: SQLContext,
+      schema: StructType,
       params: MergedParameters,
       creds: AWSCredentialsProvider,
       manifestUrl: String): String = {
@@ -98,7 +99,13 @@ private[redshift] class RedshiftWriter(
       case "AVRO" => "AVRO 'auto'"
       case csv if csv == "CSV" || csv == "CSV GZIP" => csv + s" NULL AS '${params.nullString}'"
     }
-    s"COPY ${params.table.get} FROM '$fixedUrl' CREDENTIALS '$credsString' FORMAT AS " +
+    val columns = if (params.includeColumnList) {
+        "(" + schema.fieldNames.map(name => s""""$name"""").mkString(",") + ") "
+    } else {
+      ""
+    }
+
+    s"COPY ${params.table.get} ${columns}FROM '$fixedUrl' CREDENTIALS '$credsString' FORMAT AS " +
       s"${format} manifest ${params.extraCopyOptions}"
   }
 
@@ -140,7 +147,7 @@ private[redshift] class RedshiftWriter(
 
     manifestUrl.foreach { manifestUrl =>
       // Load the temporary data into the new file
-      val copyStatement = copySql(data.sqlContext, params, creds, manifestUrl)
+      val copyStatement = copySql(data.sqlContext, data.schema, params, creds, manifestUrl)
       log.info(copyStatement)
       try {
         jdbcWrapper.executeInterruptibly(conn.prepareStatement(copyStatement))

src/test/scala/com/databricks/spark/redshift/ParametersSuite.scala

@@ -28,7 +28,8 @@ class ParametersSuite extends FunSuite with Matchers {
       "tempdir" -> "s3://foo/bar",
       "dbtable" -> "test_schema.test_table",
       "url" -> "jdbc:redshift://foo/bar?user=user&password=password",
-      "forward_spark_s3_credentials" -> "true")
+      "forward_spark_s3_credentials" -> "true",
+      "include_column_list" -> "true")
 
     val mergedParams = Parameters.mergeParameters(params)
 
@@ -37,9 +38,10 @@ class ParametersSuite extends FunSuite with Matchers {
     mergedParams.jdbcUrl shouldBe params("url")
     mergedParams.table shouldBe Some(TableName("test_schema", "test_table"))
     assert(mergedParams.forwardSparkS3Credentials)
+    assert(mergedParams.includeColumnList)
 
     // Check that the defaults have been added
-    (Parameters.DEFAULT_PARAMETERS - "forward_spark_s3_credentials").foreach {
+    (Parameters.DEFAULT_PARAMETERS - "forward_spark_s3_credentials" - "include_column_list").foreach {
       case (key, value) => mergedParams.parameters(key) shouldBe value
     }
   }

src/test/scala/com/databricks/spark/redshift/RedshiftSourceSuite.scala

@@ -435,6 +435,27 @@ class RedshiftSourceSuite
     mockRedshift.verifyThatExpectedQueriesWereIssued(expectedCommands)
   }
 
+  test("Include Column List adds the schema columns to the COPY query") {
+    val copyCommand =
+          "COPY \"PUBLIC\".\"test_table\" \\(\"testbyte\",\"testbool\",\"testdate\",\"testdouble\"" +
+          ",\"testfloat\",\"testint\",\"testlong\",\"testshort\",\"teststring\",\"testtimestamp\"\\) FROM .*"
+    val expectedCommands =
+      Seq("CREATE TABLE IF NOT EXISTS \"PUBLIC\".\"test_table\" .*".r,
+          copyCommand.r)
+
+    val params = defaultParams ++ Map("include_column_list" -> "true")
+
+    val mockRedshift = new MockRedshift(
+      defaultParams("url"),
+      Map(TableName.parseFromEscaped(defaultParams("dbtable")).toString -> null))
+
+    val source = new DefaultSource(mockRedshift.jdbcWrapper, _ => mockS3Client)
+    source.createRelation(testSqlContext, SaveMode.Append, params, expectedDataDF)
+
+    mockRedshift.verifyThatConnectionsWereClosed()
+    mockRedshift.verifyThatExpectedQueriesWereIssued(expectedCommands)
+  }
+
   test("configuring maxlength on string columns") {
     val longStrMetadata = new MetadataBuilder().putLong("maxlength", 512).build()
     val shortStrMetadata = new MetadataBuilder().putLong("maxlength", 10).build()