Support exchange spooling on GCS

plugin/trino-delta-lake/pom.xml

@@ -205,6 +205,32 @@
             <groupId>io.trino</groupId>
             <artifactId>trino-exchange-filesystem</artifactId>
             <scope>test</scope>
+            <exclusions>
+                <exclusion>
+                    <groupId>com.google.apis</groupId>
+                    <artifactId>google-api-services-storage</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.api-client</groupId>
+                    <artifactId>google-api-client</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.auto.value</groupId>
+                    <artifactId>auto-value-annotations</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.code.gson</groupId>
+                    <artifactId>gson</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.http-client</groupId>
+                    <artifactId>google-http-client</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.http-client</groupId>
+                    <artifactId>google-http-client-jackson2</artifactId>
+                </exclusion>
+            </exclusions>
         </dependency>
 
         <dependency>

plugin/trino-exchange-filesystem/pom.xml

@@ -111,6 +111,92 @@
             <artifactId>azure-storage-blob-batch</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.google.api</groupId>
+            <artifactId>gax</artifactId>
+            <version>2.17.0</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>com.google.protobuf</groupId>
+                    <artifactId>protobuf-java</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>io.opencensus</groupId>
+                    <artifactId>opencensus-api</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>org.threeten</groupId>
+                    <artifactId>threetenbp</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.auth</groupId>
+            <artifactId>google-auth-library-credentials</artifactId>
+            <version>1.6.0</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.auth</groupId>
+            <artifactId>google-auth-library-oauth2-http</artifactId>
+            <version>1.6.0</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>commons-logging</groupId>
+                    <artifactId>commons-logging</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>org.apache.httpcomponents</groupId>
+                    <artifactId>httpcore</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-core</artifactId>
+            <version>2.5.6</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>com.google.protobuf</groupId>
+                    <artifactId>protobuf-java</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>commons-logging</groupId>
+                    <artifactId>commons-logging</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-storage</artifactId>
+            <version>2.5.1</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>com.google.auto.value</groupId>
+                    <artifactId>auto-value-annotations</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.guava</groupId>
+                    <artifactId>listenablefuture</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.oauth-client</groupId>
+                    <artifactId>google-oauth-client</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.google.protobuf</groupId>
+                    <artifactId>protobuf-java</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>org.checkerframework</groupId>
+                    <artifactId>checker-qual</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+
         <dependency>
             <groupId>com.google.code.findbugs</groupId>
             <artifactId>jsr305</artifactId>

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/FileSystemExchange.java

@@ -17,8 +17,6 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.Multimap;
-import com.google.common.util.concurrent.Futures;
-import com.google.common.util.concurrent.ListenableFuture;
 import io.trino.spi.exchange.Exchange;
 import io.trino.spi.exchange.ExchangeContext;
 import io.trino.spi.exchange.ExchangeSinkHandle;
