Allow GCS-based spooling to pass JSON key as a config

trinodb · Aug 10, 2022 · 9703bf6 · 9703bf6
1 parent a631276
commit 9703bf6
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 5 deletions.
diff --git a/docs/src/main/sphinx/admin/fault-tolerant-execution.rst b/docs/src/main/sphinx/admin/fault-tolerant-execution.rst
@@ -418,7 +418,13 @@ the property may be configured for:
      - Any S3-compatible storage
    * - ``exchange.gcs.json-key-file-path``
      - Path to the JSON file that contains your Google Cloud Platform
-       service account key.
+       service account key. Not to be set together with
+       ``exchange.gcs.json-key``
+     -
+     - GCS
+   * - ``exchange.gcs.json-key``
+     - Your Google Cloud Platform service account key in JSON format.
+       Not to be set together with ``exchange.gcs.json-key-file-path``
      -
      - GCS
    * - ``exchange.azure.connection-string``

diff --git a/...nge-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java b/...nge-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java
@@ -50,6 +50,7 @@ public class ExchangeS3Config
     private int asyncClientMaxPendingConnectionAcquires = 10000;
     private Duration connectionAcquisitionTimeout = new Duration(1, MINUTES);
     private Optional<String> gcsJsonKeyFilePath = Optional.empty();
+    private Optional<String> gcsJsonKey = Optional.empty();
 
     public String getS3AwsAccessKey()
     {
@@ -228,9 +229,24 @@ public Optional<String> getGcsJsonKeyFilePath()
     }
 
     @Config("exchange.gcs.json-key-file-path")
+    @ConfigDescription("Path to the JSON file that contains your Google Cloud Platform service account key. Not to be set together with `exchange.gcs.json-key`")
     public ExchangeS3Config setGcsJsonKeyFilePath(String gcsJsonKeyFilePath)
     {
         this.gcsJsonKeyFilePath = Optional.ofNullable(gcsJsonKeyFilePath);
         return this;
     }
+
+    public Optional<String> getGcsJsonKey()
+    {
+        return gcsJsonKey;
+    }
+
+    @Config("exchange.gcs.json-key")
+    @ConfigDescription("Your Google Cloud Platform service account key in JSON format. Not to be set together with `exchange.gcs.json-key-file-path`")
+    @ConfigSecuritySensitive
+    public ExchangeS3Config setGcsJsonKey(String gcsJsonKey)
+    {
+        this.gcsJsonKey = Optional.ofNullable(gcsJsonKey);
+        return this;
+    }
 }
diff --git a/...tem/src/main/java/io/trino/plugin/exchange/filesystem/s3/S3FileSystemExchangeStorage.java b/...tem/src/main/java/io/trino/plugin/exchange/filesystem/s3/S3FileSystemExchangeStorage.java
@@ -84,10 +84,12 @@
 import javax.crypto.SecretKey;
 import javax.inject.Inject;
 
+import java.io.ByteArrayInputStream;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
@@ -169,8 +171,16 @@ public S3FileSystemExchangeStorage(S3FileSystemExchangeStorageStats stats, Excha
                 config.getConnectionAcquisitionTimeout());
 
         if (compatibilityMode == GCP) {
-            if (config.getGcsJsonKeyFilePath().isPresent()) {
-                Credentials credentials = GoogleCredentials.fromStream(new FileInputStream(config.getGcsJsonKeyFilePath().get()));
+            Optional<String> gcsJsonKeyFilePath = config.getGcsJsonKeyFilePath();
+            Optional<String> gcsJsonKey = config.getGcsJsonKey();
+            verify(!(gcsJsonKeyFilePath.isPresent() && gcsJsonKey.isPresent()),
+                    "gcsJsonKeyFilePath and gcsJsonKey shouldn't be set at the same time");
+            if (gcsJsonKeyFilePath.isPresent()) {
+                Credentials credentials = GoogleCredentials.fromStream(new FileInputStream(gcsJsonKeyFilePath.get()));
+                this.gcsClient = Optional.of(StorageOptions.newBuilder().setCredentials(credentials).build().getService());
+            }
+            else if (gcsJsonKey.isPresent()) {
+                Credentials credentials = GoogleCredentials.fromStream(new ByteArrayInputStream(gcsJsonKey.get().getBytes(StandardCharsets.UTF_8)));
                 this.gcsClient = Optional.of(StorageOptions.newBuilder().setCredentials(credentials).build().getService());
             }
             else {

diff --git a/...filesystem/src/test/java/io/trino/plugin/exchange/filesystem/s3/TestExchangeS3Config.java b/...filesystem/src/test/java/io/trino/plugin/exchange/filesystem/s3/TestExchangeS3Config.java
@@ -47,7 +47,8 @@ public void testDefaults()
                 .setAsyncClientConcurrency(100)
                 .setAsyncClientMaxPendingConnectionAcquires(10000)
                 .setConnectionAcquisitionTimeout(new Duration(1, MINUTES))
-                .setGcsJsonKeyFilePath(null));
+                .setGcsJsonKeyFilePath(null)
+                .setGcsJsonKey(null));
     }
 
     @Test
@@ -68,6 +69,7 @@ public void testExplicitPropertyMappings()
                 .put("exchange.s3.async-client-max-pending-connection-acquires", "999")
                 .put("exchange.s3.async-client-connection-acquisition-timeout", "5m")
                 .put("exchange.gcs.json-key-file-path", "/path/to/gcs_keyfile.json")
+                .put("exchange.gcs.json-key", "{}")
                 .buildOrThrow();
 
         ExchangeS3Config expected = new ExchangeS3Config()
@@ -84,7 +86,8 @@ public void testExplicitPropertyMappings()
                 .setAsyncClientConcurrency(202)
                 .setAsyncClientMaxPendingConnectionAcquires(999)
                 .setConnectionAcquisitionTimeout(new Duration(5, MINUTES))
-                .setGcsJsonKeyFilePath("/path/to/gcs_keyfile.json");
+                .setGcsJsonKeyFilePath("/path/to/gcs_keyfile.json")
+                .setGcsJsonKey("{}");
 
         assertFullMapping(properties, expected);
     }