diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java index 192ddd11c1ea..ae474e28d879 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java +++ b/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 selectList = pinotQuery.getSelectList(); + for (Expression expression : selectList) { + rejectGroovyQuery(expression); + } + List 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 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 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. */ diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java index a713eb4bdcc7..829017d52bf2 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java +++ b/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); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java index 012996d5ca14..b51840075dfe 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java +++ b/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(); } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index a85c75731614..bf2922b73e2b 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/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";