Unit tests for GCSFileIO

cdouglas · Mar 20, 2024 · 58ce5ce · 58ce5ce
1 parent 475efad
commit 58ce5ce
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 4 deletions.
diff --git a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSFileIO.java b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSFileIO.java
@@ -106,8 +106,10 @@ public OutputFile newOutputFile(String path) {
 
   @Override
   public OutputFile newOutputFile(InputFile replace) {
-    // !#! TODO
-    return null;
+    if (replace instanceof GCSInputFile) {
+      return GCSOutputFile.fromBlobId(((GCSInputFile) replace).blobId(), client(), gcpProperties, metrics);
+    }
+    return GCSOutputFile.fromLocation(replace.location(), client(), gcpProperties, metrics);
   }
 
   @Override

diff --git a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSInputFile.java b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSInputFile.java
@@ -30,7 +30,7 @@ class GCSInputFile extends BaseGCSFile implements InputFile {
 
   static GCSInputFile fromLocation(
       String location, Storage storage, GCPProperties gcpProperties, MetricsContext metrics) {
-    return new GCSInputFile(storage, BlobId.fromGsUtilUri(location), null, gcpProperties, metrics);
+    return fromLocation(location, -1, storage, gcpProperties, metrics);
   }
 
   static GCSInputFile fromLocation(
@@ -39,9 +39,26 @@ static GCSInputFile fromLocation(
       Storage storage,
       GCPProperties gcpProperties,
       MetricsContext metrics) {
+    return fromBlobId(BlobId.fromGsUtilUri(location), length, storage, gcpProperties, metrics);
+  }
+
+  static GCSInputFile fromBlobId(
+          BlobId blobId,
+          Storage storage,
+          GCPProperties gcpProperties,
+          MetricsContext metrics) {
+    return fromBlobId(blobId, -1, storage, gcpProperties, metrics);
+  }
+
+  static GCSInputFile fromBlobId(
+      BlobId blobId,
+      long length,
+      Storage storage,
+      GCPProperties gcpProperties,
+      MetricsContext metrics) {
     return new GCSInputFile(
         storage,
-        BlobId.fromGsUtilUri(location),
+        blobId,
         length > 0 ? length : null,
         gcpProperties,
         metrics);

diff --git a/gcp/src/test/java/org/apache/iceberg/gcp/gcs/GCSFileIOTest.java b/gcp/src/test/java/org/apache/iceberg/gcp/gcs/GCSFileIOTest.java
@@ -27,6 +27,7 @@
 import com.google.cloud.storage.BlobId;
 import com.google.cloud.storage.BlobInfo;
 import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
 import com.google.cloud.storage.contrib.nio.testing.LocalStorageHelper;
 import java.io.IOException;
 import java.io.InputStream;
@@ -47,6 +48,7 @@
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -105,6 +107,75 @@ public void newInputFile() throws IOException {
     assertThat(io.newInputFile(location).exists()).isFalse();
   }
 
+  @Test
+  public void newOutputFileMatch() throws IOException {
+    final String location = format("gs://%s/path/to/file.txt", TEST_BUCKET);
+    final byte[] expected = new byte[1024 * 1024];
+    random.nextBytes(expected);
+
+    final OutputFile out = io.newOutputFile(location);
+    try (OutputStream os = out.createOrOverwrite()) {
+      IOUtil.writeFully(os, ByteBuffer.wrap(expected));
+    }
+
+    final InputFile in = io.newInputFile(location);
+    assertThat(in.exists()).isTrue();
+    final byte[] actual = new byte[1024 * 1024];
+
+    try (InputStream is = in.newStream()) {
+      IOUtil.readFully(is, actual, 0, actual.length);
+    }
+    assertThat(actual).isEqualTo(expected);
+
+    OutputFile overwrite = io.newOutputFile(in);
+    final byte[] overbytes = new byte[1024 * 1024];
+    random.nextBytes(overbytes);
+    try (OutputStream os = overwrite.createOrOverwrite()) {
+      IOUtil.writeFully(os, ByteBuffer.wrap(overbytes));
+    }
+    try (InputStream is = in.newStream()) {
+      IOUtil.readFully(is, actual, 0, actual.length);
+    }
+    assertThat(actual).isEqualTo(overbytes);
+  }
+
+  @Test
+  public void newOutputFileMatchFail() throws IOException {
+    final String location = format("gs://%s/path/to/file.txt", TEST_BUCKET);
+    final byte[] expected = new byte[1024 * 1024];
+    random.nextBytes(expected);
+
+    final OutputFile out = io.newOutputFile(location);
+    try (OutputStream os = out.createOrOverwrite()) {
+      IOUtil.writeFully(os, ByteBuffer.wrap(expected));
+    }
+
+    final InputFile in = io.newInputFile(location);
+    assertThat(in.exists()).isTrue();
+    final byte[] actual = new byte[1024 * 1024];
+    try (InputStream is = in.newStream()) {
+      IOUtil.readFully(is, actual, 0, actual.length);
+    }
+    assertThat(actual).isEqualTo(expected);
+
+    // overwrite succeeds, because generation matches InputFile
+    final OutputFile overwrite = io.newOutputFile(in);
+    final byte[] overbytes = new byte[1024 * 1024];
+    random.nextBytes(overbytes);
+    try (OutputStream os = overwrite.createOrOverwrite()) {
+      IOUtil.writeFully(os, ByteBuffer.wrap(overbytes));
+    }
+    // overwrite fails, object has been overwritten
+    StorageException generationFailure = Assertions.assertThrows(
+            StorageException.class,
+            () -> {
+              try (OutputStream os = overwrite.createOrOverwrite()) {
+                IOUtil.writeFully(os, ByteBuffer.wrap(overbytes));
+              }
+            });
+    assertThat(generationFailure.getMessage()).startsWith("Generation mismatch");
+  }
+
   @Test
   public void testDelete() {
     String path = "delete/path/data.dat";