Skip to content

Commit 9ffaa0f

Browse files
kavetikaveti
committed
Support GCS vended credentials from Iceberg REST catalog
1 parent 52ab80c commit 9ffaa0f

File tree

7 files changed

+566
-9
lines changed

7 files changed

+566
-9
lines changed

‎lib/trino-filesystem-gcs/src/main/java/io/trino/filesystem/gcs/GcsFileSystem.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -76,23 +76,32 @@ public class GcsFileSystem
7676
private final long writeBlockSizeBytes;
7777
private final int pageSize;
7878
private final int batchSize;
79+
private final Optional<EncryptionKey> defaultEncryptionKey;
80+
private final Optional<EncryptionKey> defaultDecryptionKey;
7981

8082
public GcsFileSystem(ListeningExecutorService executorService, Storage storage, int readBlockSizeBytes, long writeBlockSizeBytes, int pageSize, int batchSize)
83+
{
84+
this(executorService, storage, readBlockSizeBytes, writeBlockSizeBytes, pageSize, batchSize, Optional.empty(), Optional.empty());
85+
}
86+
87+
public GcsFileSystem(ListeningExecutorService executorService, Storage storage, int readBlockSizeBytes, long writeBlockSizeBytes, int pageSize, int batchSize, Optional<EncryptionKey> defaultEncryptionKey, Optional<EncryptionKey> defaultDecryptionKey)
8188
{
8289
this.executorService = requireNonNull(executorService, "executorService is null");
8390
this.storage = requireNonNull(storage, "storage is null");
8491
this.readBlockSizeBytes = readBlockSizeBytes;
8592
this.writeBlockSizeBytes = writeBlockSizeBytes;
8693
this.pageSize = pageSize;
8794
this.batchSize = batchSize;
95+
this.defaultEncryptionKey = requireNonNull(defaultEncryptionKey, "defaultEncryptionKey is null");
96+
this.defaultDecryptionKey = requireNonNull(defaultDecryptionKey, "defaultDecryptionKey is null");
8897
}
8998

9099
@Override
91100
public TrinoInputFile newInputFile(Location location)
92101
{
93102
GcsLocation gcsLocation = new GcsLocation(location);
94103
checkIsValidFile(gcsLocation);
95-
return new GcsInputFile(gcsLocation, storage, readBlockSizeBytes, OptionalLong.empty(), Optional.empty(), Optional.empty());
104+
return new GcsInputFile(gcsLocation, storage, readBlockSizeBytes, OptionalLong.empty(), Optional.empty(), defaultDecryptionKey);
96105
}
97106

98107
@Override
@@ -108,7 +117,7 @@ public TrinoInputFile newInputFile(Location location, long length)
108117
{
109118
GcsLocation gcsLocation = new GcsLocation(location);
110119
checkIsValidFile(gcsLocation);
111-
return new GcsInputFile(gcsLocation, storage, readBlockSizeBytes, OptionalLong.of(length), Optional.empty(), Optional.empty());
120+
return new GcsInputFile(gcsLocation, storage, readBlockSizeBytes, OptionalLong.of(length), Optional.empty(), defaultDecryptionKey);
112121
}
113122

