Fixed RD-10923: SQL time/timestamps are losing milliseconds

snapi-compiler/src/test/scala/com/rawlabs/snapi/compiler/tests/builtin/credentials/MySQLPackageTest.scala

@@ -36,9 +36,9 @@ class MySQLPackageTest extends SnapiTestContext {
     |       CAST(3.14 AS DOUBLE) AS doublecol,
     |       CAST(1200000000 AS DECIMAL) AS decimalcol,
     |       '120' AS stringcol,
-    |       CAST('12:23:34' AS TIME) AS timecol,
+    |       TIME('12:23:34.123') AS timecol,
     |       CAST('2020-01-01' AS DATE) AS datecol,
-    |       CAST('2020-01-01 12:23:34' AS DATETIME) AS timestampcol,
+    |       TIMESTAMP('2020-01-01 12:23:34.123') AS timestampcol,
     |       1 = 0 AS boolcol,
     |       convert('Olala!' using utf8) AS binarycol$ttt, type collection(
     |    record(
@@ -66,9 +66,9 @@ class MySQLPackageTest extends SnapiTestContext {
       |  doublecol: 3.14,
       |  decimalcol: Decimal.From(1200000000),
       |  stringcol: "120",
-      |  timecol: Time.Build(12, 23, seconds=34),
+      |  timecol: Time.Build(12, 23, seconds=34, millis=123),
       |  datecol: Date.Build(2020, 1, 1),
-      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34),
+      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34, millis=123),
       |  boolcol: false,
       |  binarycol: Binary.FromString("Olala!")
       |}]""".stripMargin)

snapi-compiler/src/test/scala/com/rawlabs/snapi/compiler/tests/builtin/credentials/PostgreSQLPackageTest.scala

@@ -36,9 +36,9 @@ class PostgreSQLPackageTest extends SnapiTestContext {
     |       CAST('3.14' AS DOUBLE PRECISION) AS doublecol,
     |       CAST('12000000' AS DECIMAL) AS decimalcol,
     |       CAST('120' AS VARCHAR) AS stringcol,
-    |       CAST('12:23:34' AS TIME) AS timecol,
+    |       CAST('12:23:34.123' AS TIME) AS timecol,
     |       CAST('2020-01-01' AS DATE) AS datecol,
-    |       CAST('2020-01-01 12:23:34' AS TIMESTAMP) AS timestampcol,
+    |       CAST('2020-01-01 12:23:34.123' AS TIMESTAMP) AS timestampcol,
     |       CAST('false' AS BOOL) AS boolcol,
     |       decode('T2xhbGEh', 'base64') as binarycol$ttt, type collection(
     |    record(
@@ -66,9 +66,9 @@ class PostgreSQLPackageTest extends SnapiTestContext {
       |  doublecol: 3.14,
       |  decimalcol: Decimal.From(12000000),
       |  stringcol: "120",
-      |  timecol: Time.Build(12, 23, seconds=34),
+      |  timecol: Time.Build(12, 23, seconds=34, millis=123),
       |  datecol: Date.Build(2020, 1, 1),
-      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34),
+      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34, millis=123),
       |  boolcol: false,
       |  binarycol: Binary.FromString("Olala!")
       |}]""".stripMargin)

snapi-compiler/src/test/scala/com/rawlabs/snapi/compiler/tests/builtin/credentials/SQLServerPackageTest.scala

