Fix UnsupportedOperationException when RLS filter combines with expression override (#19008)

Author: navina

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

@@ -19,6 +19,7 @@
 package org.apache.pinot.broker.requesthandler;
 
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -54,8 +55,11 @@
 import org.apache.pinot.materializedview.rewrite.MaterializedViewQueryRewriteEngine;
 import org.apache.pinot.materializedview.rewrite.MaterializedViewRewritePlan;
 import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
+import org.apache.pinot.spi.auth.AuthorizationResult;
 import org.apache.pinot.spi.auth.TableAuthorizationResult;
+import org.apache.pinot.spi.auth.TableRowColAccessResult;
 import org.apache.pinot.spi.auth.TableRowColAccessResultImpl;
+import org.apache.pinot.spi.auth.broker.RequesterIdentity;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TenantConfig;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
@@ -857,6 +861,146 @@ protected BrokerResponseNative processMaterializedViewSplitBrokerRequest(long re
   /// org.apache.pinot.materializedview.handler.DefaultMaterializedViewHandler#attachFilter; their
   /// regression coverage lives in DefaultMaterializedViewHandlerTest in pinot-materialized-view.
 
+  /// Bug: `RlsFiltersRewriter` combines the RLS predicate with an existing WHERE clause via
+  /// `RequestUtils.getFunctionExpression(AND, List.of(...))`. The `List<Expression>` overload used
+  /// to store the list as-is, so the AND function's operand list was the immutable list returned by
+  /// `List.of(...)`.
+  ///
+  /// When the table also carries an expression-override map, `handleExpressionOverride` recurses
+  /// into the filter tree and calls `function.getOperands().replaceAll(...)` on that AND function,
+  /// throwing `UnsupportedOperationException` on the immutable operand list.
+  ///
+  /// This test reproduces the crash by enabling RLS (which yields a combined AND filter) on a query
+  /// that already has a WHERE clause, for a table whose expression-override map is non-null. Before
+  /// the fix, `handleRequest` throws `UnsupportedOperationException`; after the fix it completes and
+  /// reports that RLS filters were applied.
+  @Test
+  public void testRlsFilterCombinedWithExpressionOverrideDoesNotThrow()
+      throws Exception {
+    String rawTable = "faroEvents";
+    String offlineTable = "faroEvents_OFFLINE";
+
+    // Query with an existing WHERE clause so the RLS rewriter builds a combined AND filter.
+    String userSql = "SELECT COUNT(*) FROM faroEvents WHERE ts = 'x'";
+
+    Schema schema = new Schema.SchemaBuilder()
+        .setSchemaName(rawTable)
+        .addSingleValueDimension("ts", DataType.STRING)
+        .addMetric("revenue", DataType.DOUBLE)
+        .build();
+
+    TableCache tableCache = mock(TableCache.class);
+    when(tableCache.getActualTableName(rawTable)).thenReturn(rawTable);
+    when(tableCache.getSchema(rawTable)).thenReturn(schema);
+    when(tableCache.getColumnNameMap(anyString())).thenReturn(Map.of("ts", "ts", "revenue", "revenue"));
+    TableConfig tableCfg = mock(TableConfig.class);
+    when(tableCfg.getTenantConfig()).thenReturn(new TenantConfig("t_BROKER", "t_SERVER", null));
+    when(tableCache.getTableConfig(offlineTable)).thenReturn(tableCfg);
+
+    // Non-null expression-override map for the table. A non-null map is enough to make
+    // handleExpressionOverride recurse into the filter tree and call replaceAll on the AND operands.
+    Expression overrideKey = CalciteSqlParser.compileToExpression("ts");
+    Expression overrideValue = CalciteSqlParser.compileToExpression("ts");
+    when(tableCache.getExpressionOverrideMap(offlineTable)).thenReturn(Map.of(overrideKey, overrideValue));
+
+    BrokerRoutingManager routingManager = mock(BrokerRoutingManager.class);
+    when(routingManager.routingExists(offlineTable)).thenReturn(true);
+    when(routingManager.getQueryTimeoutMs(anyString())).thenReturn(10000L);
+    RoutingTable rt = mock(RoutingTable.class);
+    when(rt.getServerInstanceToSegmentsMap()).thenReturn(
+        Map.of(new ServerInstance(new InstanceConfig("server01_9000")),
+            new SegmentsToQuery(List.of("seg01"), List.of())));
+    when(routingManager.getRoutingTable(any(), Mockito.anyLong())).thenReturn(rt);
+
+    QueryQuotaManager quotaManager = mock(QueryQuotaManager.class);
+    when(quotaManager.acquire(anyString())).thenReturn(true);
+    when(quotaManager.acquireDatabase(anyString())).thenReturn(true);
+    when(quotaManager.acquireApplication(anyString())).thenReturn(true);
+
+    // AccessControl that allows everything and returns a non-empty RLS filter for the table.
+    AccessControl accessControl = new AccessControl() {
+      @Override
+      public AuthorizationResult authorize(RequesterIdentity identity, BrokerRequest request) {
+        return TableAuthorizationResult.success();
+      }
+
+      @Override
+      public TableAuthorizationResult authorize(RequesterIdentity identity, Set<String> tables) {
+        return TableAuthorizationResult.success();
+      }
+
+      @Override
+      public TableRowColAccessResult getRowColFilters(RequesterIdentity identity, String tableWithType) {
+        return new TableRowColAccessResultImpl(List.of("ts = 'allowed'"));
+      }
+    };
+    AccessControlFactory accessControlFactory = mock(AccessControlFactory.class);
+    when(accessControlFactory.create()).thenReturn(accessControl);
+
+    BrokerMetrics.register(mock(BrokerMetrics.class));
+    PinotConfiguration config = new PinotConfiguration(
+        Map.of(Broker.CONFIG_OF_BROKER_ENABLE_ROW_COLUMN_LEVEL_AUTH, "true"));
+    BrokerQueryEventListenerFactory.init(config);
+
+    AtomicReference<BrokerRequest> capturedServerRequest = new AtomicReference<>();
+    BaseSingleStageBrokerRequestHandler handler =
+        new BaseSingleStageBrokerRequestHandler(config, "broker1", new BrokerRequestIdGenerator(),
+            routingManager, accessControlFactory, quotaManager, tableCache,
+            ThreadAccountantUtils.getNoOpAccountant(), null) {
+          @Override
+          public void start() {
+          }
+
+          @Override
+          public void shutDown() {
+          }
+
+          @Override
+          protected BrokerResponseNative processBrokerRequest(long requestId,
+              BrokerRequest originalBrokerRequest, BrokerRequest serverBrokerRequest,
+              TableRouteInfo route, long timeoutMs, ServerStats serverStats,
+              RequestContext requestContext) {
+            capturedServerRequest.set(serverBrokerRequest);
+            return BrokerResponseNative.empty();
+          }
+        };
+
+    // Before the fix this throws UnsupportedOperationException from replaceAll on the immutable
+    // AND operand list produced by the RLS rewriter.
+    BrokerResponseNative response = (BrokerResponseNative) handler.handleRequest(userSql);
+
+    Assert.assertNotNull(response, "handleRequest must return a response, not throw");
+    Assert.assertTrue(response.getRLSFiltersApplied(), "response should report that RLS filters were applied");
+
+    // The server query must carry BOTH the RLS predicate and the original WHERE clause, combined
+    // under a single AND. This guards against a future regression that silently drops an operand.
+    BrokerRequest serverRequest = capturedServerRequest.get();
+    Assert.assertNotNull(serverRequest, "server query should have been captured");
+    Expression filter = serverRequest.getPinotQuery().getFilterExpression();
+    Assert.assertNotNull(filter, "server query must have a filter expression");
+    Function filterFunction = filter.getFunctionCall();
+    Assert.assertNotNull(filterFunction, "combined filter must be a function");
+    Assert.assertEquals(filterFunction.getOperator(), FilterKind.AND.name(),
+        "RLS predicate and existing WHERE clause must be combined under AND");
+    Assert.assertEquals(filterFunction.getOperands().size(), 2, "combined AND must retain both operands");
+    // Collect the RHS string literals of both EQUALS operands: one must be the RLS predicate value
+    // ('allowed') and the other the original WHERE-clause value ('x').
+    Set<String> literalValues = new HashSet<>();
+    for (Expression operand : filterFunction.getOperands()) {
+      Function eq = operand.getFunctionCall();
+      Assert.assertNotNull(eq, "each AND operand must be a comparison function");
+      for (Expression eqOperand : eq.getOperands()) {
+        if (eqOperand.getLiteral() != null) {
+          literalValues.add(eqOperand.getLiteral().getStringValue());
+        }
+      }
+    }
+    Assert.assertTrue(literalValues.contains("allowed"),
+        "combined filter must include the RLS predicate, got literals: " + literalValues);
+    Assert.assertTrue(literalValues.contains("x"),
+        "combined filter must retain the original WHERE clause, got literals: " + literalValues);
+  }
+
   /**
    * Pins the security-style defense that a user-supplied `materializedViewRewrite=true` query
    * option (e.g. via `SET materializedViewRewrite='true'`) is stripped at the broker entry

pinot-common/src/main/java/org/apache/pinot/common/utils/request/RequestUtils.java

@@ -487,9 +487,16 @@ public static String getLiteralString(Expression expression) {
     return getLiteralString(literal);
   }
 
+  /// Creates a `Function` with the given operands. The operand list stored in the returned function
+  /// is always a mutable `ArrayList`: if `operands` is not already an `ArrayList` (e.g. an immutable
+  /// `List.of(...)`), it is copied into one. Downstream query rewriters and filter optimizers mutate
+  /// operands in place (via `getOperands().replaceAll(...)`, `set(...)`, or `add(...)`), so an
+  /// immutable list would otherwise throw `UnsupportedOperationException` far from where it was
+  /// created.
   public static Function getFunction(String canonicalName, List<Expression> operands) {
     Function function = new Function(canonicalName);
-    function.setOperands(operands);
+    // Ensure a mutable ArrayList so downstream rewriters can modify operands in place.
+    function.setOperands(operands instanceof ArrayList ? operands : new ArrayList<>(operands));
     return function;
   }
 

pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/RlsFiltersRewriter.java

@@ -18,7 +18,6 @@
  */
 package org.apache.pinot.sql.parsers.rewriter;
 
-import java.util.List;
 import java.util.Map;
 import org.apache.commons.collections4.MapUtils;
 import org.apache.logging.log4j.util.Strings;
@@ -69,7 +68,7 @@ public PinotQuery rewrite(PinotQuery pinotQuery) {
     Expression existingFilterExpression = pinotQuery.getFilterExpression();
     if (existingFilterExpression != null) {
       Expression combinedFilterExpression =
-          RequestUtils.getFunctionExpression(FilterKind.AND.name(), List.of(expression, existingFilterExpression));
+          RequestUtils.getFunctionExpression(FilterKind.AND.name(), expression, existingFilterExpression);
       pinotQuery.setFilterExpression(combinedFilterExpression);
     } else {
       pinotQuery.setFilterExpression(expression);

pinot-common/src/test/java/org/apache/pinot/common/utils/request/RequestUtilsTest.java

@@ -52,6 +52,35 @@ public void testNullLiteralParsing() {
     assertTrue(nullExpr.getLiteral().getNullValue());
   }
 
+  /// Pins the factory contract that `RequestUtils.getFunction` always returns a function whose
+  /// operand list is mutable, even when the caller passes an immutable list. Downstream query
+  /// rewriters and filter optimizers mutate operands in place; a regression here would resurface as
+  /// an `UnsupportedOperationException` far from this factory (e.g. the broker RLS +
+  /// expression-override path).
+  @Test
+  public void testGetFunctionReturnsMutableOperandList() {
+    Expression a = RequestUtils.getLiteralExpression(1L);
+    Expression b = RequestUtils.getLiteralExpression(2L);
+
+    // List overload with a genuinely immutable input list must be defensively copied.
+    Function fromList = RequestUtils.getFunction("and", List.of(a, b));
+    fromList.getOperands().replaceAll(o -> o);
+    fromList.getOperands().set(0, b);
+    fromList.getOperands().add(a);
+    assertEquals(fromList.getOperands().size(), 3);
+
+    // Varargs and single-operand overloads delegate to the List overload and share the guarantee.
+    Function fromVarargs = RequestUtils.getFunction("and", a, b);
+    fromVarargs.getOperands().replaceAll(o -> o);
+    fromVarargs.getOperands().add(b);
+    assertEquals(fromVarargs.getOperands().size(), 3);
+
+    Function fromSingle = RequestUtils.getFunction("not", a);
+    fromSingle.getOperands().replaceAll(o -> o);
+    fromSingle.getOperands().add(b);
+    assertEquals(fromSingle.getOperands().size(), 2);
+  }
+
   @Test
   public void testGetLiteralExpressionForPrimitiveLong() {
     Expression literalExpression = RequestUtils.getLiteralExpression(4500L);