SQL: Whitelist SQL utility class for better scripting

x-pack/plugin/sql/build.gradle

@@ -5,7 +5,7 @@ esplugin {
     name 'x-pack-sql'
     description 'The Elasticsearch plugin that powers SQL for Elasticsearch'
     classname 'org.elasticsearch.xpack.sql.plugin.SqlPlugin'
-    extendedPlugins = ['x-pack-core']
+    extendedPlugins = ['x-pack-core', 'lang-painless']
 }
 
 configurations {
@@ -20,6 +20,7 @@ integTest.enabled = false
 
 dependencies {
     compileOnly "org.elasticsearch.plugin:x-pack-core:${version}"
+    compileOnly project(':modules:lang-painless')
     compile project('sql-proto')
     compile "org.elasticsearch.plugin:aggs-matrix-stats-client:${version}"
     compile "org.antlr:antlr4-runtime:4.5.3"

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/datetime/DateTimeFunction.java

@@ -51,8 +51,18 @@ public abstract class DateTimeFunction extends UnaryScalarFunction {
     protected final NodeInfo<DateTimeFunction> info() {
         return NodeInfo.create(this, ctorForInfo(), field(), timeZone());
     }
+
     protected abstract NodeInfo.NodeCtor2<Expression, TimeZone, DateTimeFunction> ctorForInfo();
 
+    @Override
+    protected TypeResolution resolveType() {
+        if (field().dataType() == DataType.DATE) {
+            return TypeResolution.TYPE_RESOLVED;
+        }
+        return new TypeResolution("Function [" + functionName() + "] cannot be applied on a non-date expression (["
+                + Expressions.name(field()) + "] of type [" + field().dataType().esType + "])");
+    }
+
     public TimeZone timeZone() {
         return timeZone;
     }
@@ -69,47 +79,24 @@ public Object fold() {
             return null;
         }
 
-        ZonedDateTime time = ZonedDateTime.ofInstant(
-            Instant.ofEpochMilli(folded.getMillis()), ZoneId.of(timeZone.getID()));
-        return time.get(chronoField());
+        return dateTimeChrono(folded.getMillis(), timeZone.getID(), chronoField().name());
     }
 
-    @Override
-    protected TypeResolution resolveType() {
-        if (field().dataType() == DataType.DATE) {
-            return TypeResolution.TYPE_RESOLVED;
-        }
-        return new TypeResolution("Function [" + functionName() + "] cannot be applied on a non-date expression (["
-                + Expressions.name(field()) + "] of type [" + field().dataType().esType + "])");
+    public static Integer dateTimeChrono(long millis, String tzId, String chronoName) {
+        ZonedDateTime time = ZonedDateTime.ofInstant(Instant.ofEpochMilli(millis), ZoneId.of(tzId));
+        return Integer.valueOf(time.get(ChronoField.valueOf(chronoName)));
     }
 
     @Override
     protected ScriptTemplate asScriptFrom(FieldAttribute field) {
         ParamsBuilder params = paramsBuilder();
 
         String template = null;
-        if (TimeZone.getTimeZone("UTC").equals(timeZone)) {
-            // TODO: it would be nice to be able to externalize the extract function and reuse the script across all extractors
-            template = formatTemplate("doc[{}].value.get" + extractFunction() + "()");
-            params.variable(field.name());
-        } else {
-            // TODO ewwww
-            /*
-             * This uses the Java 8 time API because Painless doesn't whitelist creation of new
-             * Joda classes.
-             *
-             * The actual script is
-             * ZonedDateTime.ofInstant(Instant.ofEpochMilli(<insert doc field>.value.millis),
-             *      ZoneId.of(<insert user tz>)).get(ChronoField.get(MONTH_OF_YEAR))
-             */
-
-            template = formatTemplate("ZonedDateTime.ofInstant(Instant.ofEpochMilli(doc[{}].value.millis), "
-                    + "ZoneId.of({})).get(ChronoField.valueOf({}))");
-            params.variable(field.name())
-                  .variable(timeZone.getID())
-                  .variable(chronoField().name());
-        }
-
+        template = formatTemplate("{sql}.dateTimeChrono(doc[{}].value.millis, {}, {})");
+        params.variable(field.name())
+              .variable(timeZone.getID())
+              .variable(chronoField().name());
+
         return new ScriptTemplate(template, params.build(), dataType());
     }
 
@@ -119,10 +106,6 @@ protected ScriptTemplate asScriptFrom(AggregateFunctionAttribute aggregate) {
         throw new UnsupportedOperationException();
     }
 
-    protected String extractFunction() {
-        return getClass().getSimpleName();
-    }
-
     /**
      * Used for generating the painless script version of this function when the time zone is not UTC
      */
@@ -164,4 +147,4 @@ public boolean equals(Object obj) {
     public int hashCode() {
         return Objects.hash(field(), timeZone);
     }
-}
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/script/ScriptTemplate.java

@@ -7,6 +7,7 @@
 
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptType;
+import org.elasticsearch.xpack.sql.expression.function.scalar.whitelist.InternalSqlScriptUtils;
 import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.util.StringUtils;
 
@@ -92,6 +93,6 @@ public String toString() {
     }
 
     public static String formatTemplate(String template) {
-        return template.replace("{}", "params.%s");
+        return template.replace("{sql}", InternalSqlScriptUtils.class.getSimpleName()).replace("{}", "params.%s");
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/whitelist/InternalSqlScriptUtils.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.sql.expression.function.scalar.whitelist;
+
+import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeFunction;
+
+/**
+ * Whitelisted class for SQL scripts.
+ * Acts as a registry of the various static methods used <b>internally</b> by the scalar functions
+ * (to simplify the whitelist definition).
+ */
+public final class InternalSqlScriptUtils {
+
+    private InternalSqlScriptUtils() {}
+
+    public static Integer dateTimeChrono(long millis, String tzId, String chronoName) {
+        return DateTimeFunction.dateTimeChrono(millis, tzId, chronoName);
+    }
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/SqlPainlessExtension.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.sql.plugin;
+
+import org.elasticsearch.painless.spi.PainlessExtension;
+import org.elasticsearch.painless.spi.Whitelist;
+import org.elasticsearch.painless.spi.WhitelistLoader;
+import org.elasticsearch.script.FilterScript;
+import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.script.SearchScript;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+import static java.util.Collections.singletonList;
+
+public class SqlPainlessExtension implements PainlessExtension {
+
+    private static final Whitelist WHITELIST = WhitelistLoader.loadFromResourceFiles(SqlPainlessExtension.class, "sql_whitelist.txt");
+
+    @Override
+    public Map<ScriptContext<?>, List<Whitelist>> getContextWhitelists() {
+        Map<ScriptContext<?>, List<Whitelist>> whitelist = new LinkedHashMap<>();
+        List<Whitelist> list = singletonList(WHITELIST);
+        whitelist.put(FilterScript.CONTEXT, list);
+        whitelist.put(SearchScript.AGGS_CONTEXT, list);
+        whitelist.put(SearchScript.CONTEXT, list);
+        whitelist.put(SearchScript.SCRIPT_SORT_CONTEXT, list);
+        return whitelist;
+    }
+}

x-pack/plugin/sql/src/main/resources/META-INF/services/org.elasticsearch.painless.spi.PainlessExtension

@@ -0,0 +1 @@
+org.elasticsearch.xpack.sql.plugin.SqlPainlessExtension

x-pack/plugin/sql/src/main/resources/org/elasticsearch/xpack/sql/plugin/sql_whitelist.txt

@@ -0,0 +1,12 @@
+#
+# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+# or more contributor license agreements. Licensed under the Elastic License;
+# you may not use this file except in compliance with the Elastic License.
+#
+
+# This file contains a whitelist for SQL specific utilities available inside SQL scripting
+
+class org.elasticsearch.xpack.sql.expression.function.scalar.whitelist.InternalSqlScriptUtils {
+
+  Integer dateTimeChrono(long, String, String)
+}

x-pack/plugin/src/test/resources/rest-api-spec/test/sql/whitelist.yml

@@ -0,0 +1,27 @@
+---
+"SQL methods available in Painless":
+    - do:
+        bulk:
+          refresh: true
+          body:
+            - index:
+                _index: test
+                _type:  doc
+                _id:    1
+            - str: test1
+              int: 1
+    - do:
+        index: test
+        search:
+            body:
+                query:
+                    match_all: {}
+                script_fields:
+                    sNum1:
+                        script:
+                            source: "return 
+                            InternalSqlScriptUtils.dateTimeChrono(doc['int'].value, 'UTC', 'YEAR')"
+                            lang: painless
+
+    - match: { hits.total: 1 }
+    - match: { hits.hits.0.fields.sNum1.0: 1970 }