@@ -307,11 +305,7 @@ public ExchangeSourceStatistics getExchangeSourceStatistics(ExchangeSourceHandle
     @Override
     public void close()
     {
-        ImmutableList.Builder<ListenableFuture<Void>> futures = ImmutableList.builder();
-        for (Integer taskPartitionId : allSinks) {
-            futures.add(exchangeStorage.deleteRecursively(getTaskOutputDirectory(taskPartitionId)));
-        }
-        stats.getCloseExchange().record(Futures.allAsList(futures.build()));
+        stats.getCloseExchange().record(exchangeStorage.deleteRecursively(allSinks.stream().map(this::getTaskOutputDirectory).collect(toImmutableList())));
     }
 
     private static String generateRandomizedPrefix()

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/FileSystemExchangeModule.java

@@ -29,6 +29,8 @@
 import java.util.List;
 
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static io.trino.plugin.exchange.filesystem.s3.S3FileSystemExchangeStorage.CompatibilityMode.AWS;
+import static io.trino.plugin.exchange.filesystem.s3.S3FileSystemExchangeStorage.CompatibilityMode.GCP;
 import static io.trino.spi.StandardErrorCode.CONFIGURATION_INVALID;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.lang.String.format;
@@ -53,11 +55,13 @@ protected void setup(Binder binder)
         if (scheme == null || scheme.equals("file")) {
             binder.bind(FileSystemExchangeStorage.class).to(LocalFileSystemExchangeStorage.class).in(Scopes.SINGLETON);
         }
-        else if (ImmutableSet.of("s3", "s3a", "s3n").contains(scheme)) {
+        else if (ImmutableSet.of("s3", "gs").contains(scheme)) {
             binder.bind(S3FileSystemExchangeStorageStats.class).in(Scopes.SINGLETON);
             newExporter(binder).export(S3FileSystemExchangeStorageStats.class).withGeneratedName();
             binder.bind(FileSystemExchangeStorage.class).to(S3FileSystemExchangeStorage.class).in(Scopes.SINGLETON);
             configBinder(binder).bindConfig(ExchangeS3Config.class);
+            S3FileSystemExchangeStorage.CompatibilityMode compatibilityMode = scheme.equals("gs") ? GCP : AWS;
+            binder.bind(S3FileSystemExchangeStorage.CompatibilityMode.class).toInstance(compatibilityMode);
         }
         else if (ImmutableSet.of("abfs", "abfss").contains(scheme)) {
             binder.bind(FileSystemExchangeStorage.class).to(AzureBlobFileSystemExchangeStorage.class).in(Scopes.SINGLETON);

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/FileSystemExchangeSink.java

@@ -13,6 +13,7 @@
  */
 package io.trino.plugin.exchange.filesystem;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -203,7 +204,7 @@ public synchronized CompletableFuture<Void> abort()
 
         return stats.getExchangeSinkAbort().record(toCompletableFuture(Futures.transformAsync(
                 abortFuture,
-                ignored -> exchangeStorage.deleteRecursively(outputDirectory),
+                ignored -> exchangeStorage.deleteRecursively(ImmutableList.of(outputDirectory)),
                 directExecutor())));
     }
 

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/FileSystemExchangeStorage.java

@@ -36,7 +36,7 @@ public interface FileSystemExchangeStorage
 
     ListenableFuture<Void> createEmptyFile(URI file);
 
-    ListenableFuture<Void> deleteRecursively(URI dir);
+    ListenableFuture<Void> deleteRecursively(List<URI> directories);
 
     List<FileStatus> listFiles(URI dir) throws IOException;
 

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/azure/AzureBlobFileSystemExchangeStorage.java

@@ -158,12 +158,14 @@ public ListenableFuture<Void> createEmptyFile(URI file)
     }
 
     @Override
-    public ListenableFuture<Void> deleteRecursively(URI dir)
+    public ListenableFuture<Void> deleteRecursively(List<URI> directories)
     {
-        return asVoid(Futures.transformAsync(
-                toListenableFuture(listObjectsRecursively(dir).byPage().collectList().toFuture()),
-                pagedResponseList -> deleteObjects(getContainerName(dir), pagedResponseList),
-                directExecutor()));
+        return asVoid(Futures.allAsList(directories.stream()
+                .map(dir -> Futures.transformAsync(
+                        toListenableFuture(listObjectsRecursively(dir).byPage().collectList().toFuture()),
+                        pagedResponseList -> deleteObjects(getContainerName(dir), pagedResponseList),
+                        directExecutor()))
+                .collect(toImmutableList())));
     }
 
     @Override

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/local/LocalFileSystemExchangeStorage.java

@@ -107,13 +107,15 @@ public ListenableFuture<Void> createEmptyFile(URI file)
     }
 
     @Override
-    public ListenableFuture<Void> deleteRecursively(URI dir)
+    public ListenableFuture<Void> deleteRecursively(List<URI> directories)
     {
-        try {
-            MoreFiles.deleteRecursively(Paths.get(dir.getPath()), ALLOW_INSECURE);
-        }
-        catch (IOException | RuntimeException e) {
-            return immediateFailedFuture(e);
+        for (URI dir : directories) {
+            try {
+                MoreFiles.deleteRecursively(Paths.get(dir.getPath()), ALLOW_INSECURE);
+            }
+            catch (IOException | RuntimeException e) {
+                return immediateFailedFuture(e);
+            }
         }
         return immediateVoidFuture();
     }

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java

@@ -48,6 +48,7 @@ public class ExchangeS3Config
     private int asyncClientConcurrency = 100;
     private int asyncClientMaxPendingConnectionAcquires = 10000;
     private Duration connectionAcquisitionTimeout = new Duration(1, MINUTES);
+    private Optional<String> gcsJsonKeyFilePath = Optional.empty();
 
     public String getS3AwsAccessKey()
     {
@@ -205,4 +206,16 @@ public ExchangeS3Config setConnectionAcquisitionTimeout(Duration connectionAcqui
         this.connectionAcquisitionTimeout = connectionAcquisitionTimeout;
         return this;
     }
+
+    public Optional<String> getGcsJsonKeyFilePath()
+    {
+        return gcsJsonKeyFilePath;
+    }
+
+    @Config("exchange.gcs.json-key-file-path")
+    public ExchangeS3Config setGcsJsonKeyFilePath(String gcsJsonKeyFilePath)
+    {
+        this.gcsJsonKeyFilePath = Optional.ofNullable(gcsJsonKeyFilePath);
+        return this;
+    }
 }