Add faker FTE query runner

wendigo · Feb 3, 2025 · 4a5e5a1 · 4a5e5a1
1 parent b35a1ba
commit 4a5e5a1
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/plugin/trino-faker/pom.xml b/plugin/trino-faker/pom.xml
@@ -142,6 +142,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-exchange-filesystem</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-main</artifactId>

diff --git a/plugin/trino-faker/src/test/java/io/trino/plugin/faker/FakerQueryRunner.java b/plugin/trino-faker/src/test/java/io/trino/plugin/faker/FakerQueryRunner.java
@@ -17,14 +17,17 @@
 import io.airlift.log.Level;
 import io.airlift.log.Logger;
 import io.airlift.log.Logging;
+import io.trino.plugin.exchange.filesystem.FileSystemExchangePlugin;
 import io.trino.plugin.tpch.TpchPlugin;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.QueryRunner;
 
+import java.io.File;
 import java.util.Map;
 
 import static io.airlift.testing.Closeables.closeAllSuppress;
 import static io.trino.testing.TestingSession.testSessionBuilder;
+import static java.nio.file.Files.createTempDirectory;
 import static java.util.Objects.requireNonNullElse;
 
 public class FakerQueryRunner
@@ -95,4 +98,37 @@ public static void main(String[] args)
         log.info("======== SERVER STARTED ========");
         log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl());
     }
+
+    public static final class FakerQueryRunnerWithTaskRetries
+    {
+        private FakerQueryRunnerWithTaskRetries() {}
+
+        public static void main(String[] args)
+                throws Exception
+        {
+            Logger log = Logger.get(FakerQueryRunnerWithTaskRetries.class);
+
+            File exchangeManagerDirectory = createTempDirectory("exchange_manager").toFile();
+            Map<String, String> exchangeManagerProperties = ImmutableMap.<String, String>builder()
+                    .put("exchange.base-directories", exchangeManagerDirectory.getAbsolutePath())
+                    .buildOrThrow();
+            exchangeManagerDirectory.deleteOnExit();
+
+            @SuppressWarnings("resource")
+            QueryRunner queryRunner = builder()
+                    .setExtraProperties(ImmutableMap.<String, String>builder()
+                            .put("http-server.http.port", requireNonNullElse(System.getenv("TRINO_PORT"), "8080"))
+                            .put("retry-policy", "TASK")
+                            .put("fault-tolerant-execution-task-memory", "1GB")
+                            .buildOrThrow())
+                    .setAdditionalSetup(runner -> {
+                        runner.installPlugin(new FileSystemExchangePlugin());
+                        runner.loadExchangeManager("filesystem", exchangeManagerProperties);
+                    })
+                    .build();
+
+            log.info("======== SERVER STARTED ========");
+            log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl());
+        }
+    }
 }