SNOW-1002378 Add new Parameter "string_timestamp_format" (#539)

* add new parameter

* fix test

Author: sfc-gh-bli

src/it/scala/net/snowflake/spark/snowflake/DataTypesIntegrationSuite.scala

@@ -16,14 +16,13 @@
 
 package net.snowflake.spark.snowflake
 
-import java.sql.Timestamp
+import java.sql.{SQLException, Timestamp}
 import java.text.{DateFormat, SimpleDateFormat}
-
 import org.apache.spark.sql.{Row, SaveMode}
 import net.snowflake.spark.snowflake.Utils.SNOWFLAKE_SOURCE_NAME
 import net.snowflake.spark.snowflake.Utils.SNOWFLAKE_SOURCE_SHORT_NAME
 import org.apache.spark.sql.functions._
-import org.apache.spark.sql.types.{DateType, TimestampType}
+import org.apache.spark.sql.types.{DateType, StringType, StructField, StructType, TimestampType}
 
 /**
   * Created by mzukowski on 8/12/16.
@@ -98,6 +97,50 @@ class DataTypesIntegrationSuite extends IntegrationSuiteBase {
     )
   }
 
+  test("change timestamp format") {
+    jdbcUpdate(s"CREATE OR REPLACE TABLE $test_table (START_TIME TIMESTAMP)")
+
+    val spark = sparkSession
+    import spark.implicits._
+
+    val df1 = Seq("2023-11-28 10:23:59.123456").toDF()
+    df1.write.format(SNOWFLAKE_SOURCE_NAME)
+      .options(connectorOptionsNoTable).option("dbtable", test_table)
+      .option(Parameters.PARAM_STRING_TIMESTAMP_FORMAT, "YYYY-MM-DD HH24:MI:SS.FF9")
+      .mode(SaveMode.Append).save()
+
+    val df = sparkSession.read.format(SNOWFLAKE_SOURCE_NAME)
+      .options(connectorOptionsNoTable).option("dbtable", test_table).load()
+    assert(df.collect().head.getTimestamp(0).toString.startsWith("2023-11-28 10:23:59.123"))
+
+    // it doesn't work if source dataframe has timestamp column
+    val dateFormat: DateFormat = new SimpleDateFormat("dd/MM/yyyy HH:mm:ss.SSS")
+    val time: Timestamp =
+      new Timestamp(dateFormat.parse("28/10/1996 00:00:00.000").getTime)
+
+    val data = sc.parallelize(Seq(
+      Row(time, "2023-11-28 10:23:59.123456")
+    ))
+
+    val schema = new StructType(
+      Array(
+        StructField("time", TimestampType),
+        StructField("str", StringType)
+      )
+    )
+
+    jdbcUpdate(s"CREATE OR REPLACE TABLE $test_table (START_TIME TIMESTAMP, END_TIME TIMESTAMP)")
+
+    val df2 = sparkSession.createDataFrame(data, schema)
+
+    assertThrows[SQLException]{
+      df2.write.format(SNOWFLAKE_SOURCE_NAME).options(connectorOptionsNoTable)
+        .options(connectorOptionsNoTable).option("dbtable", test_table)
+        .option(Parameters.PARAM_STRING_TIMESTAMP_FORMAT, "YYYY-MM-DD HH24:MI:SS.FF9")
+        .mode(SaveMode.Append).save()
+    }
+  }
+
   test("insert timestamp into date") {
     jdbcUpdate(s"create or replace table $test_table(i date)")
 

src/main/scala/net/snowflake/spark/snowflake/Parameters.scala

@@ -114,6 +114,9 @@ object Parameters {
   val PARAM_TIMESTAMP_NTZ_OUTPUT_FORMAT: String = knownParam("timestamp_ntz_output_format")
   val PARAM_TIMESTAMP_LTZ_OUTPUT_FORMAT: String = knownParam("timestamp_ltz_output_format")
   val PARAM_TIMESTAMP_TZ_OUTPUT_FORMAT: String = knownParam("timestamp_tz_output_format")
+  // changed the timestamp format when copy data from string column to timestamp column.
+  // it only works when source dataframe doesn't have timestamp columns.
+  val PARAM_STRING_TIMESTAMP_FORMAT: String = knownParam("string_timestamp_format")
   val TIMESTAMP_OUTPUT_FORMAT_SF_CURRENT: String = "sf_current"
   val PARAM_JDBC_QUERY_RESULT_FORMAT: String = knownParam(
     "jdbc_query_result_format"
@@ -1017,6 +1020,9 @@ object Parameters {
       }
     }
 
+    def getStringTimestampFormat: Option[String] =
+      parameters.get(PARAM_STRING_TIMESTAMP_FORMAT)
+
     def streamingStage: Option[String] = parameters.get(PARAM_STREAMING_STAGE)
 
     def storagePath: Option[String] = {

src/main/scala/net/snowflake/spark/snowflake/io/StageWriter.scala

@@ -25,7 +25,7 @@ import net.snowflake.spark.snowflake.io.SupportedFormat.SupportedFormat
 import net.snowflake.spark.snowflake.DefaultJDBCWrapper.DataBaseOperations
 import net.snowflake.spark.snowflake.test.{TestHook, TestHookFlag}
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.types.{BinaryType, StructType}
+import org.apache.spark.sql.types.{BinaryType, StructType, TimestampType}
 import org.apache.spark.sql.{SQLContext, SaveMode}
 import org.slf4j.LoggerFactory
 
@@ -191,6 +191,8 @@ private[io] object StageWriter {
 
   private[io] val log = LoggerFactory.getLogger(getClass)
 
+  private val DEFAULT_TIMESTAMP_FORMAT: String = "TZHTZM YYYY-MM-DD HH24:MI:SS.FF9"
+
   def writeToStage(sqlContext: SQLContext,
                    rdd: RDD[String],
                    schema: StructType,
@@ -892,6 +894,21 @@ private[io] object StageWriter {
 
     val mappingFromString = getMappingFromString(mappingList, fromString)
 
+    val hasTimestampColumn: Boolean =
+      schema.exists(field => field.dataType == TimestampType)
+
+    val timestampFormat: String =
+      if (params.getStringTimestampFormat.isEmpty) {
+        DEFAULT_TIMESTAMP_FORMAT
+      } else if (hasTimestampColumn) {
+        log.error("The source Dataframe contains timestamp columns, "
+          + "the parameter `string_timestamp_format` has been ignored.")
+        DEFAULT_TIMESTAMP_FORMAT
+      } else {
+        params.getStringTimestampFormat.get
+      }
+
+
     val formatString =
       format match {
         case SupportedFormat.CSV =>
@@ -901,7 +918,7 @@ private[io] object StageWriter {
                |    FIELD_DELIMITER='|'
                |    NULL_IF=()
                |    FIELD_OPTIONALLY_ENCLOSED_BY='"'
-               |    TIMESTAMP_FORMAT='TZHTZM YYYY-MM-DD HH24:MI:SS.FF9'
+               |    TIMESTAMP_FORMAT='$timestampFormat'
                |    DATE_FORMAT='TZHTZM YYYY-MM-DD HH24:MI:SS.FF9'
                |    BINARY_FORMAT=BASE64
                |  )