diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java index 1ad08f0ba1a2..cfeb12dcae5c 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java @@ -261,12 +261,18 @@ public boolean shouldRetry(IOException e) { this.credentials = credentials; this.maxBytesRewrittenPerCall = null; this.numRewriteTokensUsed = null; - googleCloudStorageOptions = + GoogleCloudStorageOptions.Builder optionsBuilder = GoogleCloudStorageOptions.builder() .setAppName("Beam") .setReadChannelOptions(gcsReadOptions) - .setGrpcEnabled(shouldUseGrpc) - .build(); + .setGrpcEnabled(shouldUseGrpc); + if (storageClient.getRootUrl() != null) { + optionsBuilder.setStorageRootUrl(storageClient.getRootUrl()); + } + if (storageClient.getServicePath() != null) { + optionsBuilder.setStorageServicePath(storageClient.getServicePath()); + } + googleCloudStorageOptions = optionsBuilder.build(); try { googleCloudStorage = createGoogleCloudStorage(googleCloudStorageOptions, storageClient, credentials); @@ -495,6 +501,11 @@ private Long toFileSize(StorageObjectOrIOException storageObjectOrIOException) } } + @VisibleForTesting + GoogleCloudStorage getGoogleCloudStorage() { + return googleCloudStorage; + } + @VisibleForTesting void setCloudStorageImpl(GoogleCloudStorage g) { googleCloudStorage = g; diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java index d32ca162e3fd..2f77f15dcffc 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java @@ -63,6 +63,7 @@ import com.google.auth.Credentials; import com.google.cloud.hadoop.gcsio.CreateObjectOptions; import com.google.cloud.hadoop.gcsio.GoogleCloudStorage; +import com.google.cloud.hadoop.gcsio.GoogleCloudStorageImpl; import com.google.cloud.hadoop.gcsio.GoogleCloudStorageOptions; import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions; import com.google.cloud.hadoop.gcsio.StorageResourceId; @@ -1866,6 +1867,18 @@ public void testReadMetricsAreNotCollectedWhenNotEnabledOpenWithOptions() throws testReadMetrics(false, GoogleCloudStorageReadOptions.DEFAULT); } + @Test + public void testGcsEndpoint() throws IOException { + GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); + pipelineOptions.setGcsEndpoint("http://localhost:4443/storage/v1/"); + + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + GoogleCloudStorageImpl gcsImpl = + (GoogleCloudStorageImpl) gcsUtil.delegate.getGoogleCloudStorage(); + assertEquals("http://localhost:4443/", gcsImpl.getOptions().getStorageRootUrl()); + assertEquals("storage/v1/", gcsImpl.getOptions().getStorageServicePath()); + } + /** A helper to wrap a {@link GenericJson} object in a content stream. */ private static InputStream toStream(String content) throws IOException { return new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8));