114123
@Override
@@ -124,7 +133,7 @@ public TrinoInputFile newInputFile(Location location, long length, Instant lastM
124133
{
125134
GcsLocation gcsLocation = new GcsLocation(location);
126135
checkIsValidFile(gcsLocation);
127-
return new GcsInputFile(gcsLocation, storage, readBlockSizeBytes, OptionalLong.of(length), Optional.of(lastModified), Optional.empty());
136+
return new GcsInputFile(gcsLocation, storage, readBlockSizeBytes, OptionalLong.of(length), Optional.of(lastModified), defaultDecryptionKey);
128137
}
129138

130139
@Override
@@ -140,7 +149,7 @@ public TrinoOutputFile newOutputFile(Location location)
140149
{
141150
GcsLocation gcsLocation = new GcsLocation(location);
142151
checkIsValidFile(gcsLocation);
143-
return new GcsOutputFile(gcsLocation, storage, writeBlockSizeBytes, Optional.empty());
152+
return new GcsOutputFile(gcsLocation, storage, writeBlockSizeBytes, defaultEncryptionKey);
144153
}
145154

146155
@Override

‎lib/trino-filesystem-gcs/src/main/java/io/trino/filesystem/gcs/GcsFileSystemFactory.java‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,17 @@
1717
import com.google.inject.Inject;
1818
import io.trino.filesystem.TrinoFileSystem;
1919
import io.trino.filesystem.TrinoFileSystemFactory;
20+
import io.trino.filesystem.encryption.EncryptionKey;
2021
import io.trino.spi.security.ConnectorIdentity;
2122
import jakarta.annotation.PreDestroy;
2223

24+
import java.util.Base64;
25+
import java.util.Optional;
26+
2327
import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator;
2428
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
29+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_DECRYPTION_KEY_PROPERTY;
30+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_ENCRYPTION_KEY_PROPERTY;
2531
import static java.lang.Math.toIntExact;
2632
import static java.util.Objects.requireNonNull;
2733
import static java.util.concurrent.Executors.newCachedThreadPool;
@@ -56,6 +62,18 @@ public void stop()
5662
@Override
5763
public TrinoFileSystem create(ConnectorIdentity identity)
5864
{
59-
return new GcsFileSystem(executorService, storageFactory.create(identity), readBlockSizeBytes, writeBlockSizeBytes, pageSize, batchSize);
65+
Optional<EncryptionKey> defaultEncryptionKey = extractEncryptionKey(identity, EXTRA_CREDENTIALS_GCS_ENCRYPTION_KEY_PROPERTY);
66+
Optional<EncryptionKey> defaultDecryptionKey = extractEncryptionKey(identity, EXTRA_CREDENTIALS_GCS_DECRYPTION_KEY_PROPERTY);
67+
return new GcsFileSystem(executorService, storageFactory.create(identity), readBlockSizeBytes, writeBlockSizeBytes, pageSize, batchSize, defaultEncryptionKey, defaultDecryptionKey);
68+
}
69+
70+
private static Optional<EncryptionKey> extractEncryptionKey(ConnectorIdentity identity, String extraCredentialKey)
71+
{
72+
String base64Key = identity.getExtraCredentials().get(extraCredentialKey);
73+
if (base64Key == null) {
74+
return Optional.empty();
75+
}
76+
byte[] keyBytes = Base64.getDecoder().decode(base64Key);
77+
return Optional.of(new EncryptionKey(keyBytes, "AES256"));
6078
}
6179
}

‎lib/trino-filesystem-gcs/src/main/java/io/trino/filesystem/gcs/GcsStorageFactory.java‎

Lines changed: 45 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@
1414
package io.trino.filesystem.gcs;
1515

1616
import com.google.api.gax.retrying.RetrySettings;
17+
import com.google.auth.oauth2.AccessToken;
18+
import com.google.auth.oauth2.GoogleCredentials;
19+
import com.google.cloud.NoCredentials;
1720
import com.google.cloud.storage.Storage;
1821
import com.google.cloud.storage.StorageOptions;
1922
import com.google.inject.Inject;
@@ -22,11 +25,18 @@
2225
import java.io.IOException;
2326
import java.io.UncheckedIOException;
2427
import java.time.Duration;
28+
import java.util.Date;
2529
import java.util.Map;
2630
import java.util.Optional;
2731

2832
import static com.google.cloud.storage.StorageRetryStrategy.getUniformStorageRetryStrategy;
2933
import static com.google.common.net.HttpHeaders.USER_AGENT;
34+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_NO_AUTH_PROPERTY;
35+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_EXPIRES_AT_PROPERTY;
36+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY;
37+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_PROJECT_ID_PROPERTY;
38+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_SERVICE_HOST_PROPERTY;
39+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_USER_PROJECT_PROPERTY;
3040
import static java.util.Objects.requireNonNull;
3141

