Implement the setters

Author: wnob

core/src/main/java/org/apache/calcite/avatica/AvaticaSite.java

@@ -37,6 +37,8 @@
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.sql.Types;
+import java.time.OffsetDateTime;
+import java.time.OffsetTime;
 import java.util.Calendar;
 
 /**
@@ -191,10 +193,18 @@ public void setTimestamp(Timestamp x, Calendar calendar) {
     slots[index] = wrap(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, x, calendar);
   }
 
+  public void setTimestampWithTimezone(OffsetDateTime x) {
+    slots[index] = wrap(ColumnMetaData.Rep.JAVA_TIME_OFFSETDATETIME, x);
+  }
+
   public void setTime(Time x, Calendar calendar) {
     slots[index] = wrap(ColumnMetaData.Rep.JAVA_SQL_TIME, x, calendar);
   }
 
+  public void setTimeWithTimezone(OffsetTime x) {
+    slots[index] = wrap(ColumnMetaData.Rep.JAVA_TIME_OFFSETTIME, x);
+  }
+
   public void setDate(Date x, Calendar calendar) {
     slots[index] = wrap(ColumnMetaData.Rep.JAVA_SQL_DATE, x, calendar);
   }
@@ -284,6 +294,12 @@ public void setObject(Object x, int targetSqlType) {
     case Types.TINYINT:
       setByte(toByte(x));
       break;
+    case Types.TIME_WITH_TIMEZONE:
+      setTimeWithTimezone(toTimeWithTimezone(x));
+      break;
+    case Types.TIMESTAMP_WITH_TIMEZONE:
+      setTimestampWithTimezone(toTimestampWithTimezone(x));
+      break;
     default:
       throw notImplemented();
     }
@@ -502,13 +518,37 @@ private static Time toTime(Object x) {
     return new Time(toLong(x));
   }
 
+  private static OffsetTime toTimeWithTimezone(Object x) {
+    if (x instanceof OffsetTime) {
+      return (OffsetTime) x;
+    }
+    if (x instanceof String) {
+      // TODO: This parses the string only to serialize it again in `TypeValue.jdbcToSerial`.
+      //       Think about this more.
+      return OffsetTime.parse((String) x);
+    }
+    throw unsupportedCast(x.getClass(), OffsetTime.class);
+  }
+
   private static Timestamp toTimestamp(Object x) {
     if (x instanceof String) {
       return Timestamp.valueOf((String) x);
     }
     return new Timestamp(toLong(x));
   }
 
+  private static OffsetDateTime toTimestampWithTimezone(Object x) {
+    if (x instanceof OffsetDateTime) {
+      return (OffsetDateTime) x;
+    }
+    if (x instanceof String) {
+      // TODO: This parses the string only to serialize it again in `TypeValue.jdbcToSerial`.
+      //       Think about this more.
+      return OffsetDateTime.parse((String) x);
+    }
+    throw unsupportedCast(x.getClass(), OffsetDateTime.class);
+  }
+
   private static double toDouble(Object x) {
     if (x instanceof Number) {
       return ((Number) x).doubleValue();

core/src/main/java/org/apache/calcite/avatica/remote/TypedValue.java

@@ -39,6 +39,8 @@
 import java.sql.SQLException;
 import java.sql.Time;
 import java.sql.Timestamp;
+import java.time.OffsetDateTime;
+import java.time.OffsetTime;
 import java.util.ArrayList;
 import java.util.Calendar;
 import java.util.Date;
@@ -114,6 +116,18 @@
  *     <td>ARRAY</td>
  *                  <td>Array</td> <td>List&lt;Object&gt;</td> <td>List&lt;Object&gt;</td>
  *   </tr>
+ *   <tr>
+ *     <td>TIME_WITH_TIMEZONE</td>
+ *     <td>{@link java.time.OffsetTime}</td>
+ *     <td>{@link String}</td>
+ *     <td>{@link java.time.OffsetTime}</td>
+ *   </tr>
+ *   <tr>
+ *     <td>TIMESTAMP_WITH_TIMEZONE</td>
+ *     <td>{@link java.time.OffsetDateTime}</td>
+ *     <td>{@link String}</td>
+ *     <td>{@link java.time.OffsetDateTime}</td>
+ *   </tr>
  * </table>
  *
  * <p>Note:
@@ -182,6 +196,8 @@ private boolean isSerial(ColumnMetaData.Rep rep, Object value) {
     }
     switch (rep) {
     case BYTE_STRING:
+    case JAVA_TIME_OFFSETTIME:
+    case JAVA_TIME_OFFSETDATETIME:
       return value instanceof String;
     case JAVA_SQL_DATE:
     case JAVA_SQL_TIME:
@@ -412,6 +428,10 @@ private static Object jdbcToSerial(ColumnMetaData.Rep rep, Object value,
       default:
         return t;
       }
+    case JAVA_TIME_OFFSETTIME:
+      return ((OffsetTime) value).toString();
+    case JAVA_TIME_OFFSETDATETIME:
+      return ((OffsetDateTime) value).toString();
     case ARRAY:
       Array array = (Array) value;
       Objects.requireNonNull(componentType, "Component Type must not be null for ARRAYs");

core/src/test/java/org/apache/calcite/avatica/AvaticaResultSetConversionsTest.java

@@ -1142,7 +1142,8 @@ private TimestampAccessorTestHelper(Getter g) {
 
     @Override public void testGetTimestamp(ResultSet resultSet, Calendar calendar)
         throws SQLException {
-      assertEquals(new Timestamp(TEST_TIMESTAMP_MILLIS_SINCE_EPOCH), g.getTimestamp(resultSet, calendar));
+      assertEquals(
+          new Timestamp(TEST_TIMESTAMP_MILLIS_SINCE_EPOCH), g.getTimestamp(resultSet, calendar));
     }
   }