From e892bda57a39bbb40abdc58d241a0162e067510d Mon Sep 17 00:00:00 2001
From: cip999 <andcip999@gmail.com>
Date: Fri, 20 Aug 2021 17:31:17 +0200
Subject: [PATCH] Added "datetime_format" URL parameter to handle jdbc date and
 time formatting with new engine

---
 .../protocol/response/QueryResultTest.java    | 44 +++++++++++++++++++
 .../sql/sql/domain/SQLQueryRequestTest.java   | 14 ++++++
 2 files changed, 58 insertions(+)

diff --git a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java
index 5373650b25..945603e463 100644
--- a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java
+++ b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java
@@ -16,18 +16,26 @@
 
 package com.amazon.opendistroforelasticsearch.sql.protocol.response;
 
+import static com.amazon.opendistroforelasticsearch.sql.data.model.ExprTimestampValue.INCLUDE_TIME_WHEN_NONZERO;
 import static com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils.tupleValue;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.INTEGER;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRING;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.TIMESTAMP;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
 
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTimestampValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTupleValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils;
 import com.amazon.opendistroforelasticsearch.sql.executor.ExecutionEngine;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.LinkedHashMap;
+
 import org.junit.jupiter.api.Test;
 
 class QueryResultTest {
@@ -36,6 +44,10 @@ class QueryResultTest {
       new ExecutionEngine.Schema.Column("name", null, STRING),
       new ExecutionEngine.Schema.Column("age", null, INTEGER)));
 
+  private ExecutionEngine.Schema schemaTimestamp = new ExecutionEngine.Schema(ImmutableList.of(
+          new ExecutionEngine.Schema.Column("name", null, STRING),
+          new ExecutionEngine.Schema.Column("birthdate", null, TIMESTAMP)));
+
 
   @Test
   void size() {
@@ -125,4 +137,36 @@ void iterate() {
     }
   }
 
+  @Test
+  void setDateTimeFormatTest() {
+    LinkedHashMap<String, ExprValue> valueMap1 = new LinkedHashMap<>();
+    valueMap1.put("name", ExprValueUtils.fromObjectValue("John"));
+    valueMap1.put("birthdate", ExprValueUtils.fromObjectValue("1987-10-23 00:00:00", TIMESTAMP));
+
+    LinkedHashMap<String, ExprValue> valueMap2 = new LinkedHashMap<>();
+    valueMap2.put("name", ExprValueUtils.fromObjectValue("Andrea"));
+    valueMap2.put("birthdate", ExprValueUtils.fromObjectValue("1999-03-14 03:30:00", TIMESTAMP));
+
+    QueryResult response = new QueryResult(
+            schemaTimestamp,
+            Arrays.asList(
+                    new ExprTupleValue(valueMap1),
+                    new ExprTupleValue(valueMap2)
+            ));
+
+    response.setDatetimeFormat(INCLUDE_TIME_WHEN_NONZERO);
+
+    int i = 0;
+    for (Object[] objects : response) {
+      if (i == 0) {
+        assertArrayEquals(new Object[] {"John", "1987-10-23"}, objects);
+      } else if (i == 1) {
+        assertArrayEquals(new Object[] {"Andrea", "1999-03-14 03:30:00"}, objects);
+      } else {
+        fail("More rows returned than expected");
+      }
+      i++;
+    }
+  }
+
 }
\ No newline at end of file
diff --git a/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/domain/SQLQueryRequestTest.java b/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/domain/SQLQueryRequestTest.java
index b7ea848224..5d16757dad 100644
--- a/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/domain/SQLQueryRequestTest.java
+++ b/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/domain/SQLQueryRequestTest.java
@@ -23,6 +23,8 @@
 
 import com.amazon.opendistroforelasticsearch.sql.protocol.response.format.Format;
 import com.google.common.collect.ImmutableMap;
+
+import java.util.HashMap;
 import java.util.Map;
 import org.json.JSONObject;
 import org.junit.jupiter.api.Test;
@@ -44,6 +46,18 @@ public void shouldSupportQueryWithJDBCFormat() {
     assertEquals(request.format(), Format.JDBC);
   }
 
+  @Test
+  public void shouldSupportQueryDatetimeFormat() {
+    Map<String, String> params = new HashMap<>();
+    params.put("format", "jdbc");
+    params.put("datetime_format", "never_include_time");
+
+    SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1")
+            .params(params)
+            .build();
+    assertEquals("never_include_time", request.getDatetimeFormat());
+  }
+
   @Test
   public void shouldSupportQueryWithQueryFieldOnly() {
     SQLQueryRequest request =