3242
public class GcsStorageFactory
@@ -59,14 +69,45 @@ public GcsStorageFactory(GcsFileSystemConfig config, GcsAuth gcsAuth)
5969
public Storage create(ConnectorIdentity identity)
6070
{
6171
try {
72+
Map<String, String> extraCredentials = identity.getExtraCredentials();
73+
boolean noAuth = Boolean.parseBoolean(extraCredentials.getOrDefault(EXTRA_CREDENTIALS_GCS_NO_AUTH_PROPERTY, "false"));
74+
String vendedOAuthToken = extraCredentials.get(EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY);
75+
6276
StorageOptions.Builder storageOptionsBuilder = StorageOptions.newBuilder();
63-
if (projectId != null) {
64-
storageOptionsBuilder.setProjectId(projectId);
77+
78+
String effectiveProjectId = extraCredentials.getOrDefault(EXTRA_CREDENTIALS_GCS_PROJECT_ID_PROPERTY, projectId);
79+
if (effectiveProjectId != null) {
80+
storageOptionsBuilder.setProjectId(effectiveProjectId);
6581
}
6682

67-
gcsAuth.setAuth(storageOptionsBuilder, identity);
83+
if (noAuth) {
84+
storageOptionsBuilder.setCredentials(NoCredentials.getInstance());
85+
}
86+
else if (vendedOAuthToken != null) {
87+
Date expirationTime = null;
88+
String expiresAt = extraCredentials.get(EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_EXPIRES_AT_PROPERTY);
89+
if (expiresAt != null) {
90+
expirationTime = new Date(Long.parseLong(expiresAt));
91+
}
92+
AccessToken accessToken = new AccessToken(vendedOAuthToken, expirationTime);
93+
storageOptionsBuilder.setCredentials(GoogleCredentials.create(accessToken));
94+
}
95+
else {
96+
gcsAuth.setAuth(storageOptionsBuilder, identity);
97+
}
6898

69-
endpoint.ifPresent(storageOptionsBuilder::setHost);
99+
String vendedServiceHost = extraCredentials.get(EXTRA_CREDENTIALS_GCS_SERVICE_HOST_PROPERTY);
100+
if (vendedServiceHost != null) {
101+
storageOptionsBuilder.setHost(vendedServiceHost);
102+
}
103+
else {
104+
endpoint.ifPresent(storageOptionsBuilder::setHost);
105+
}
106+
107+
String vendedUserProject = extraCredentials.get(EXTRA_CREDENTIALS_GCS_USER_PROJECT_PROPERTY);
108+
if (vendedUserProject != null) {
109+
storageOptionsBuilder.setQuotaProjectId(vendedUserProject);
110+
}
70111

71112
// Note: without uniform strategy we cannot retry idempotent operations.
72113
// The trino-filesystem api does not violate the conditions for idempotency, see https://cloud.google.com/storage/docs/retry-strategy#java for details.

‎lib/trino-filesystem-gcs/src/test/java/io/trino/filesystem/gcs/TestGcsStorageFactory.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,21 @@
1414
package io.trino.filesystem.gcs;
1515

1616
import com.google.auth.Credentials;
17+
import com.google.auth.oauth2.AccessToken;
18+
import com.google.auth.oauth2.GoogleCredentials;
1719
import com.google.cloud.NoCredentials;
1820
import com.google.cloud.storage.Storage;
21+
import com.google.common.collect.ImmutableMap;
1922
import io.trino.spi.security.ConnectorIdentity;
2023
import org.junit.jupiter.api.Test;
2124

2225
import static io.trino.filesystem.gcs.GcsFileSystemConfig.AuthType;
26+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_NO_AUTH_PROPERTY;
27+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_EXPIRES_AT_PROPERTY;
28+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY;
29+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_PROJECT_ID_PROPERTY;
30+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_SERVICE_HOST_PROPERTY;
31+
import static io.trino.filesystem.gcs.GcsFileSystemConstants.EXTRA_CREDENTIALS_GCS_USER_PROJECT_PROPERTY;
2332
import static org.assertj.core.api.Assertions.assertThat;
2433

2534
final class TestGcsStorageFactory
@@ -38,4 +47,159 @@ void testApplicationDefaultCredentials()
3847

3948
assertThat(actualCredentials).isEqualTo(NoCredentials.getInstance());
4049
}
50+
51+
@Test
52+
void testVendedOAuthToken()
53+
throws Exception
54+
{
55+
GcsFileSystemConfig config = new GcsFileSystemConfig().setAuthType(AuthType.APPLICATION_DEFAULT);
56+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
57+
58+
ConnectorIdentity identity = ConnectorIdentity.forUser("test")
59+
.withExtraCredentials(ImmutableMap.of(
60+
EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY, "ya29.test-token"))
61+
.build();
62+
63+
try (Storage storage = storageFactory.create(identity)) {
64+
Credentials credentials = storage.getOptions().getCredentials();
65+
assertThat(credentials).isInstanceOf(GoogleCredentials.class);
66+
GoogleCredentials googleCredentials = (GoogleCredentials) credentials;
67+
AccessToken accessToken = googleCredentials.getAccessToken();
68+
assertThat(accessToken).isNotNull();
69+
assertThat(accessToken.getTokenValue()).isEqualTo("ya29.test-token");
70+
}
71+
}
72+
73+
@Test
74+
void testVendedOAuthTokenWithExpiration()
75+
throws Exception
76+
{
77+
GcsFileSystemConfig config = new GcsFileSystemConfig().setAuthType(AuthType.APPLICATION_DEFAULT);
78+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
79+
80+
ConnectorIdentity identity = ConnectorIdentity.forUser("test")
81+
.withExtraCredentials(ImmutableMap.of(
82+
EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY, "ya29.test-token",
83+
EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_EXPIRES_AT_PROPERTY, "1700000000000"))
84+
.build();
85+
86+
try (Storage storage = storageFactory.create(identity)) {
87+
Credentials credentials = storage.getOptions().getCredentials();
88+
assertThat(credentials).isInstanceOf(GoogleCredentials.class);
89+
GoogleCredentials googleCredentials = (GoogleCredentials) credentials;
90+
AccessToken accessToken = googleCredentials.getAccessToken();
91+
assertThat(accessToken).isNotNull();
92+
assertThat(accessToken.getTokenValue()).isEqualTo("ya29.test-token");
93+
assertThat(accessToken.getExpirationTime()).isNotNull();
94+
assertThat(accessToken.getExpirationTime().getTime()).isEqualTo(1700000000000L);
95+
}
96+
}
97+
98+
@Test
99+
void testVendedProjectId()
100+
throws Exception
101+
{
102+
GcsFileSystemConfig config = new GcsFileSystemConfig()
103+
.setAuthType(AuthType.APPLICATION_DEFAULT)
104+
.setProjectId("static-project");
105+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
106+
107+
ConnectorIdentity identity = ConnectorIdentity.forUser("test")
108+
.withExtraCredentials(ImmutableMap.of(
109+
EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY, "ya29.test-token",
110+
EXTRA_CREDENTIALS_GCS_PROJECT_ID_PROPERTY, "vended-project"))
111+
.build();
112+
113+
try (Storage storage = storageFactory.create(identity)) {
114+
assertThat(storage.getOptions().getProjectId()).isEqualTo("vended-project");
115+
}
116+
}
117+
118+
@Test
119+
void testVendedServiceHost()
120+
throws Exception
121+
{
122+
GcsFileSystemConfig config = new GcsFileSystemConfig()
123+
.setAuthType(AuthType.APPLICATION_DEFAULT);
124+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
125+
126+
ConnectorIdentity identity = ConnectorIdentity.forUser("test")
127+
.withExtraCredentials(ImmutableMap.of(
128+
EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY, "ya29.test-token",
129+
EXTRA_CREDENTIALS_GCS_SERVICE_HOST_PROPERTY, "https://custom-storage.googleapis.com"))
130+
.build();
131+
132+
try (Storage storage = storageFactory.create(identity)) {
133+
assertThat(storage.getOptions().getHost()).isEqualTo("https://custom-storage.googleapis.com");
134+
}
135+
}
136+
137+
@Test
138+
void testVendedNoAuth()
139+
throws Exception
140+
{
141+
GcsFileSystemConfig config = new GcsFileSystemConfig().setAuthType(AuthType.APPLICATION_DEFAULT);
142+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
143+
144+
ConnectorIdentity identity = ConnectorIdentity.forUser("test")
145+
.withExtraCredentials(ImmutableMap.of(
146+
EXTRA_CREDENTIALS_GCS_NO_AUTH_PROPERTY, "true"))
147+
.build();
148+
149+
try (Storage storage = storageFactory.create(identity)) {
150+
assertThat(storage.getOptions().getCredentials()).isEqualTo(NoCredentials.getInstance());
151+
}
152+
}
153+
154+
@Test
155+
void testNoAuthTakesPriorityOverOAuthToken()
156+
throws Exception
157+
{
158+
GcsFileSystemConfig config = new GcsFileSystemConfig().setAuthType(AuthType.APPLICATION_DEFAULT);
159+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
160+
161+
ConnectorIdentity identity = ConnectorIdentity.forUser("test")
162+
.withExtraCredentials(ImmutableMap.of(
163+
EXTRA_CREDENTIALS_GCS_NO_AUTH_PROPERTY, "true",
164+
EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY, "ya29.test-token"))
165+
.build();
166+
167+
try (Storage storage = storageFactory.create(identity)) {
168+
assertThat(storage.getOptions().getCredentials()).isEqualTo(NoCredentials.getInstance());
169+
}
170+
}
171+
172+
@Test
173+
void testVendedUserProject()
174+
throws Exception
175+
{
176+
GcsFileSystemConfig config = new GcsFileSystemConfig()
177+
.setAuthType(AuthType.APPLICATION_DEFAULT);
178+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
179+
180+
ConnectorIdentity identity = ConnectorIdentity.forUser("test")
181+
.withExtraCredentials(ImmutableMap.of(
182+
EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY, "ya29.test-token",
183+
EXTRA_CREDENTIALS_GCS_USER_PROJECT_PROPERTY, "billing-project"))
184+
.build();
185+
186+
try (Storage storage = storageFactory.create(identity)) {
187+
assertThat(storage.getOptions().getQuotaProjectId()).isEqualTo("billing-project");
188+
}
189+
}
190+
191+
@Test
192+
void testStaticConfigUsedWithoutVendedCredentials()
193+
throws Exception
194+
{
195+
GcsFileSystemConfig config = new GcsFileSystemConfig()
196+
.setAuthType(AuthType.APPLICATION_DEFAULT)
197+
.setProjectId("static-project");
198+
GcsStorageFactory storageFactory = new GcsStorageFactory(config, new ApplicationDefaultAuth());
199+
200+
try (Storage storage = storageFactory.create(ConnectorIdentity.ofUser("test"))) {
201+
assertThat(storage.getOptions().getProjectId()).isEqualTo("static-project");
202+
assertThat(storage.getOptions().getCredentials()).isEqualTo(NoCredentials.getInstance());
203+
}
204+
}
41205
}

‎lib/trino-filesystem/src/main/java/io/trino/filesystem/gcs/GcsFileSystemConstants.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.filesystem.gcs;
15+
16+
public final class GcsFileSystemConstants
17+
{
18+
public static final String EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_PROPERTY = "internal$gcs_oauth2_token";
19+
public static final String EXTRA_CREDENTIALS_GCS_OAUTH_TOKEN_EXPIRES_AT_PROPERTY = "internal$gcs_oauth2_token_expires_at";
20+
public static final String EXTRA_CREDENTIALS_GCS_PROJECT_ID_PROPERTY = "internal$gcs_project_id";
21+
public static final String EXTRA_CREDENTIALS_GCS_SERVICE_HOST_PROPERTY = "internal$gcs_service_host";
22+
public static final String EXTRA_CREDENTIALS_GCS_NO_AUTH_PROPERTY = "internal$gcs_no_auth";
23+
public static final String EXTRA_CREDENTIALS_GCS_USER_PROJECT_PROPERTY = "internal$gcs_user_project";
24+
public static final String EXTRA_CREDENTIALS_GCS_ENCRYPTION_KEY_PROPERTY = "internal$gcs_encryption_key";
25+
public static final String EXTRA_CREDENTIALS_GCS_DECRYPTION_KEY_PROPERTY = "internal$gcs_decryption_key";
26+
27+
private GcsFileSystemConstants() {}
28+
}

0 commit comments

Comments
 (0)