Code review: use StandaloneQueryRunner

trinodb · Mar 6, 2024 · f896e2b · f896e2b
1 parent dc8a1dc
commit f896e2b
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 21 deletions.
diff --git a/...gin/opa/DistributedQueryRunnerHelper.java → ...o/trino/plugin/opa/QueryRunnerHelper.java b/...gin/opa/DistributedQueryRunnerHelper.java → ...o/trino/plugin/opa/QueryRunnerHelper.java
@@ -15,32 +15,32 @@
 
 import io.trino.Session;
 import io.trino.spi.security.Identity;
-import io.trino.testing.DistributedQueryRunner;
+import io.trino.testing.QueryRunner;
+import io.trino.testing.StandaloneQueryRunner;
 import org.intellij.lang.annotations.Language;
 
 import java.util.Set;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.trino.plugin.opa.TestHelpers.opaConfigToDict;
+import static io.trino.testing.TestingSession.testSession;
 import static io.trino.testing.TestingSession.testSessionBuilder;
 
-public final class DistributedQueryRunnerHelper
+public final class QueryRunnerHelper
 {
-    private final DistributedQueryRunner runner;
+    private final QueryRunner runner;
 
-    private DistributedQueryRunnerHelper(DistributedQueryRunner runner)
+    private QueryRunnerHelper(QueryRunner runner)
     {
         this.runner = runner;
     }
 
-    public static DistributedQueryRunnerHelper withOpaConfig(OpaConfig opaConfig)
-            throws Exception
+    public static QueryRunnerHelper withOpaConfig(OpaConfig opaConfig)
     {
-        return new DistributedQueryRunnerHelper(
-                DistributedQueryRunner.builder(testSessionBuilder().build())
-                        .setSystemAccessControl(new OpaAccessControlFactory().create(opaConfigToDict(opaConfig)))
-                        .setNodeCount(1)
-                        .build());
+        return new QueryRunnerHelper(
+                new StandaloneQueryRunner(
+                        testSession(),
+                        builder -> builder.setSystemAccessControl(new OpaAccessControlFactory().create(opaConfigToDict(opaConfig)))));
     }
 
     public Set<String> querySetOfStrings(String user, String query)
@@ -53,7 +53,7 @@ public Set<String> querySetOfStrings(Session session, @Language("SQL") String qu
         return runner.execute(session, query).getMaterializedRows().stream().map(row -> row.getField(0) == null ? "<NULL>" : row.getField(0).toString()).collect(toImmutableSet());
     }
 
-    public DistributedQueryRunner getBaseQueryRunner()
+    public QueryRunner getBaseQueryRunner()
     {
         return runner;
     }
@@ -65,6 +65,6 @@ public void teardown()
 
     private static Session userSession(String user)
     {
-        return testSessionBuilder().setIdentity(Identity.ofUser(user)).build();
+        return testSessionBuilder().setOriginalIdentity(Identity.ofUser(user)).setIdentity(Identity.ofUser(user)).build();
     }
 }
diff --git a/.../trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControlDataFilteringSystem.java b/.../trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControlDataFilteringSystem.java
@@ -98,7 +98,7 @@ public class TestOpaAccessControlDataFilteringSystem
             .addAll(DUMMY_CUSTOMERS_IN_TABLE)
             .build();
 
-    private DistributedQueryRunnerHelper runner;
+    private QueryRunnerHelper runner;
 
     @AfterEach
     public void teardown()
@@ -249,9 +249,8 @@ private void assertResultsForUser(String asUser, String query, Set<String> expec
     }
 
     private void setupTrinoWithOpa(OpaConfig opaConfig)
-            throws Exception
     {
-        this.runner = DistributedQueryRunnerHelper.withOpaConfig(opaConfig);
+        this.runner = QueryRunnerHelper.withOpaConfig(opaConfig);
         MockConnectorFactory connectorFactory = MockConnectorFactory.builder()
                 .withListSchemaNames(session -> ImmutableList.of("sample_schema"))
                 .withListTables((session, schema) -> ImmutableList.<String>builder()

diff --git a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControlSystem.java b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControlSystem.java
@@ -34,7 +34,7 @@
 @TestInstance(PER_CLASS)
 public class TestOpaAccessControlSystem
 {
-    private DistributedQueryRunnerHelper runner;
+    private QueryRunnerHelper runner;
     private static final String OPA_ALLOW_POLICY_NAME = "allow";
     private static final String OPA_BATCH_ALLOW_POLICY_NAME = "batchAllow";
     @Container
@@ -47,7 +47,6 @@ class UnbatchedAuthorizerTests
     {
         @BeforeAll
         public void setupTrino()
-                throws Exception
         {
             setupTrinoWithOpa(new OpaConfig().setOpaUri(OPA_CONTAINER.getOpaUriForPolicyPath(OPA_ALLOW_POLICY_NAME)));
         }
@@ -122,7 +121,6 @@ class BatchedAuthorizerTests
     {
         @BeforeAll
         public void setupTrino()
-                throws Exception
         {
             setupTrinoWithOpa(new OpaConfig()
                     .setOpaUri(OPA_CONTAINER.getOpaUriForPolicyPath(OPA_ALLOW_POLICY_NAME))
@@ -214,9 +212,8 @@ public void testAllowUnbatchedQuery()
     }
 
     private void setupTrinoWithOpa(OpaConfig opaConfig)
-            throws Exception
     {
-        this.runner = DistributedQueryRunnerHelper.withOpaConfig(opaConfig);
+        this.runner = QueryRunnerHelper.withOpaConfig(opaConfig);
         runner.getBaseQueryRunner().installPlugin(new BlackHolePlugin());
         runner.getBaseQueryRunner().createCatalog("catalog_one", "blackhole");
         runner.getBaseQueryRunner().createCatalog("catalog_two", "blackhole");