SQL: Whitelist SQL utility class for better scripting (#30681)

Add SQL class for reusing code inside SQL functions within Painless

Fix #29832

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/licenses/antlr4-runtime-4.5.3.jar.sha1

@@ -1 +0,0 @@
-2609e36f18f7e8d593cc1cddfb2ac776dc96b8e0

x-pack/plugin/sql/licenses/antlr4-runtime-LICENSE.txt

@@ -1,26 +0,0 @@
-[The "BSD license"]
-Copyright (c) 2015 Terence Parr, Sam Harwell
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions
-are met:
-
- 1. Redistributions of source code must retain the above copyright
-    notice, this list of conditions and the following disclaimer.
- 2. Redistributions in binary form must reproduce the above copyright
-    notice, this list of conditions and the following disclaimer in the
-    documentation and/or other materials provided with the distribution.
- 3. The name of the author may not be used to endorse or promote products
-    derived from this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
-IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
-OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
-IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
-INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
-NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
-THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

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,18 +79,12 @@ public abstract class DateTimeFunction extends UnaryScalarFunction {
             return null;
         }
 
-        ZonedDateTime time = ZonedDateTime.ofInstant(
+        return dateTimeChrono(folded.getMillis(), timeZone.getID(), chronoField().name());
-            Instant.ofEpochMilli(folded.getMillis()), ZoneId.of(timeZone.getID()));
-        return time.get(chronoField());
     }
 
-    @Override
+    public static Integer dateTimeChrono(long millis, String tzId, String chronoName) {
-    protected TypeResolution resolveType() {
+        ZonedDateTime time = ZonedDateTime.ofInstant(Instant.ofEpochMilli(millis), ZoneId.of(tzId));
-        if (field().dataType() == DataType.DATE) {
+        return Integer.valueOf(time.get(ChronoField.valueOf(chronoName)));
-            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 + "])");
     }
 
     @Override
@@ -88,27 +92,10 @@ public abstract class DateTimeFunction extends UnaryScalarFunction {
         ParamsBuilder params = paramsBuilder();
 
         String template = null;
-        if (TimeZone.getTimeZone("UTC").equals(timeZone)) {
+        template = formatTemplate("{sql}.dateTimeChrono(doc[{}].value.millis, {}, {})");
-            // 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());
-        }
         
         return new ScriptTemplate(template, params.build(), dataType());
     }
@@ -119,10 +106,6 @@ public abstract class DateTimeFunction extends UnaryScalarFunction {
         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
      */

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

@@ -7,6 +7,7 @@ package org.elasticsearch.xpack.sql.expression.function.scalar.script;
 
 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 class ScriptTemplate {
     }
 
     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.HashMap;
+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 HashMap<>();
+        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)
+}