Skip to content

Commit

Permalink
Allow configuring gcs service endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
wendigo committed Jan 6, 2025
1 parent 2ef6dc1 commit d4be693
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import jakarta.validation.constraints.NotNull;
import jakarta.validation.constraints.Size;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkState;
Expand All @@ -39,6 +40,7 @@ public class GcsFileSystemConfig
private int batchSize = 100;

private String projectId;
private Optional<String> endpoint = Optional.empty();

private boolean useGcsAccessToken;
private String jsonKey;
Expand Down Expand Up @@ -120,6 +122,19 @@ public GcsFileSystemConfig setProjectId(String projectId)
return this;
}

public Optional<String> getEndpoint()
{
return endpoint;
}

@ConfigDescription("Endpoint to use for GCS requests")
@Config("gcs.endpoint")
public GcsFileSystemConfig setEndpoint(Optional<String> endpoint)
{
this.endpoint = endpoint;
return this;
}

public boolean isUseGcsAccessToken()
{
return useGcsAccessToken;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ public class GcsStorageFactory
public static final String GCS_OAUTH_KEY = "gcs.oauth";
public static final List<String> DEFAULT_SCOPES = ImmutableList.of("https://www.googleapis.com/auth/cloud-platform");
private final String projectId;
private final Optional<String> endpoint;
private final boolean useGcsAccessToken;
private final Optional<GoogleCredentials> jsonGoogleCredential;
private final int maxRetries;
Expand All @@ -56,6 +57,7 @@ public GcsStorageFactory(GcsFileSystemConfig config)
{
config.validate();
projectId = config.getProjectId();
endpoint = config.getEndpoint();
useGcsAccessToken = config.isUseGcsAccessToken();
String jsonKey = config.getJsonKey();
String jsonKeyFilePath = config.getJsonKeyFilePath();
Expand Down Expand Up @@ -105,6 +107,9 @@ public Storage create(ConnectorIdentity identity)
if (projectId != null) {
storageOptionsBuilder.setProjectId(projectId);
}

endpoint.ifPresent(storageOptionsBuilder::setHost);

// Note: without uniform strategy we cannot retry idempotent operations.
// The trino-filesystem api does not violate the conditions for idempotency, see https://cloud.google.com/storage/docs/retry-strategy#java for details.
return storageOptionsBuilder
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Map;
import java.util.Optional;

import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
Expand All @@ -44,6 +45,7 @@ void testDefaults()
.setPageSize(100)
.setBatchSize(100)
.setProjectId(null)
.setEndpoint(Optional.empty())
.setUseGcsAccessToken(false)
.setJsonKey(null)
.setJsonKeyFilePath(null)
Expand All @@ -67,6 +69,7 @@ void testExplicitPropertyMappings()
.put("gcs.page-size", "10")
.put("gcs.batch-size", "11")
.put("gcs.project-id", "project")
.put("gcs.endpoint", "http://custom.dns.org:8000")
.put("gcs.use-access-token", "true")
.put("gcs.json-key", "{}")
.put("gcs.json-key-file-path", jsonKeyFile.toString())
Expand All @@ -84,6 +87,7 @@ void testExplicitPropertyMappings()
.setPageSize(10)
.setBatchSize(11)
.setProjectId("project")
.setEndpoint(Optional.of("http://custom.dns.org:8000"))
.setUseGcsAccessToken(true)
.setJsonKey("{}")
.setJsonKeyFilePath(jsonKeyFile.toString())
Expand Down

0 comments on commit d4be693

Please sign in to comment.