GCP: Add gcs.write.object-context.* catalog property support(parity with S3 write.tags)

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.
+   *
+   * <p>Each catalog property of the form {@code gcs.write.object-context.<key>=<value>} results in
+   * an object context entry {@code key → value} being set on the new GCS object created by this
+   * FileIO.
+   *
+   * <p>Example catalog property:
+   *
+   * <pre>gcs.write.object-context.env=prod</pre>
+   */
+  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<String, String> writeObjectContexts = ImmutableMap.of();
+
   @VisibleForTesting
   List<String> parseCommaSeparatedList(String input, List<String> defaultValue) {
     if (input == null || input.trim().isEmpty()) {
@@ -197,6 +212,24 @@ public GCPProperties(Map<String, String> properties) {
 
     gcsAnalyticsCoreEnabled =
         PropertyUtil.propertyAsBoolean(properties, GCS_ANALYTICS_CORE_ENABLED, false);
+
+    ImmutableMap.Builder<String, String> contexts = ImmutableMap.builder();
+
+    for (Map.Entry<String, String> 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<key>=<value>'",
+                  entry.getKey(), GCS_WRITE_OBJECT_CONTEXT_PREFIX));
+        }
+        contexts.put(contextKey, entry.getValue());
+      }
+    }
+
+    this.writeObjectContexts = contexts.buildOrThrow();
   }
 
   public Optional<Integer> channelReadChunkSize() {
@@ -278,4 +311,8 @@ public Map<String, String> properties() {
   public boolean isGcsAnalyticsCoreEnabled() {
     return gcsAnalyticsCoreEnabled;
   }
+
+  public Map<String, String> writeObjectContexts() {
+    return writeObjectContexts;
+  }
 }

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.
+   *
+   * <p>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.
+   *
+   * <p>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<String, String> writeContexts = gcpProperties.writeObjectContexts();
+    if (writeContexts != null && !writeContexts.isEmpty()) {
+      Map<String, BlobInfo.ObjectCustomContextPayload> 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);
 

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");
+  }
 }

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<String, String> blankContexts = ImmutableMap.of();
+    Map<String, String> 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<String, String> 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<String, String> contexts) {
+    if (contexts == null || contexts.isEmpty()) {
+      return properties;
+    }
+
+    Map<String, String> 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<String, String> contexts) {
+    if (contexts == null || contexts.isEmpty()) {
+      return;
+    }
+
+    Map<String, ObjectCustomContextPayload> customContexts =
+        storage.get(blobId).asBlobInfo().getContexts().getCustom();
+
+    assertThat(customContexts).as("GCS object should have custom contexts attached").isNotNull();
+
+    Map<String, String> 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) {