Adding broker level config for disabling Pinot queries with Groovy

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java

@@ -76,6 +76,7 @@
 import org.apache.pinot.common.response.broker.ResultTable;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils;
 import org.apache.pinot.core.query.optimizer.QueryOptimizer;
 import org.apache.pinot.core.requesthandler.PinotQueryParserFactory;
@@ -128,6 +129,7 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
   private final RateLimiter _numDroppedLogRateLimiter;
   private final AtomicInteger _numDroppedLog;
 
+  private final boolean _disableGroovy;
   private final int _defaultHllLog2m;
   private final boolean _enableQueryLimitOverride;
   private final boolean _enableDistinctCountBitmapOverride;
@@ -142,6 +144,7 @@ public BaseBrokerRequestHandler(PinotConfiguration config, RoutingManager routin
     _tableCache = tableCache;
     _brokerMetrics = brokerMetrics;
 
+    _disableGroovy = _config.getProperty(CommonConstants.Broker.DISABLE_GROOVY, false);
     _defaultHllLog2m = _config.getProperty(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M_KEY,
         CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M);
     _enableQueryLimitOverride = _config.getProperty(Broker.CONFIG_OF_ENABLE_QUERY_LIMIT_OVERRIDE, false);
@@ -270,6 +273,9 @@ private BrokerResponseNative handleSQLRequest(long requestId, String query, Json
       LOGGER.warn("Caught exception while updating column names in request {}: {}, {}", requestId, query,
           e.getMessage());
     }
+    if (_disableGroovy) {
+      rejectGroovyQuery(pinotQuery);
+    }
     if (_defaultHllLog2m > 0) {
       handleHLLLog2mOverride(pinotQuery, _defaultHllLog2m);
     }
@@ -1168,6 +1174,57 @@ private static void handleHLLLog2mOverride(PinotQuery pinotQuery, int hllLog2mOv
     }
   }
 
+  /**
+   * Verifies that no groovy is present in the PinotQuery when disabled.
+   */
+  @VisibleForTesting
+  static void rejectGroovyQuery(PinotQuery pinotQuery) {
+    List<Expression> selectList = pinotQuery.getSelectList();
+    for (Expression expression : selectList) {
+      rejectGroovyQuery(expression);
+    }
+    List<Expression> orderByList = pinotQuery.getOrderByList();
+    if (orderByList != null) {
+      for (Expression expression : orderByList) {
+        // NOTE: Order-by is always a Function with the ordering of the Expression
+        rejectGroovyQuery(expression.getFunctionCall().getOperands().get(0));
+      }
+    }
+    Expression havingExpression = pinotQuery.getHavingExpression();
+    if (havingExpression != null) {
+      rejectGroovyQuery(havingExpression);
+    }
+    Expression filterExpression = pinotQuery.getFilterExpression();
+    if (filterExpression != null) {
+      rejectGroovyQuery(filterExpression);
+    }
+    List<Expression> groupByList = pinotQuery.getGroupByList();
+    if (groupByList != null) {
+      for (Expression expression : groupByList) {
+        rejectGroovyQuery(expression);
+      }
+    }
+  }
+
+  private static void rejectGroovyQuery(Expression expression) {
+    Function functionCall = expression.getFunctionCall();
+    if (functionCall == null) {
+      return;
+    }
+
+    if (TransformFunctionFactory.canonicalize(functionCall.getOperator())
+        .equals(TransformFunctionType.GROOVY.getName())) {
+      throw new BadQueryRequestException("Groovy transform functions are disabled for queries");
+    }
+
+    List<Expression> operands = functionCall.getOperands();
+    if (operands != null) {
+      for (Expression operandExpression : operands) {
+        rejectGroovyQuery(operandExpression);
+      }
+    }
+  }
+
   /**
    * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given SQL expression.
    */

pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java

@@ -181,6 +181,40 @@ public void testUnsupportedNonExistColumnsQueries() {
     testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "COL1", "b", "B", "c", "C"), sql);
   }
 
+  @Test
+  public void testRejectGroovyQuery() {
+    testRejectGroovyQuery(
+        "SELECT groovy('{\"returnType\":\"INT\",\"isSingleValue\":true}', 'arg0 + arg1', colA, colB) FROM foo", true);
+    testRejectGroovyQuery(
+        "SELECT GROOVY('{\"returnType\":\"INT\",\"isSingleValue\":true}', 'arg0 + arg1', colA, colB) FROM foo", true);
+    testRejectGroovyQuery(
+        "SELECT groo_vy('{\"returnType\":\"INT\",\"isSingleValue\":true}', 'arg0 + arg1', colA, colB) FROM foo", true);
+    testRejectGroovyQuery(
+        "SELECT foo FROM bar WHERE GROOVY('{\"returnType\":\"STRING\",\"isSingleValue\":true}', 'arg0 + arg1', colA,"
+            + " colB) = 'foobarval'", true);
+    testRejectGroovyQuery(
+        "SELECT COUNT(colA) FROM bar GROUP BY GROOVY('{\"returnType\":\"STRING\",\"isSingleValue\":true}', "
+            + "'arg0 + arg1', colA, colB)", true);
+    testRejectGroovyQuery(
+        "SELECT foo FROM bar HAVING GROOVY('{\"returnType\":\"STRING\",\"isSingleValue\":true}', 'arg0 + arg1', colA,"
+            + " colB) = 'foobarval'", true);
+
+    testRejectGroovyQuery("SELECT foo FROM bar", false);
+  }
+
+  private void testRejectGroovyQuery(String query, boolean queryContainsGroovy) {
+    PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+
+    try {
+      BaseBrokerRequestHandler.rejectGroovyQuery(pinotQuery);
+      if (queryContainsGroovy) {
+        Assert.fail("Query should have failed since groovy was found in query: " + pinotQuery);
+      }
+    } catch (Exception e) {
+      Assert.assertEquals(e.getMessage(), "Groovy transform functions are disabled for queries");
+    }
+  }
+
   private void testUnsupportedPQLQuery(String query, String errorMessage) {
     try {
       BrokerRequest brokerRequest = PQL_COMPILER.compileToBrokerRequest(query);

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java

@@ -131,7 +131,6 @@ private TransformFunctionFactory() {
           put(canonicalize(TransformFunctionType.AND.getName().toLowerCase()), AndOperatorTransformFunction.class);
           put(canonicalize(TransformFunctionType.OR.getName().toLowerCase()), OrOperatorTransformFunction.class);
 
-
           // geo functions
           // geo constructors
           put(canonicalize(TransformFunctionType.ST_GEOG_FROM_TEXT.getName().toLowerCase()),
@@ -269,7 +268,13 @@ public static TransformFunction get(@Nullable QueryContext queryContext, Express
     }
   }
 
-  private static String canonicalize(String functionName) {
+  /**
+   * Converts the transform function name into its canonical form
+   *
+   * @param functionName Name of the transform function
+   * @return canonicalized transform function name
+   */
+  public static String canonicalize(String functionName) {
     return StringUtils.remove(functionName, '_').toLowerCase();
   }
 }

pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java

@@ -228,6 +228,8 @@ public static class Broker {
     //Set to true to load all services tagged and compiled with hk2-metadata-generator. Default to False
     public static final String BROKER_SERVICE_AUTO_DISCOVERY = "pinot.broker.service.auto.discovery";
 
+    public static final String DISABLE_GROOVY = "pinot.broker.disable.query.groovy";
+
     public static class Request {
       public static final String PQL = "pql";
       public static final String SQL = "sql";