Add policy enforcer to sanity check on policy in query execution

benchmarks/src/test/java/org/apache/druid/benchmark/query/InPlanningBenchmark.java

@@ -33,6 +33,7 @@
 import org.apache.druid.math.expr.ExpressionProcessing;
 import org.apache.druid.query.DruidProcessingConfig;
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
+import org.apache.druid.query.policy.NoopPolicyEnforcer;
 import org.apache.druid.segment.AutoTypeColumnSchema;
 import org.apache.druid.segment.IndexSpec;
 import org.apache.druid.segment.QueryableIndex;
@@ -202,6 +203,7 @@ public void setup() throws JsonProcessingException
         CalciteTests.createJoinableFactoryWrapper(),
         CatalogResolver.NULL_RESOLVER,
         new AuthConfig(),
+        NoopPolicyEnforcer.instance(),
         new DruidHookDispatcher()
     );
 

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlBaseBenchmark.java

@@ -60,6 +60,7 @@
 import org.apache.druid.query.aggregation.datasketches.theta.sql.ThetaSketchEstimateOperatorConversion;
 import org.apache.druid.query.aggregation.datasketches.tuple.ArrayOfDoublesSketchModule;
 import org.apache.druid.query.lookup.LookupExtractor;
+import org.apache.druid.query.policy.NoopPolicyEnforcer;
 import org.apache.druid.segment.IncrementalIndexSegment;
 import org.apache.druid.segment.IndexSpec;
 import org.apache.druid.segment.PhysicalSegmentInspector;
@@ -452,6 +453,7 @@ public static Pair<PlannerFactory, SqlEngine> createSqlSystem(
         new JoinableFactoryWrapper(QueryFrameworkUtils.createDefaultJoinableFactory(injector)),
         CatalogResolver.NULL_RESOLVER,
         new AuthConfig(),
+        NoopPolicyEnforcer.instance(),
         new DruidHookDispatcher()
     );
 

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlVsNativeBenchmark.java

@@ -32,6 +32,7 @@
 import org.apache.druid.query.dimension.DefaultDimensionSpec;
 import org.apache.druid.query.groupby.GroupByQuery;
 import org.apache.druid.query.groupby.ResultRow;
+import org.apache.druid.query.policy.NoopPolicyEnforcer;
 import org.apache.druid.segment.QueryableIndex;
 import org.apache.druid.segment.generator.GeneratorBasicSchemas;
 import org.apache.druid.segment.generator.GeneratorSchemaInfo;
@@ -129,6 +130,7 @@ public void setup()
         CalciteTests.createJoinableFactoryWrapper(),
         CatalogResolver.NULL_RESOLVER,
         new AuthConfig(),
+        NoopPolicyEnforcer.instance(),
         new DruidHookDispatcher()
     );
     groupByQuery = GroupByQuery

extensions-core/multi-stage-query/pom.xml

@@ -334,6 +334,10 @@
             <version>${project.parent.version}</version>
             <scope>test</scope>
         </dependency>
+      <dependency>
+        <groupId>com.google.inject.extensions</groupId>
+        <artifactId>guice-testlib</artifactId>
+      </dependency>
     </dependencies>
 
     <build>

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java

@@ -32,6 +32,7 @@
 import org.apache.druid.msq.kernel.StageId;
 import org.apache.druid.msq.querykit.DataSegmentProvider;
 import org.apache.druid.query.groupby.GroupingEngine;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.segment.IndexIO;
 import org.apache.druid.segment.IndexMergerV9;
 import org.apache.druid.segment.SegmentWrangler;
@@ -78,6 +79,12 @@ public DartFrameContext(
     this.storageParameters = storageParameters;
   }
 
