From 7c91e477bdf60f555f6fbf399c019bedf6543a7f Mon Sep 17 00:00:00 2001
From: "Mateusz \"Serafin\" Gajewski" <github@wendigo.pl>
Date: Tue, 24 Sep 2024 21:13:52 +0200
Subject: [PATCH] Use localstack for spooled tests

In comparison to Minio it supports SSE-C without the need to configure
TLS to secure connection.
---
 testing/trino-tests/pom.xml                   | 41 +++++++++++++++++
 ...actSpooledQueryDataDistributedQueries.java | 44 ++++++++++++-------
 2 files changed, 70 insertions(+), 15 deletions(-)

diff --git a/testing/trino-tests/pom.xml b/testing/trino-tests/pom.xml
index 866ad75b9cb33..679c1b8a40d72 100644
--- a/testing/trino-tests/pom.xml
+++ b/testing/trino-tests/pom.xml
@@ -305,12 +305,53 @@
             <artifactId>jmh-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>localstack</artifactId>
+            <scope>test</scope>
+        </dependency>
 
         <dependency>
             <groupId>org.yaml</groupId>
             <artifactId>snakeyaml</artifactId>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>auth</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>aws-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>regions</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>s3</artifactId>
+            <scope>test</scope>
+            <exclusions>
+                <exclusion>
+                    <groupId>commons-logging</groupId>
+                    <artifactId>commons-logging</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>sdk-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>
diff --git a/testing/trino-tests/src/test/java/io/trino/server/protocol/AbstractSpooledQueryDataDistributedQueries.java b/testing/trino-tests/src/test/java/io/trino/server/protocol/AbstractSpooledQueryDataDistributedQueries.java
index e9995f520ea98..07b8426036418 100644
--- a/testing/trino-tests/src/test/java/io/trino/server/protocol/AbstractSpooledQueryDataDistributedQueries.java
+++ b/testing/trino-tests/src/test/java/io/trino/server/protocol/AbstractSpooledQueryDataDistributedQueries.java
@@ -27,9 +27,15 @@
 import io.trino.testing.QueryRunner;
 import io.trino.testing.TestingStatementClientFactory;
 import io.trino.testing.TestingTrinoClient;
-import io.trino.testing.containers.Minio;
 import io.trino.tpch.TpchTable;
 import okhttp3.OkHttpClient;
+import org.testcontainers.containers.localstack.LocalStackContainer;
+import org.testcontainers.containers.localstack.LocalStackContainer.Service;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
 
 import java.util.Map;
 import java.util.Optional;
@@ -37,16 +43,13 @@
 
 import static io.airlift.testing.Closeables.closeAllSuppress;
 import static io.trino.client.StatementClientFactory.newStatementClient;
-import static io.trino.testing.containers.Minio.MINIO_ACCESS_KEY;
-import static io.trino.testing.containers.Minio.MINIO_REGION;
-import static io.trino.testing.containers.Minio.MINIO_SECRET_KEY;
 import static io.trino.util.Ciphers.createRandomAesEncryptionKey;
 import static java.util.Base64.getEncoder;
 
 public abstract class AbstractSpooledQueryDataDistributedQueries
         extends AbstractTestEngineOnlyQueries
 {
-    private Minio minio;
+    private LocalStackContainer localstack;
 
     protected abstract String encoding();
 
@@ -59,11 +62,14 @@ protected Map<String, String> spoolingConfig()
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        minio = closeAfterClass(Minio.builder().build());
-        minio.start();
+        localstack = closeAfterClass(new LocalStackContainer("s3-latest"));
+        localstack.start();
 
         String bucketName = "segments" + UUID.randomUUID();
-        minio.createBucket(bucketName, true);
+
+        try (S3Client client = createS3Client(localstack)) {
+            client.createBucket(CreateBucketRequest.builder().bucket(bucketName).build());
+        }
 
         DistributedQueryRunner queryRunner = MemoryQueryRunner.builder()
                 .setInitialTables(TpchTable.getTables())
@@ -75,13 +81,11 @@ protected QueryRunner createQueryRunner()
                     Map<String, String> spoolingConfig = ImmutableMap.<String, String>builder()
                             .put("fs.s3.enabled", "true")
                             .put("fs.location", "s3://" + bucketName + "/")
-                            // Direct storage access with encryption requires SSE-c which is not yet implemented
-                            .put("fs.segment.encryption", "false")
-                            .put("s3.endpoint", minio.getMinioAddress())
-                            .put("s3.region", MINIO_REGION)
-                            .put("s3.aws-access-key", MINIO_ACCESS_KEY)
-                            .put("s3.aws-secret-key", MINIO_SECRET_KEY)
-                            .put("s3.path-style-access", "true")
+                            .put("fs.segment.encryption", "true")
+                            .put("s3.endpoint", localstack.getEndpointOverride(Service.S3).toString())
+                            .put("s3.region", localstack.getRegion())
+                            .put("s3.aws-access-key", localstack.getAccessKey())
+                            .put("s3.aws-secret-key", localstack.getSecretKey())
                             .putAll(spoolingConfig())
                             .buildKeepingLast();
                     runner.loadSpoolingManager("filesystem", spoolingConfig);
@@ -119,4 +123,14 @@ private static String randomAES256Key()
     {
         return getEncoder().encodeToString(createRandomAesEncryptionKey().getEncoded());
     }
+
+    protected S3Client createS3Client(LocalStackContainer localstack)
+    {
+        return S3Client.builder()
+                .endpointOverride(localstack.getEndpointOverride(Service.S3))
+                .region(Region.of(localstack.getRegion()))
+                .credentialsProvider(StaticCredentialsProvider.create(
+                        AwsBasicCredentials.create(localstack.getAccessKey(), localstack.getSecretKey())))
+                .build();
+    }
 }