diff --git a/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java b/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java index 2702ce565d4e..1659fc7b30a0 100644 --- a/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java +++ b/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java @@ -73,6 +73,19 @@ public class GCPProperties implements Serializable { /** Controls whether analytics core library is enabled or not. Defaults to false. */ public static final String GCS_ANALYTICS_CORE_ENABLED = "gcs.analytics-core.enabled"; + /** + * Prefix for GCS object context entries attached to objects at write time. + * + *

Each catalog property of the form {@code gcs.write.object-context.=} results in + * an object context entry {@code key → value} being set on the new GCS object created by this + * FileIO. + * + *

Example catalog property: + * + *

gcs.write.object-context.env=prod
+ */ + public static final String GCS_WRITE_OBJECT_CONTEXT_PREFIX = "gcs.write.object-context."; + /** * Max possible batch size for deletion. Currently, a max of 100 keys is advised, so we default to * a number below that. https://cloud.google.com/storage/docs/batch @@ -106,6 +119,8 @@ public class GCPProperties implements Serializable { private int gcsDeleteBatchSize = GCS_DELETE_BATCH_SIZE_DEFAULT; + private Map writeObjectContexts = ImmutableMap.of(); + @VisibleForTesting List parseCommaSeparatedList(String input, List defaultValue) { if (input == null || input.trim().isEmpty()) { @@ -197,6 +212,24 @@ public GCPProperties(Map properties) { gcsAnalyticsCoreEnabled = PropertyUtil.propertyAsBoolean(properties, GCS_ANALYTICS_CORE_ENABLED, false); + + ImmutableMap.Builder contexts = ImmutableMap.builder(); + + for (Map.Entry entry : properties.entrySet()) { + if (entry.getKey().startsWith(GCS_WRITE_OBJECT_CONTEXT_PREFIX)) { + String contextKey = entry.getKey().substring(GCS_WRITE_OBJECT_CONTEXT_PREFIX.length()); + if (contextKey.isEmpty()) { + throw new IllegalArgumentException( + String.format( + "Invalid GCS write object-context property '%s': context key must not be empty. " + + "Use the format '%s='", + entry.getKey(), GCS_WRITE_OBJECT_CONTEXT_PREFIX)); + } + contexts.put(contextKey, entry.getValue()); + } + } + + this.writeObjectContexts = contexts.buildOrThrow(); } public Optional channelReadChunkSize() { @@ -278,4 +311,8 @@ public Map properties() { public boolean isGcsAnalyticsCoreEnabled() { return gcsAnalyticsCoreEnabled; } + + public Map writeObjectContexts() { + return writeObjectContexts; + } } diff --git a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java index 3fb8aa4801ae..e2360cedc05c 100644 --- a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java +++ b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java @@ -28,6 +28,8 @@ import java.nio.channels.Channels; import java.util.Arrays; import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; import org.apache.iceberg.gcp.GCPProperties; import org.apache.iceberg.io.FileIOMetricsContext; import org.apache.iceberg.io.PositionOutputStream; @@ -48,7 +50,7 @@ class GCSOutputStream extends PositionOutputStream { private final StackTraceElement[] createStack; private final Storage storage; - private final BlobId blobId; + private final BlobInfo blobInfo; private final GCPProperties gcpProperties; private OutputStream stream; @@ -63,8 +65,8 @@ class GCSOutputStream extends PositionOutputStream { Storage storage, BlobId blobId, GCPProperties gcpProperties, MetricsContext metrics) throws IOException { this.storage = storage; - this.blobId = blobId; this.gcpProperties = gcpProperties; + this.blobInfo = buildBlobInfo(blobId, gcpProperties); createStack = Thread.currentThread().getStackTrace(); @@ -74,6 +76,38 @@ class GCSOutputStream extends PositionOutputStream { openStream(); } + /** + * Builds the {@link BlobInfo} used to initiate the resumable upload. + * + *

When {@code gcs.write.object-context.*} properties are present, builds an {@link + * BlobInfo.ObjectContexts} from them and sets it on the builder via {@code + * BlobInfo.Builder.setContexts()}. The SDK then includes the contexts in the {@code + * WriteObjectRequest} / {@code StartResumableWriteRequest} wire payload. + * + *

When no context properties are configured this method behaves identically to the original + * implementation — {@code setContexts} is never called and there is zero overhead. + */ + private static BlobInfo buildBlobInfo(BlobId blobId, GCPProperties gcpProperties) { + BlobInfo.Builder builder = BlobInfo.newBuilder(blobId); + + Map writeContexts = gcpProperties.writeObjectContexts(); + if (writeContexts != null && !writeContexts.isEmpty()) { + Map payloads = + writeContexts.entrySet().stream() + .collect( + Collectors.toMap( + Map.Entry::getKey, + e -> + BlobInfo.ObjectCustomContextPayload.newBuilder() + .setValue(e.getValue()) + .build())); + + builder.setContexts(BlobInfo.ObjectContexts.newBuilder().setCustom(payloads).build()); + } + + return builder.build(); + } + @Override public long getPos() { return pos; @@ -110,9 +144,7 @@ private void openStream() { .userProject() .ifPresent(userProject -> writeOptions.add(BlobWriteOption.userProject(userProject))); - WriteChannel channel = - storage.writer( - BlobInfo.newBuilder(blobId).build(), writeOptions.toArray(new BlobWriteOption[0])); + WriteChannel channel = storage.writer(blobInfo, writeOptions.toArray(new BlobWriteOption[0])); gcpProperties.channelWriteChunkSize().ifPresent(channel::setChunkSize); diff --git a/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java b/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java index 0ec2183fc355..7311764f439b 100644 --- a/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java +++ b/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java @@ -23,6 +23,7 @@ import static org.apache.iceberg.gcp.GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT; import static org.apache.iceberg.gcp.GCPProperties.GCS_OAUTH2_TOKEN; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException; import static org.assertj.core.api.Assertions.assertThatIllegalStateException; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; @@ -77,4 +78,73 @@ public void refreshCredentialsEndpointSetButRefreshDisabled() { .get() .isEqualTo("/v1/credentials"); } + + @Test + public void testWriteObjectContextsEmpty() { + GCPProperties props = new GCPProperties(ImmutableMap.of("gcs.project-id", "my-project")); + assertThat(props.writeObjectContexts()).isEmpty(); + } + + @Test + public void testWriteObjectContextsSingleEntry() { + GCPProperties props = + new GCPProperties(ImmutableMap.of("gcs.write.object-context.env", "prod")); + + assertThat(props.writeObjectContexts()).isEqualTo(ImmutableMap.of("env", "prod")); + } + + @Test + public void testWriteObjectContextsMultipleEntries() { + GCPProperties props = + new GCPProperties( + ImmutableMap.of( + "gcs.write.object-context.env", "prod", + "gcs.write.object-context.iceberg-table", "orders")); + + assertThat(props.writeObjectContexts()) + .isEqualTo( + ImmutableMap.of( + "env", "prod", + "iceberg-table", "orders")); + } + + @Test + public void testWriteObjectContextsDoNotIncludeOtherProperties() { + // unrelated gcs.* properties must not leak into the context map + GCPProperties props = + new GCPProperties( + ImmutableMap.of( + "gcs.project-id", "my-project", + "gcs.user-project", "billing-project", + "gcs.write.object-context.team", "data-platform")); + + assertThat(props.writeObjectContexts()).containsOnlyKeys("team"); + } + + @Test + public void testWriteObjectContextsDoNotAffectExistingProperties() { + // adding context properties must not disturb parsing of pre-existing fields + GCPProperties props = + new GCPProperties( + ImmutableMap.of( + "gcs.project-id", "my-project", + "gcs.channel.write.chunk-size-bytes", "16777216", + "gcs.write.object-context.env", "staging")); + + assertThat(props.projectId()).contains("my-project"); + assertThat(props.channelWriteChunkSize()).contains(16_777_216); + + // FIX: Replaced ImmutableMap.entry with AssertJ's fluid entry syntax + assertThat(props.writeObjectContexts()) + .containsExactlyEntriesOf(ImmutableMap.of("env", "staging")); + } + + @Test + public void testWriteObjectContextsEmptyKeyThrows() { + assertThatIllegalArgumentException() + .isThrownBy( + () -> new GCPProperties(ImmutableMap.of("gcs.write.object-context.", "some-value"))) + .withMessageContaining("gcs.write.object-context.") + .withMessageContaining("context key must not be empty"); + } } diff --git a/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java b/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java index 9a68cfc9a4d5..815e4be5f72c 100644 --- a/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java +++ b/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java @@ -21,15 +21,19 @@ import static org.assertj.core.api.Assertions.assertThat; import com.google.cloud.storage.BlobId; +import com.google.cloud.storage.BlobInfo.ObjectCustomContextPayload; import com.google.cloud.storage.Storage; import com.google.cloud.storage.contrib.nio.testing.LocalStorageHelper; import java.io.IOException; import java.io.UncheckedIOException; +import java.util.Map; import java.util.Random; import java.util.UUID; +import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.iceberg.gcp.GCPProperties; import org.apache.iceberg.metrics.MetricsContext; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,15 +48,33 @@ public class TestGCSOutputStream { @Test public void testWrite() { - // Run tests for both byte and array write paths + Map blankContexts = ImmutableMap.of(); + Map multipleContexts = + ImmutableMap.of( + "env", "prod", + "analytics-id", "123456"); + Stream.of(true, false) .forEach( arrayWrite -> { - // Test small file write - writeAndVerify(storage, randomBlobId(), randomData(1024), arrayWrite); + // Test small file write with no contexts + writeAndVerify(storage, randomBlobId(), randomData(1024), arrayWrite, blankContexts); + + // Test large file write with no contexts + writeAndVerify( + storage, randomBlobId(), randomData(10 * 1024 * 1024), arrayWrite, blankContexts); - // Test large file - writeAndVerify(storage, randomBlobId(), randomData(10 * 1024 * 1024), arrayWrite); + // Test small file write with context + writeAndVerify( + storage, randomBlobId(), randomData(1024), arrayWrite, multipleContexts); + + // Test large file write with context + writeAndVerify( + storage, + randomBlobId(), + randomData(10 * 1024 * 1024), + arrayWrite, + multipleContexts); }); } @@ -64,9 +86,11 @@ public void testMultipleClose() throws IOException { stream.close(); } - private void writeAndVerify(Storage client, BlobId uri, byte[] data, boolean arrayWrite) { + private void writeAndVerify( + Storage client, BlobId uri, byte[] data, boolean arrayWrite, Map contexts) { try (GCSOutputStream stream = - new GCSOutputStream(client, uri, properties, MetricsContext.nullMetrics())) { + new GCSOutputStream( + client, uri, propertiesWithContexts(contexts), MetricsContext.nullMetrics())) { if (arrayWrite) { stream.write(data); assertThat(stream.getPos()).isEqualTo(data.length); @@ -82,6 +106,46 @@ private void writeAndVerify(Storage client, BlobId uri, byte[] data, boolean arr byte[] actual = readGCSData(uri); assertThat(actual).isEqualTo(data); + verifyContexts(uri, contexts); + } + + /** Builds a {@link GCPProperties} with the given context entries added under the write prefix. */ + private GCPProperties propertiesWithContexts(Map contexts) { + if (contexts == null || contexts.isEmpty()) { + return properties; + } + + Map props = + contexts.entrySet().stream() + .collect( + Collectors.toMap( + e -> GCPProperties.GCS_WRITE_OBJECT_CONTEXT_PREFIX + e.getKey(), + Map.Entry::getValue)); + return new GCPProperties(props); + } + + /** Asserts that every entry in {@code contexts} is present on the committed GCS object. */ + private void verifyContexts(BlobId blobId, Map contexts) { + if (contexts == null || contexts.isEmpty()) { + return; + } + + Map customContexts = + storage.get(blobId).asBlobInfo().getContexts().getCustom(); + + assertThat(customContexts).as("GCS object should have custom contexts attached").isNotNull(); + + Map actualContextValues = + customContexts.entrySet().stream() + .collect( + Collectors.toMap( + Map.Entry::getKey, + e -> e.getValue().getValue() // Extract the string value from the payload + )); + + assertThat(actualContextValues) + .as("GCS object metadata should contain the configured contexts") + .containsAllEntriesOf(contexts); } private byte[] readGCSData(BlobId blobId) {