Specify system access control for testing presto server

presto-main-base/src/main/java/com/facebook/presto/testing/TestingAccessControlManager.java

@@ -101,9 +101,16 @@ public class TestingAccessControlManager
 
     @Inject
     public TestingAccessControlManager(TransactionManager transactionManager)
+    {
+        this(transactionManager, true);
+    }
+
+    public TestingAccessControlManager(TransactionManager transactionManager, boolean loadDefaultSystemAccessControl)
     {
         super(transactionManager);
-        setSystemAccessControl(AllowAllSystemAccessControl.NAME, ImmutableMap.of());
+        if (loadDefaultSystemAccessControl) {
+            setSystemAccessControl(AllowAllSystemAccessControl.NAME, ImmutableMap.of());
+        }
     }
 
     public static TestingPrivilege privilege(String entityName, TestingPrivilegeType type)

presto-main-base/src/main/java/com/facebook/presto/testing/TestingPrestoServerModule.java

@@ -0,0 +1,63 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.testing;
+
+import com.facebook.presto.eventlistener.EventListenerConfig;
+import com.facebook.presto.eventlistener.EventListenerManager;
+import com.facebook.presto.security.AccessControlManager;
+import com.facebook.presto.server.GracefulShutdownHandler;
+import com.facebook.presto.server.security.PrestoAuthenticatorManager;
+import com.facebook.presto.spi.security.AccessControl;
+import com.facebook.presto.storage.TempStorageManager;
+import com.facebook.presto.transaction.TransactionManager;
+import com.google.inject.Binder;
+import com.google.inject.Module;
+import com.google.inject.Provides;
+import com.google.inject.Scopes;
+import com.google.inject.Singleton;
+
+import static java.util.Objects.requireNonNull;
+
+public class TestingPrestoServerModule
+        implements Module
+{
+    private final boolean loadDefaultSystemAccessControl;
+
+    public TestingPrestoServerModule(boolean loadDefaultSystemAccessControl)
+    {
+        this.loadDefaultSystemAccessControl = loadDefaultSystemAccessControl;
+    }
+
+    @Override
+    public void configure(Binder binder)
+    {
+        binder.bind(PrestoAuthenticatorManager.class).in(Scopes.SINGLETON);
+        binder.bind(TestingEventListenerManager.class).in(Scopes.SINGLETON);
+        binder.bind(TestingTempStorageManager.class).in(Scopes.SINGLETON);
+        binder.bind(EventListenerManager.class).to(TestingEventListenerManager.class).in(Scopes.SINGLETON);
+        binder.bind(EventListenerConfig.class).in(Scopes.SINGLETON);
+        binder.bind(TempStorageManager.class).to(TestingTempStorageManager.class).in(Scopes.SINGLETON);
+        binder.bind(AccessControl.class).to(AccessControlManager.class).in(Scopes.SINGLETON);
+        binder.bind(GracefulShutdownHandler.class).in(Scopes.SINGLETON);
+        binder.bind(ProcedureTester.class).in(Scopes.SINGLETON);
+    }
+
+    @Provides
+    @Singleton
+    public AccessControlManager createAccessControlManager(TransactionManager transactionManager)
+    {
+        requireNonNull(transactionManager, "transactionManager is null");
+        return new TestingAccessControlManager(transactionManager, loadDefaultSystemAccessControl);
+    }
+}

presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java

@@ -39,7 +39,6 @@
 import com.facebook.presto.cost.StatsCalculator;
 import com.facebook.presto.dispatcher.DispatchManager;
 import com.facebook.presto.dispatcher.QueryPrerequisitesManagerModule;
-import com.facebook.presto.eventlistener.EventListenerConfig;
 import com.facebook.presto.eventlistener.EventListenerManager;
 import com.facebook.presto.execution.QueryInfo;
 import com.facebook.presto.execution.QueryManager;
@@ -74,7 +73,6 @@
 import com.facebook.presto.spi.eventlistener.EventListener;
 import com.facebook.presto.spi.function.SqlFunction;
 import com.facebook.presto.spi.memory.ClusterMemoryPoolManager;
-import com.facebook.presto.spi.security.AccessControl;
 import com.facebook.presto.split.PageSourceManager;
 import com.facebook.presto.split.SplitManager;
 import com.facebook.presto.sql.analyzer.FeaturesConfig;
@@ -85,11 +83,10 @@
 import com.facebook.presto.sql.planner.NodePartitioningManager;
 import com.facebook.presto.sql.planner.Plan;
 import com.facebook.presto.sql.planner.sanity.PlanCheckerProviderManager;
-import com.facebook.presto.storage.TempStorageManager;
 import com.facebook.presto.testing.ProcedureTester;
 import com.facebook.presto.testing.TestingAccessControlManager;
 import com.facebook.presto.testing.TestingEventListenerManager;