@@ -36,9 +36,9 @@ class SQLServerPackageTest extends SnapiTestContext {
     |       CAST('3.14' AS DOUBLE PRECISION) AS doublecol,
     |       CAST('12000000' AS DECIMAL) AS decimalcol,
     |       CAST('120' AS VARCHAR) AS stringcol,
-    |       CAST('12:23:34' AS TIME) AS timecol,
+    |       CAST('12:23:34.123' AS TIME) AS timecol,
     |       CAST('2020-01-01' AS DATE) AS datecol,
-    |       CAST('2020-01-01 12:23:34' AS DATETIME) AS timestampcol,
+    |       CAST('2020-01-01 12:23:34.123' AS DATETIME) AS timestampcol,
     |       CAST('Olala!' AS VARBINARY(MAX)) AS binarycol $ttt)""".stripMargin) { it =>
     it should typeAs("""collection(
       |    record(
@@ -65,9 +65,9 @@ class SQLServerPackageTest extends SnapiTestContext {
       |  doublecol: 3.14,
       |  decimalcol: Decimal.From(12000000),
       |  stringcol: "120",
-      |  timecol: Time.Build(12, 23, seconds=34),
+      |  timecol: Time.Build(12, 23, seconds=34, millis=123),
       |  datecol: Date.Build(2020, 1, 1),
-      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34),
+      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34, millis=123),
       |  binarycol: Binary.FromString("Olala!")
       |}]""".stripMargin)
   }

snapi-compiler/src/test/scala/com/rawlabs/snapi/compiler/tests/builtin/credentials/SnowflakePackageTest.scala

@@ -37,9 +37,9 @@ class SnowflakePackageTest extends SnapiTestContext {
     |       CAST('3.14' AS FLOAT8) AS "doublecol",
     |       CAST('12000000' AS DECIMAL) AS "decimalcol",
     |       CAST('120' AS VARCHAR) AS "stringcol",
-    |       CAST('12:23:34' AS TIME) AS "timecol",
+    |       CAST('12:23:34.123' AS TIME) AS "timecol",
     |       CAST('2020-01-01' AS DATE) AS "datecol",
-    |       CAST('2020-01-01 12:23:34' AS DATETIME) AS "timestampcol",
+    |       CAST('2020-01-01 12:23:34.123' AS DATETIME) AS "timestampcol",
     |       1 = 0 AS "boolcol",
     |       to_binary('tralala', 'utf-8') AS "binarycol" $ttt, type collection(
     |    record(
@@ -67,9 +67,9 @@ class SnowflakePackageTest extends SnapiTestContext {
       |  doublecol: 3.14,
       |  decimalcol: Decimal.From(12000000),
       |  stringcol: "120",
-      |  timecol: Time.Build(12, 23, seconds=34),
+      |  timecol: Time.Build(12, 23, seconds=34, millis=123),
       |  datecol: Date.Build(2020, 1, 1),
-      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34),
+      |  timestampcol: Timestamp.Build(2020, 1, 1, 12, 23, seconds=34, millis=123),
       |  boolcol: false,
       |  binarycol: String.Encode("tralala", "utf-8")
       |}]""".stripMargin)

snapi-truffle/src/main/java/com/rawlabs/snapi/truffle/ast/io/jdbc/JdbcQuery.java