+  @Override
+  public PolicyEnforcer policyEnforcer()
+  {
+    return workerContext.policyEnforcer();
+  }
+
   @Override
   public SegmentWrangler segmentWrangler()
   {

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContext.java

@@ -44,6 +44,7 @@
 import org.apache.druid.query.DruidProcessingConfig;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.groupby.GroupingEngine;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.segment.SegmentWrangler;
 import org.apache.druid.server.DruidNode;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
@@ -61,6 +62,7 @@ public class DartWorkerContext implements WorkerContext
   private final WorkerId workerId;
   private final DruidNode selfNode;
   private final ObjectMapper jsonMapper;
+  private final PolicyEnforcer policyEnforcer;
   private final Injector injector;
   private final DartWorkerClient workerClient;
   private final DruidProcessingConfig processingConfig;
@@ -84,6 +86,7 @@ public class DartWorkerContext implements WorkerContext
       final String controllerHost,
       final DruidNode selfNode,
       final ObjectMapper jsonMapper,
+      final PolicyEnforcer policyEnforcer,
       final Injector injector,
       final DartWorkerClient workerClient,
       final DruidProcessingConfig processingConfig,
@@ -102,6 +105,7 @@ public class DartWorkerContext implements WorkerContext
     this.workerId = WorkerId.fromDruidNode(selfNode, queryId);
     this.selfNode = selfNode;
     this.jsonMapper = jsonMapper;
+    this.policyEnforcer = policyEnforcer;
     this.injector = injector;
     this.workerClient = workerClient;
     this.processingConfig = processingConfig;
@@ -133,6 +137,12 @@ public ObjectMapper jsonMapper()
     return jsonMapper;
   }
 
