Adding RBAC authorization checks for multi-stage query engine

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

@@ -30,6 +30,7 @@
 import javax.ws.rs.core.HttpHeaders;
 import javax.ws.rs.core.Response;
 import org.apache.calcite.jdbc.CalciteSchemaBuilder;
+import org.apache.pinot.broker.api.AccessControl;
 import org.apache.pinot.broker.api.RequesterIdentity;
 import org.apache.pinot.broker.broker.AccessControlFactory;
 import org.apache.pinot.broker.querylog.QueryLogger;
@@ -51,6 +52,8 @@
 import org.apache.pinot.common.utils.ExceptionUtils;
 import org.apache.pinot.common.utils.config.QueryOptionsUtils;
 import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.core.auth.Actions;
+import org.apache.pinot.core.auth.TargetType;
 import org.apache.pinot.core.query.reduce.ExecutionStatsAggregator;
 import org.apache.pinot.core.transport.ServerInstance;
 import org.apache.pinot.query.QueryEnvironment;
@@ -128,7 +131,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
           queryPlanResult = _queryEnvironment.explainQuery(query, sqlNodeAndOptions, requestId);
           String plan = queryPlanResult.getExplainPlan();
           Set<String> tableNames = queryPlanResult.getTableNames();
-          if (!hasTableAccess(requesterIdentity, tableNames, requestContext)) {
+          if (!hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders)) {
             throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN);
           }
 
@@ -164,7 +167,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
     updatePhaseTimingForTables(tableNames, BrokerQueryPhase.REQUEST_COMPILATION, compilationTimeNs);
 
     // Validate table access.
-    if (!hasTableAccess(requesterIdentity, tableNames, requestContext)) {
+    if (!hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders)) {
       throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN);
     }
     updatePhaseTimingForTables(tableNames, BrokerQueryPhase.AUTHORIZATION, System.nanoTime() - compilationEndTimeNs);
@@ -251,8 +254,10 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
    * Validates whether the requester has access to all the tables.
    */
   private boolean hasTableAccess(RequesterIdentity requesterIdentity, Set<String> tableNames,
-      RequestContext requestContext) {
-    boolean hasAccess = _accessControlFactory.create().hasAccess(requesterIdentity, tableNames);
+      RequestContext requestContext, HttpHeaders httpHeaders) {
+    AccessControl accessControl = _accessControlFactory.create();
+    boolean hasAccess = accessControl.hasAccess(requesterIdentity, tableNames) && tableNames.stream()
+        .allMatch(table -> accessControl.hasAccess(httpHeaders, TargetType.TABLE, table, Actions.Table.QUERY));
     if (!hasAccess) {
       _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1);
       LOGGER.warn("Access denied for requestId {}", requestContext.getRequestId());