Added "datetime_format" URL parameter to handle jdbj date and time fo…

…rmatting with new engine
opendistro-for-elasticsearch · Aug 22, 2021 · 89a6567 · 89a6567
1 parent 3644545
commit 89a6567
Show file tree

Hide file tree

Showing 10 changed files with 68 additions and 11 deletions.
diff --git a/...rc/main/java/com/amazon/opendistroforelasticsearch/sql/data/model/ExprTimestampValue.java b/...rc/main/java/com/amazon/opendistroforelasticsearch/sql/data/model/ExprTimestampValue.java
@@ -46,14 +46,22 @@ public class ExprTimestampValue extends AbstractExprValue {
   /**
    * todo. only support timestamp in format yyyy-MM-dd HH:mm:ss.
    */
-  private static final DateTimeFormatter FORMATTER_WITNOUT_NANO = DateTimeFormatter
+  private static final DateTimeFormatter FORMATTER_WITHOUT_NANO = DateTimeFormatter
       .ofPattern("yyyy-MM-dd HH:mm:ss");
+  private static final DateTimeFormatter FORMATTER_WITHOUT_TIME = DateTimeFormatter
+          .ofPattern("yyyy-MM-dd");
   private final Instant timestamp;
 
+  private String datetimeFormat;
+
   private static final DateTimeFormatter FORMATTER_VARIABLE_MICROS;
   private static final int MIN_FRACTION_SECONDS = 0;
   private static final int MAX_FRACTION_SECONDS = 6;
 
+  private static final String ALWAYS_INCLUDE_TIME = "always_include_time";
+  private static final String NEVER_INCLUDE_TIME = "never_include_time";
+  private static final String INCLUDE_TIME_WHEN_NONZERO = "include_time_when_nonzero";
+
   static {
     FORMATTER_VARIABLE_MICROS = new DateTimeFormatterBuilder()
         .appendPattern("yyyy-MM-dd HH:mm:ss")
@@ -69,6 +77,8 @@ public class ExprTimestampValue extends AbstractExprValue {
    * Constructor.
    */
   public ExprTimestampValue(String timestamp) {
+    this.datetimeFormat = ALWAYS_INCLUDE_TIME;
+
     try {
       this.timestamp = LocalDateTime.parse(timestamp,
           ExprDateFormatters.TOLERANT_PARSER_DATE_TIME_FORMATTER)
@@ -83,9 +93,29 @@ public ExprTimestampValue(String timestamp) {
 
   @Override
   public String value() {
-    return timestamp.getNano() == 0 ? FORMATTER_WITNOUT_NANO.withZone(ZONE)
-        .format(timestamp.truncatedTo(ChronoUnit.SECONDS))
-        : FORMATTER_VARIABLE_MICROS.withZone(ZONE).format(timestamp);
+    switch (datetimeFormat) {
+      case NEVER_INCLUDE_TIME:
+        return valueWithoutTime();
+
+      case INCLUDE_TIME_WHEN_NONZERO:
+        LocalTime time = timeValue();
+        return (time.getHour() == 0 && time.getMinute() == 0 && time.getSecond() == 0) ?
+                valueWithoutTime() : valueWithTime();
+
+      case ALWAYS_INCLUDE_TIME:
+      default:
+        return valueWithTime();
+    }
+  }
+
+  private String valueWithTime() {
+    return timestamp.getNano() == 0 ? FORMATTER_WITHOUT_NANO.withZone(ZONE)
+            .format(timestamp.truncatedTo(ChronoUnit.SECONDS))
+            : FORMATTER_VARIABLE_MICROS.withZone(ZONE).format(timestamp);
+  }
+
+  private String valueWithoutTime() {
+    return FORMATTER_WITHOUT_TIME.withZone(ZONE).format(timestamp);
   }
 
   @Override
@@ -132,4 +162,6 @@ public boolean equal(ExprValue other) {
   public int hashCode() {
     return Objects.hashCode(timestamp);
   }
+
+  public void setDatetimeFormat(String format) { this.datetimeFormat = format; }
 }
diff --git a/...test/src/test/java/com/amazon/opendistroforelasticsearch/sql/legacy/SQLIntegTestCase.java b/...test/src/test/java/com/amazon/opendistroforelasticsearch/sql/legacy/SQLIntegTestCase.java
@@ -223,8 +223,12 @@ protected Request getSqlCursorCloseRequest(String cursorRequest) {
   }
 
   protected String executeQuery(String query, String requestType) {
+    return executeQuery(query, requestType, "always_include_time");
+  }
+
+  protected String executeQuery(String query, String requestType, String datetimeFormat) {
     try {
-      String endpoint = "/_opendistro/_sql?format=" + requestType;
+      String endpoint = "/_opendistro/_sql?format=" + requestType + "&datetime_format=" + datetimeFormat;
       String requestBody = makeRequest(query);
 
       Request sqlRequest = new Request("POST", endpoint);

diff --git a/...main/java/com/amazon/opendistroforelasticsearch/sql/legacy/plugin/RestSQLQueryAction.java b/...main/java/com/amazon/opendistroforelasticsearch/sql/legacy/plugin/RestSQLQueryAction.java
@@ -172,7 +172,7 @@ private ResponseListener<QueryResponse> createQueryResponseListener(RestChannel
     } else if (format.equals(Format.RAW)) {
       formatter = new RawResponseFormatter();
     } else {
-      formatter = new JdbcResponseFormatter(PRETTY);
+      formatter = new JdbcResponseFormatter(PRETTY, request.getDatetimeFormat());
     }
     return new ResponseListener<QueryResponse>() {
       @Override

diff --git a/.../src/main/java/com/amazon/opendistroforelasticsearch/sql/legacy/plugin/RestSqlAction.java b/.../src/main/java/com/amazon/opendistroforelasticsearch/sql/legacy/plugin/RestSqlAction.java
@@ -172,7 +172,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
     @Override
     protected Set<String> responseParams() {
         Set<String> responseParams = new HashSet<>(super.responseParams());
-        responseParams.addAll(Arrays.asList("sql", "flat", "separator", "_score", "_type", "_id", "newLine", "format", "sanitize"));
+        responseParams.addAll(Arrays.asList("sql", "flat", "separator", "_score", "_type", "_id", "newLine", "format", "datetime_format", "sanitize"));
         return responseParams;
     }
 

diff --git a/...main/java/com/amazon/opendistroforelasticsearch/sql/legacy/plugin/RestSqlStatsAction.java b/...main/java/com/amazon/opendistroforelasticsearch/sql/legacy/plugin/RestSqlStatsAction.java
@@ -84,7 +84,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
     @Override
     protected Set<String> responseParams() {
         Set<String> responseParams = new HashSet<>(super.responseParams());
-        responseParams.addAll(Arrays.asList("sql", "flat", "separator", "_score", "_type", "_id", "newLine", "format", "sanitize"));
+        responseParams.addAll(Arrays.asList("sql", "flat", "separator", "_score", "_type", "_id", "newLine", "format", "datetime_format", "sanitize"));
         return responseParams;
     }
 

diff --git a/...c/main/java/com/amazon/opendistroforelasticsearch/sql/plugin/rest/RestPPLQueryAction.java b/...c/main/java/com/amazon/opendistroforelasticsearch/sql/plugin/rest/RestPPLQueryAction.java
@@ -115,7 +115,7 @@ public String getName() {
   @Override
   protected Set<String> responseParams() {
     Set<String> responseParams = new HashSet<>(super.responseParams());
-    responseParams.addAll(Arrays.asList("format", "sanitize"));
+    responseParams.addAll(Arrays.asList("format", "datetime_format", "sanitize"));
     return responseParams;
   }
 

diff --git a/...c/main/java/com/amazon/opendistroforelasticsearch/sql/plugin/rest/RestPPLStatsAction.java b/...c/main/java/com/amazon/opendistroforelasticsearch/sql/plugin/rest/RestPPLStatsAction.java
@@ -85,7 +85,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
   @Override
   protected Set<String> responseParams() {
     Set<String> responseParams = new HashSet<>(super.responseParams());
-    responseParams.addAll(Arrays.asList("format", "sanitize"));
+    responseParams.addAll(Arrays.asList("format", "datetime_format", "sanitize"));
     return responseParams;
   }
 }
diff --git a/...rc/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResult.java b/...rc/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResult.java
@@ -16,8 +16,10 @@
 
 package com.amazon.opendistroforelasticsearch.sql.protocol.response;
 
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTimestampValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType;
 import com.amazon.opendistroforelasticsearch.sql.executor.ExecutionEngine;
 import com.amazon.opendistroforelasticsearch.sql.executor.ExecutionEngine.Schema.Column;
 import java.util.Collection;
@@ -75,6 +77,15 @@ public Iterator<Object[]> iterator() {
         .iterator();
   }
 
+  public void setDatetimeFormat(String datetimeFormat) {
+    for (ExprValue exprValue : exprValues) {
+      for (ExprValue v : ExprValueUtils.getTupleValue(exprValue).values())
+        if (v.type() == ExprCoreType.TIMESTAMP) {
+          ((ExprTimestampValue) v).setDatetimeFormat(datetimeFormat);
+        }
+    }
+  }
+
   private String getColumnName(Column column) {
     return (column.getAlias() != null) ? column.getAlias() : column.getName();
   }

diff --git a/...amazon/opendistroforelasticsearch/sql/protocol/response/format/JdbcResponseFormatter.java b/...amazon/opendistroforelasticsearch/sql/protocol/response/format/JdbcResponseFormatter.java
@@ -36,12 +36,17 @@
  */
 public class JdbcResponseFormatter extends JsonResponseFormatter<QueryResult> {
 
-  public JdbcResponseFormatter(Style style) {
+  private final String datetimeFormat;
+
+  public JdbcResponseFormatter(Style style, String datetimeFormat) {
     super(style);
+    this.datetimeFormat = datetimeFormat;
   }
 
   @Override
   protected Object buildJsonObject(QueryResult response) {
+    response.setDatetimeFormat(datetimeFormat);
+
     JdbcResponse.JdbcResponseBuilder json = JdbcResponse.builder();
 
     // Fetch schema and data rows

diff --git a/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/domain/SQLQueryRequest.java b/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/domain/SQLQueryRequest.java
@@ -42,6 +42,7 @@ public class SQLQueryRequest {
   private static final Set<String> SUPPORTED_FIELDS = ImmutableSet.of(
       "query", "fetch_size", "parameters");
   private static final String QUERY_PARAMS_FORMAT = "format";
+  private static final String QUERY_PARAMS_DATETIMEFORMAT = "datetime_format";
   private static final String QUERY_PARAMS_SANITIZE = "sanitize";
 
   /**
@@ -142,6 +143,10 @@ private String getFormat(Map<String, String> params) {
     return "jdbc";
   }
 
+  public String getDatetimeFormat() {
+    return params.getOrDefault(QUERY_PARAMS_DATETIMEFORMAT, "always_include_time");
+  }
+
   private boolean shouldSanitize(Map<String, String> params) {
     if (params.containsKey(QUERY_PARAMS_SANITIZE)) {
       return Boolean.parseBoolean(params.get(QUERY_PARAMS_SANITIZE));