+  @Override
+  public PolicyEnforcer policyEnforcer()
+  {
+    return policyEnforcer;
+  }
+
   @Override
   public Injector injector()
   {

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java

@@ -38,6 +38,7 @@
 import org.apache.druid.query.DruidProcessingConfig;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.groupby.GroupingEngine;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.rpc.ServiceClientFactory;
 import org.apache.druid.segment.SegmentWrangler;
 import org.apache.druid.server.DruidNode;
@@ -52,6 +53,7 @@ public class DartWorkerFactoryImpl implements DartWorkerFactory
   private final DruidNode selfNode;
   private final ObjectMapper jsonMapper;
   private final ObjectMapper smileMapper;
+  private final PolicyEnforcer policyEnforcer;
   private final Injector injector;
   private final ServiceClientFactory serviceClientFactory;
   private final DruidProcessingConfig processingConfig;
@@ -67,6 +69,7 @@ public DartWorkerFactoryImpl(
       @Self DruidNode selfNode,
       @Json ObjectMapper jsonMapper,
       @Smile ObjectMapper smileMapper,
+      PolicyEnforcer policyEnforcer,
       Injector injector,
       @EscalatedGlobal ServiceClientFactory serviceClientFactory,
       DruidProcessingConfig processingConfig,
@@ -81,6 +84,7 @@ public DartWorkerFactoryImpl(
     this.selfNode = selfNode;
     this.jsonMapper = jsonMapper;
     this.smileMapper = smileMapper;
+    this.policyEnforcer = policyEnforcer;
     this.injector = injector;
     this.serviceClientFactory = serviceClientFactory;
     this.processingConfig = processingConfig;
@@ -100,6 +104,7 @@ public Worker build(String queryId, String controllerHost, File tempDir, QueryCo
         controllerHost,
         selfNode,
         jsonMapper,
+        policyEnforcer,
         injector,
         new DartWorkerClientImpl(queryId, serviceClientFactory, smileMapper, null),
         processingConfig,

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerContext.java

@@ -27,6 +27,7 @@
 import org.apache.druid.msq.kernel.FrameProcessorFactory;
 import org.apache.druid.msq.kernel.WorkOrder;
 import org.apache.druid.msq.util.MultiStageQueryContext;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.server.DruidNode;
 
 import java.io.File;
@@ -51,6 +52,8 @@ public interface WorkerContext
 
   ObjectMapper jsonMapper();
 
+  PolicyEnforcer policyEnforcer();
+
   // Using an Injector directly because tasks do not have a way to provide their own Guice modules.
   Injector injector();
 

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java

@@ -31,6 +31,7 @@
 import org.apache.druid.msq.kernel.StageId;
 import org.apache.druid.msq.querykit.DataSegmentProvider;
 import org.apache.druid.query.groupby.GroupingEngine;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.segment.IndexIO;
 import org.apache.druid.segment.IndexMergerV9;
 import org.apache.druid.segment.SegmentWrangler;
@@ -73,6 +74,12 @@ public IndexerFrameContext(
     this.dataServerQueryHandlerFactory = dataServerQueryHandlerFactory;
   }
 
+  @Override
+  public PolicyEnforcer policyEnforcer()
+  {
+    return context.policyEnforcer();
+  }
+
   @Override
   public SegmentWrangler segmentWrangler()
   {

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java

@@ -52,6 +52,7 @@
 import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryToolChestWarehouse;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.rpc.ServiceClientFactory;
 import org.apache.druid.rpc.ServiceLocations;
 import org.apache.druid.rpc.ServiceLocator;
@@ -191,6 +192,12 @@ public ObjectMapper jsonMapper()
     return toolbox.getJsonMapper();
   }
 
+  @Override
+  public PolicyEnforcer policyEnforcer()
+  {
+    return toolbox.getPolicyEnforcer();
+  }
+
   @Override
   public Injector injector()
   {

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/kernel/FrameContext.java

@@ -27,6 +27,7 @@
 import org.apache.druid.msq.exec.WorkerStorageParameters;
 import org.apache.druid.msq.querykit.DataSegmentProvider;
 import org.apache.druid.query.groupby.GroupingEngine;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.segment.IndexIO;
 import org.apache.druid.segment.IndexMergerV9;
 import org.apache.druid.segment.SegmentWrangler;
@@ -44,6 +45,8 @@
  */
 public interface FrameContext extends Closeable
 {
+  PolicyEnforcer policyEnforcer();
+
   SegmentWrangler segmentWrangler();
 
   GroupingEngine groupingEngine();

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessorFactory.java

@@ -162,8 +162,8 @@ public ProcessorsAndChannels<Object, Long> makeProcessors(
     final ProcessorManager processorManager;
 
     if (segmentMapFnProcessor == null) {
-      final Function<SegmentReference, SegmentReference> segmentMapFn =
-          query.getDataSource().createSegmentMapFunction(query);
+      final Function<SegmentReference, SegmentReference> segmentMapFn = ExecutionVertex.of(query)
+                                                                                       .createSegmentMapFunction(frameContext.policyEnforcer());
       processorManager = processorManagerFn.apply(ImmutableList.of(segmentMapFn));
     } else {
       processorManager = new ChainedProcessorManager<>(ProcessorManagers.of(() -> segmentMapFnProcessor), processorManagerFn);
@@ -342,7 +342,7 @@ private FrameProcessor<Function<SegmentReference, SegmentReference>> makeSegment
     if (broadcastInputs.isEmpty()) {
       if (ExecutionVertex.of(query).isSegmentMapFunctionExpensive()) {
         // Joins may require significant computation to compute the segmentMapFn. Offload it to a processor.
-        return new SimpleSegmentMapFnProcessor(query);
+        return new SimpleSegmentMapFnProcessor(query, frameContext.policyEnforcer());
       } else {
         // Non-joins are expected to have cheap-to-compute segmentMapFn. Do the computation in the factory thread,
         // without offloading to a processor.
@@ -351,6 +351,7 @@ private FrameProcessor<Function<SegmentReference, SegmentReference>> makeSegment
     } else {
       return BroadcastJoinSegmentMapFnProcessor.create(
           query,
+          frameContext.policyEnforcer(),
           broadcastInputs,
           frameContext.memoryParameters().getBroadcastBufferMemory()
       );

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessor.java

@@ -40,6 +40,8 @@
 import org.apache.druid.query.InlineDataSource;
 import org.apache.druid.query.JoinAlgorithm;
 import org.apache.druid.query.Query;
+import org.apache.druid.query.planning.ExecutionVertex;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.Cursor;
 import org.apache.druid.segment.SegmentReference;
@@ -64,6 +66,7 @@
 public class BroadcastJoinSegmentMapFnProcessor implements FrameProcessor<Function<SegmentReference, SegmentReference>>
 {
   private final Query<?> query;
+  private final PolicyEnforcer policyEnforcer;
   private final Int2IntMap inputNumberToProcessorChannelMap;
   private final List<ReadableFrameChannel> channels;
   private final List<FrameReader> channelReaders;
@@ -87,13 +90,15 @@ public class BroadcastJoinSegmentMapFnProcessor implements FrameProcessor<Functi
    */
   public BroadcastJoinSegmentMapFnProcessor(
       final Query<?> query,
+      final PolicyEnforcer policyEnforcer,
       final Int2IntMap inputNumberToProcessorChannelMap,
       final List<ReadableFrameChannel> channels,
       final List<FrameReader> channelReaders,
       final long memoryReservedForBroadcastJoin
   )
   {
     this.query = query;
+    this.policyEnforcer = policyEnforcer;
     this.inputNumberToProcessorChannelMap = inputNumberToProcessorChannelMap;
     this.channels = channels;
     this.channelReaders = channelReaders;
@@ -117,6 +122,7 @@ public BroadcastJoinSegmentMapFnProcessor(
    */
   public static BroadcastJoinSegmentMapFnProcessor create(
       final Query<?> query,
+      final PolicyEnforcer policyEnforcer,
       final Int2ObjectMap<ReadableInput> sideChannels,
       final long memoryReservedForBroadcastJoin
   )
@@ -134,6 +140,7 @@ public static BroadcastJoinSegmentMapFnProcessor create(
 
     return new BroadcastJoinSegmentMapFnProcessor(
         query,
+        policyEnforcer,
         inputNumberToProcessorChannelMap,
         inputChannels,
         channelReaders,
@@ -193,7 +200,8 @@ private void addFrame(final int channelNumber, final Frame frame)
 
   private Function<SegmentReference, SegmentReference> createSegmentMapFunction()
   {
-    return inlineChannelData(query.getDataSource()).createSegmentMapFunction(query);
+    DataSource transformed = inlineChannelData(query.getDataSource());
+    return ExecutionVertex.of(query.withDataSource(transformed)).createSegmentMapFunction(policyEnforcer);
   }
 
   DataSource inlineChannelData(final DataSource originalDataSource)
@@ -230,7 +238,6 @@ DataSource inlineChannelData(final DataSource originalDataSource)
    * broadcast tables.
    *
    * @param readableInputs all readable input channel numbers, including non-side-channels
-   *
    * @return whether side channels have been fully read
    */
   boolean buildBroadcastTablesIncrementally(final IntSet readableInputs)

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/SimpleSegmentMapFnProcessor.java

@@ -26,6 +26,7 @@
 import org.apache.druid.frame.processor.ReturnOrAwait;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.planning.ExecutionVertex;
+import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.segment.SegmentReference;
 
 import java.util.Collections;
@@ -44,10 +45,13 @@
 public class SimpleSegmentMapFnProcessor implements FrameProcessor<Function<SegmentReference, SegmentReference>>
 {
   private final Query<?> query;
+  private final PolicyEnforcer policyEnforcer;
 
-  public SimpleSegmentMapFnProcessor(final Query<?> query)
+  public SimpleSegmentMapFnProcessor(final Query<?> query,
+                                     final PolicyEnforcer policyEnforcer)
   {
     this.query = query;
+    this.policyEnforcer = policyEnforcer;
   }
 
   @Override
@@ -66,7 +70,7 @@ public List<WritableFrameChannel> outputChannels()
   public ReturnOrAwait<Function<SegmentReference, SegmentReference>> runIncrementally(final IntSet readableInputs)
   {
     ExecutionVertex ev = ExecutionVertex.of(query);
-    return ReturnOrAwait.returnObject(ev.createSegmentMapFunction());
+    return ReturnOrAwait.returnObject(ev.createSegmentMapFunction(policyEnforcer));
   }
 
   @Override

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java

@@ -51,6 +51,7 @@
 import org.apache.druid.query.DefaultQueryConfig;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.policy.NoopPolicyEnforcer;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.QueryStackTests;
 import org.apache.druid.server.ResponseContextConfig;
@@ -215,6 +216,7 @@ public void register(ControllerHolder holder)
         CalciteTests.createJoinableFactoryWrapper(),
         CatalogResolver.NULL_RESOLVER,
         new AuthConfig(),
+        NoopPolicyEnforcer.instance(),
         new DruidHookDispatcher()
     );
 

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java

@@ -29,6 +29,7 @@
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.query.QuerySegmentWalker;
+import org.apache.druid.query.policy.NoopPolicyEnforcer;
 import org.apache.druid.server.security.AuthConfig;
 import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
 import org.apache.druid.sql.calcite.planner.CatalogResolver;
@@ -77,6 +78,7 @@ NamedViewSchema.NAME, new NamedViewSchema(EasyMock.createMock(ViewSchema.class))
         new CalciteRulesManager(ImmutableSet.of()),
         CalciteTests.TEST_AUTHORIZER_MAPPER,
         AuthConfig.newBuilder().build(),
+        NoopPolicyEnforcer.instance(),
         new DruidHookDispatcher()
     );
     final NativeSqlEngine engine = CalciteTests.createMockSqlEngine(