-import com.facebook.presto.testing.TestingTempStorageManager;
+import com.facebook.presto.testing.TestingPrestoServerModule;
 import com.facebook.presto.testing.TestingWarningCollectorModule;
 import com.facebook.presto.transaction.TransactionManager;
 import com.facebook.presto.ttl.clusterttlprovidermanagers.ClusterTtlProviderManagerModule;
@@ -289,6 +286,42 @@ public TestingPrestoServer(
             List<Module> additionalModules,
             Optional<Path> dataDirectory)
             throws Exception
+    {
+        this(
+                resourceManager,
+                resourceManagerEnabled,
+                catalogServer,
+                catalogServerEnabled,
+                coordinatorSidecar,
+                coordinatorSidecarEnabled,
+                coordinator,
+                skipLoadingResourceGroupConfigurationManager,
+                true,
+                properties,
+                environment,
+                discoveryUri,
+                parserOptions,
+                additionalModules,
+                dataDirectory);
+    }
+
+    public TestingPrestoServer(
+            boolean resourceManager,
+            boolean resourceManagerEnabled,
+            boolean catalogServer,
+            boolean catalogServerEnabled,
+            boolean coordinatorSidecar,
+            boolean coordinatorSidecarEnabled,
+            boolean coordinator,
+            boolean skipLoadingResourceGroupConfigurationManager,
+            boolean loadDefaultSystemAccessControl,
+            Map<String, String> properties,
+            String environment,
+            URI discoveryUri,
+            SqlParserOptions parserOptions,
+            List<Module> additionalModules,
+            Optional<Path> dataDirectory)
+            throws Exception
     {
         this.resourceManager = resourceManager;
         this.catalogServer = catalogServer;
@@ -328,18 +361,9 @@ public TestingPrestoServer(
                 .add(new NodeTtlFetcherManagerModule())
                 .add(new ClusterTtlProviderManagerModule())
                 .add(new ClientRequestFilterModule())
+                .add(new TestingPrestoServerModule(loadDefaultSystemAccessControl))
                 .add(binder -> {
-                    binder.bind(TestingAccessControlManager.class).in(Scopes.SINGLETON);
-                    binder.bind(TestingEventListenerManager.class).in(Scopes.SINGLETON);
-                    binder.bind(TestingTempStorageManager.class).in(Scopes.SINGLETON);
-                    binder.bind(AccessControlManager.class).to(TestingAccessControlManager.class).in(Scopes.SINGLETON);
-                    binder.bind(EventListenerManager.class).to(TestingEventListenerManager.class).in(Scopes.SINGLETON);
-                    binder.bind(EventListenerConfig.class).in(Scopes.SINGLETON);
-                    binder.bind(TempStorageManager.class).to(TestingTempStorageManager.class).in(Scopes.SINGLETON);
-                    binder.bind(AccessControl.class).to(AccessControlManager.class).in(Scopes.SINGLETON);
                     binder.bind(ShutdownAction.class).to(TestShutdownAction.class).in(Scopes.SINGLETON);
-                    binder.bind(GracefulShutdownHandler.class).in(Scopes.SINGLETON);
-                    binder.bind(ProcedureTester.class).in(Scopes.SINGLETON);
                     binder.bind(RequestBlocker.class).in(Scopes.SINGLETON);
                     newSetBinder(binder, Filter.class, TheServlet.class).addBinding()
                             .to(RequestBlocker.class).in(Scopes.SINGLETON);
@@ -385,7 +409,7 @@ public TestingPrestoServer(
         transactionManager = injector.getInstance(TransactionManager.class);
         sqlParser = injector.getInstance(SqlParser.class);
         metadata = injector.getInstance(Metadata.class);
-        accessControl = injector.getInstance(TestingAccessControlManager.class);
+        accessControl = (TestingAccessControlManager) injector.getInstance(AccessControlManager.class);
         procedureTester = injector.getInstance(ProcedureTester.class);
         splitManager = injector.getInstance(SplitManager.class);
         pageSourceManager = injector.getInstance(PageSourceManager.class);

presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java

@@ -31,6 +31,7 @@
 import com.facebook.presto.metadata.InternalNode;
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.security.AllowAllSystemAccessControl;
 import com.facebook.presto.server.BasicQueryInfo;
 import com.facebook.presto.server.testing.TestingPrestoServer;
 import com.facebook.presto.spi.ConnectorId;
@@ -134,6 +135,8 @@ public class DistributedQueryRunner
     private final int resourceManagerCount;
     private final AtomicReference<Handle> testFunctionNamespacesHandle = new AtomicReference<>();
 
+    private final Map<String, String> accessControlProperties;
+
     @Deprecated
     public DistributedQueryRunner(Session defaultSession, int nodeCount)
             throws Exception
@@ -163,7 +166,8 @@ public DistributedQueryRunner(Session defaultSession, int nodeCount, Map<String,
                 ENVIRONMENT,
                 Optional.empty(),
                 Optional.empty(),
-                ImmutableList.of());
+                ImmutableList.of(),
+                ImmutableMap.of("access-control.name", AllowAllSystemAccessControl.NAME));
     }
 
     public static Builder builder(Session defaultSession)
@@ -189,11 +193,13 @@ private DistributedQueryRunner(
             String environment,
             Optional<Path> dataDirectory,
             Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher,
-            List<Module> extraModules)
+            List<Module> extraModules,
+            Map<String, String> accessControlProperties)
             throws Exception
     {
         requireNonNull(defaultSession, "defaultSession is null");
         this.extraModules = requireNonNull(extraModules, "extraModules is null");
+        this.accessControlProperties = requireNonNull(accessControlProperties, "accessControlProperties is null");
 
         try {
             long start = nanoTime();
@@ -243,6 +249,7 @@ private DistributedQueryRunner(
                             coordinatorSidecarEnabled,
                             false,
                             skipLoadingResourceGroupConfigurationManager,
+                            false,
                             workerProperties,
                             parserOptions,
                             environment,
@@ -273,6 +280,7 @@ private DistributedQueryRunner(
                             false,
                             false,
                             skipLoadingResourceGroupConfigurationManager,
+                            false,
                             rmProperties,
                             parserOptions,
                             environment,
@@ -294,6 +302,7 @@ private DistributedQueryRunner(
                         false,
                         false,
                         skipLoadingResourceGroupConfigurationManager,
+                        false,
                         catalogServerProperties,
                         parserOptions,
                         environment,
@@ -313,6 +322,7 @@ private DistributedQueryRunner(
                         true,
                         false,
                         skipLoadingResourceGroupConfigurationManager,
+                        false,
                         coordinatorSidecarProperties,
                         parserOptions,
                         environment,
@@ -321,6 +331,8 @@ private DistributedQueryRunner(
                 servers.add(coordinatorSidecar.get());
             }
 
+            final boolean loadDefaultSystemAccessControl = !accessControlProperties.containsKey("access-control.name") ||
+                    accessControlProperties.get("access-control.name").equals("allow-all");
             for (int i = 0; i < coordinatorCount; i++) {
                 TestingPrestoServer coordinator = closer.register(createTestingPrestoServer(
                         discoveryUrl,
@@ -332,6 +344,7 @@ private DistributedQueryRunner(
                         false,
                         true,
                         skipLoadingResourceGroupConfigurationManager,
+                        loadDefaultSystemAccessControl,
                         extraCoordinatorProperties,
                         parserOptions,
                         environment,
@@ -464,6 +477,7 @@ private static TestingPrestoServer createTestingPrestoServer(
             boolean coordinatorSidecarEnabled,
             boolean coordinator,
             boolean skipLoadingResourceGroupConfigurationManager,
+            boolean loadDefaultSystemAccessControl,
             Map<String, String> extraProperties,
             SqlParserOptions parserOptions,
             String environment,
@@ -495,6 +509,7 @@ private static TestingPrestoServer createTestingPrestoServer(
                 coordinatorSidecarEnabled,
                 coordinator,
                 skipLoadingResourceGroupConfigurationManager,
+                loadDefaultSystemAccessControl,
                 properties,
                 environment,
                 discoveryUri,
@@ -798,6 +813,15 @@ public void createTestFunctionNamespace(String catalogName, String schemaName)
         testFunctionNamespacesHandle.get().execute("INSERT INTO function_namespaces SELECT ?, ?", catalogName, schemaName);
     }
 
+    public void loadSystemAccessControl()
+    {
+        for (TestingPrestoServer server : servers) {
+            if (server.isCoordinator()) {
+                server.getAccessControl().loadSystemAccessControl(accessControlProperties);
+            }
+        }
+    }
+
     private boolean isConnectorVisibleToAllNodes(ConnectorId connectorId)
     {
         if (!externalWorkers.isEmpty()) {
@@ -1097,6 +1121,7 @@ public static class Builder
         private boolean skipLoadingResourceGroupConfigurationManager;
         private List<Module> extraModules = ImmutableList.of();
         private int resourceManagerCount = 1;
+        private Map<String, String> accessControlProperties = ImmutableMap.of("access-control.name", AllowAllSystemAccessControl.NAME);
 
         protected Builder(Session defaultSession)
         {
@@ -1232,6 +1257,12 @@ public Builder setSkipLoadingResourceGroupConfigurationManager(boolean skipLoadi
             return this;
         }
 
+        public Builder setAccessControlProperties(Map<String, String> accessControlProperties)
+        {
+            this.accessControlProperties = accessControlProperties;
+            return this;
+        }
+
         public DistributedQueryRunner build()
                 throws Exception
         {
@@ -1253,7 +1284,8 @@ public DistributedQueryRunner build()
                     environment,
                     dataDirectory,
                     externalWorkerLauncher,
-                    extraModules);
+                    extraModules,
+                    accessControlProperties);
         }
     }
 }