@@ -23,6 +23,7 @@
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.time.temporal.ChronoField;
 
 public class JdbcQuery {
 
@@ -172,8 +173,20 @@ DateObject getDate(String colName, Node node) {
   @TruffleBoundary
   TimeObject getTime(String colName, Node node) {
     try {
+      // Extract the SQL time (a JDBC object) from the result set.
       java.sql.Time sqlTime = rs.getTime(colName);
-      return new TimeObject(sqlTime.toLocalTime());
+      // Turn it into LocalTime. It does something proper with potential timezone conversion, but
+      // doesn't have the milliseconds (toLocalTime's doc says it sets the LocalTime nanoseconds
+      // field to zero).
+      java.time.LocalTime withoutMilliseconds = sqlTime.toLocalTime();
+      // Get the value as milliseconds (possibly shifted by a certain timezone) but we have the
+      // milliseconds.
+      long asMillis = sqlTime.getTime();
+      // Extract the actual milliseconds.
+      long millis = asMillis % 1000;
+      // Fix the LocalTime milliseconds.
+      java.time.LocalTime localTime = withoutMilliseconds.with(ChronoField.MILLI_OF_SECOND, millis);
+      return new TimeObject(localTime);
     } catch (SQLException e) {
       throw exceptionHandler.columnParseError(e, colName, node);
     }

sql-compiler/src/main/scala/com/rawlabs/sql/compiler/NamedParametersPreparedStatement.scala

@@ -46,6 +46,7 @@ import org.bitbucket.inkytonik.kiama.util.Position
 import org.postgresql.util.PSQLException
 
 import java.sql.{Connection, ResultSet, ResultSetMetaData}
+import java.time.LocalTime
 import scala.collection.mutable
 
 class NamedParametersPreparedStatementException(val errors: List[ErrorMessage]) extends Exception
@@ -493,7 +494,7 @@ class NamedParametersPreparedStatement(
         case RawString(v) => setString(paramName, v)
         case RawDecimal(v) => setBigDecimal(paramName, v)
         case RawDate(v) => setDate(paramName, java.sql.Date.valueOf(v))
-        case RawTime(v) => setTime(paramName, java.sql.Time.valueOf(v))
+        case RawTime(v) => setTime(paramName, new java.sql.Time(v.toNanoOfDay / 1000000))
         case RawTimestamp(v) => setTimestamp(paramName, java.sql.Timestamp.valueOf(v))
         case RawInterval(years, months, weeks, days, hours, minutes, seconds, millis) => ???
         case RawBinary(v) => setBytes(paramName, v)
@@ -624,7 +625,7 @@ class NamedParametersPreparedStatement(
       case java.sql.Types.FLOAT => RawFloat(rs.getFloat(1))
       case java.sql.Types.DOUBLE => RawDouble(rs.getDouble(1))
       case java.sql.Types.DATE => RawDate(rs.getDate(1).toLocalDate)
-      case java.sql.Types.TIME => RawTime(rs.getTime(1).toLocalTime)
+      case java.sql.Types.TIME => RawTime(LocalTime.ofNanoOfDay(rs.getTime(1).getTime * 1000000))
       case java.sql.Types.TIMESTAMP => RawTimestamp(rs.getTimestamp(1).toLocalDateTime)
       case java.sql.Types.BOOLEAN => RawBool(rs.getBoolean(1))
       case java.sql.Types.VARCHAR => RawString(rs.getString(1))

sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetCsvWriter.scala

@@ -39,6 +39,7 @@ import com.rawlabs.compiler.utils.RecordFieldsNaming
 import java.io.{IOException, OutputStream}
 import java.sql.ResultSet
 import java.time.format.DateTimeFormatter
+import java.time.temporal.ChronoField
 import scala.annotation.tailrec
 
 object TypedResultSetCsvWriter {
@@ -143,7 +144,17 @@ class TypedResultSetCsvWriter(os: OutputStream, lineSeparator: String, maxRows:
         val date = v.getDate(i).toLocalDate
         gen.writeString(dateFormatter.format(date))
       case _: RawTimeType =>
-        val time = v.getTime(i).toLocalTime
+        // Extract the SQL time (a JDBC object) from the result set.
+        val sqlTime = v.getTime(i)
+        // Turn it into LocalTime. It does something proper with potential timezone conversion, but
+        // doesn't have the milliseconds (toLocalTime's doc says it sets the LocalTime nanoseconds field to zero).
+        val withoutMilliseconds = sqlTime.toLocalTime
+        // Get the value as milliseconds (possibly shifted by a certain timezone) but we have the milliseconds.
+        val asMillis = sqlTime.getTime
+        // Extract the actual milliseconds.
+        val millis = asMillis % 1000
+        // Fix the LocalTime milliseconds.
+        val time = withoutMilliseconds.`with`(ChronoField.MILLI_OF_SECOND, millis)
         val formatter = if (time.getNano > 0) timeFormatter else timeFormatterNoMs
         val formatted = formatter.format(time)
         gen.writeString(formatted)

sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala

@@ -40,6 +40,7 @@ import com.rawlabs.compiler.utils.RecordFieldsNaming
 import java.io.{IOException, OutputStream}
 import java.sql.ResultSet
 import java.time.format.DateTimeFormatter
+import java.time.temporal.ChronoField
 import scala.annotation.tailrec
 
 object TypedResultSetJsonWriter {
@@ -134,7 +135,17 @@ class TypedResultSetJsonWriter(os: OutputStream, maxRows: Option[Long]) {
         val date = v.getDate(i).toLocalDate
         gen.writeString(dateFormatter.format(date))
       case _: RawTimeType =>
-        val time = v.getTime(i).toLocalTime
+        // Extract the SQL time (a JDBC object) from the result set.
+        val sqlTime = v.getTime(i)
+        // Turn it into LocalTime. It does something proper with potential timezone conversion, but
+        // doesn't have the milliseconds (toLocalTime's doc says it sets the LocalTime nanoseconds field to zero).
+        val withoutMilliseconds = sqlTime.toLocalTime
+        // Get the value as milliseconds (possibly shifted by a certain timezone) but we have the milliseconds.
+        val asMillis = sqlTime.getTime
+        // Extract the actual milliseconds.
+        val millis = asMillis % 1000
+        // Fix the LocalTime milliseconds.
+        val time = withoutMilliseconds.`with`(ChronoField.MILLI_OF_SECOND, millis)
         val formatted = timeFormatter.format(time)
         gen.writeString(formatted)
       case _: RawTimestampType =>

sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala

@@ -111,7 +111,7 @@ class TestSqlCompilerServiceAirports
       jdbcUrl = Some(jdbcUrl)
     )
   }
-  private def asCsv(params: Map[String, RawValue], scopes: Set[String] = Set.empty): ProgramEnvironment = {
+  private def asCsv(params: Map[String, RawValue] = Map.empty, scopes: Set[String] = Set.empty): ProgramEnvironment = {
     ProgramEnvironment(
       user,
       if (params.isEmpty) None else Some(params.toArray),
@@ -991,4 +991,94 @@ class TestSqlCompilerServiceAirports
     assert(v.messages.size == 1)
     assert(v.messages(0).message == "schema \"country\" does not exist")
   }
+
+  test("""SELECT pg_typeof(NOW())""".stripMargin) { t =>
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionFailure(errors2) = compilerService.getProgramDescription(t.q, asJson())
+    errors2.map(_.message).contains("unsupported type: regtype")
+  }
+
+  test("""SELECT CAST(pg_typeof(NOW()) AS VARCHAR)""".stripMargin) { t =>
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionSuccess(_) = compilerService.getProgramDescription(t.q, asJson())
+    val baos = new ByteArrayOutputStream()
+    baos.reset()
+    assert(compilerService.execute(t.q, asJson(), None, baos) == ExecutionSuccess(true))
+    assert(baos.toString() == """[{"pg_typeof":"timestamp with time zone"}]""")
+
+  }
+
+  test("""SELECT NOW()""".stripMargin) { t =>
+    // NOW() is a timestamp with timezone. The one of the SQL connection. This test is to make sure
+    // it works (we cannot assert on the result).
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionSuccess(description) = compilerService.getProgramDescription(t.q, asJson())
+    val baos = new ByteArrayOutputStream()
+    baos.reset()
+    assert(compilerService.execute(t.q, asJson(), None, baos) == ExecutionSuccess(true))
+  }
+
+  test("""SELECT TIMESTAMP '2001-07-01 12:13:14.567' AS t""".stripMargin) { t =>
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionSuccess(_) = compilerService.getProgramDescription(t.q, asJson())
+    val baos = new ByteArrayOutputStream()
+    for (env <- Seq(asJson(), asCsv())) {
+      baos.reset()
+      assert(compilerService.execute(t.q, env, None, baos) == ExecutionSuccess(true))
+      assert(baos.toString().contains("12:13:14.567"))
+    }
+  }
+
+  test("""SELECT TIMESTAMP '2001-07-01 12:13:14' AS t""".stripMargin) { t =>
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionSuccess(_) = compilerService.getProgramDescription(t.q, asJson())
+    val baos = new ByteArrayOutputStream()
+    baos.reset()
+    assert(compilerService.execute(t.q, asCsv(), None, baos) == ExecutionSuccess(true))
+    assert(baos.toString().contains("12:13:14"))
+    assert(!baos.toString().contains("12:13:14.000"))
+  }
+
+  test("""SELECT TIME '12:13:14.567' AS t""".stripMargin) { t =>
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionSuccess(_) = compilerService.getProgramDescription(t.q, asJson())
+    val baos = new ByteArrayOutputStream()
+    for (env <- Seq(asJson(), asCsv())) {
+      baos.reset()
+      assert(compilerService.execute(t.q, env, None, baos) == ExecutionSuccess(true))
+      assert(baos.toString().contains("12:13:14.567"))
+    }
+  }
+
+  test("""SELECT TIME '12:13:14' AS t""".stripMargin) { t =>
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionSuccess(_) = compilerService.getProgramDescription(t.q, asJson())
+    val baos = new ByteArrayOutputStream()
+    baos.reset()
+    assert(compilerService.execute(t.q, asCsv(), None, baos) == ExecutionSuccess(true))
+    assert(baos.toString().contains("12:13:14"))
+    assert(!baos.toString().contains("12:13:14.000"))
+  }
+
+  test("""-- @default t TIME '12:13:14.567'
+    |SELECT :t AS t""".stripMargin) { t =>
+    val ValidateResponse(errors) = compilerService.validate(t.q, asJson())
+    assert(errors.isEmpty)
+    val GetProgramDescriptionSuccess(_) = compilerService.getProgramDescription(t.q, asJson())
+    val baos = new ByteArrayOutputStream()
+    baos.reset()
+    for (env <- Seq(asJson(), asCsv())) {
+      baos.reset()
+      assert(compilerService.execute(t.q, env, None, baos) == ExecutionSuccess(true))
+      assert(baos.toString().contains("12:13:14.567"))
+    }
+  }
+
 }