Spark: Don't allow branch_ usage with VERSION AS OF

apache · Dec 5, 2023 · d1fc9c3 · d1fc9c3
1 parent a4d4756
commit d1fc9c3
Show file tree

Hide file tree

Showing 9 changed files with 60 additions and 3 deletions.
diff --git a/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/SparkCatalog.java b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/SparkCatalog.java
@@ -171,7 +171,7 @@ public Table loadTable(Identifier ident, String version) throws NoSuchTableExcep
       SparkTable sparkTable = (SparkTable) table;
 
       Preconditions.checkArgument(
-          sparkTable.snapshotId() == null,
+          sparkTable.snapshotId() == null && sparkTable.branch() == null,
           "Cannot do time-travel based on both table identifier and AS OF");
 
       try {

diff --git a/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java
@@ -173,6 +173,10 @@ public Long snapshotId() {
     return snapshotId;
   }
 
+  public String branch() {
+    return branch;
+  }
+
   public SparkTable copyWithSnapshotId(long newSnapshotId) {
     return new SparkTable(icebergTable, newSnapshotId, refreshEagerly);
   }

diff --git a/spark/v3.3/spark/src/test/java/org/apache/iceberg/spark/sql/TestSelect.java b/spark/v3.3/spark/src/test/java/org/apache/iceberg/spark/sql/TestSelect.java
@@ -426,6 +426,21 @@ public void testInvalidTimeTravelBasedOnBothAsOfAndTableIdentifier() {
         });
   }
 
+  @Test
+  public void testInvalidTimeTravelAgainstBranchIdentifierWithAsOf() {
+    long snapshotId = validationCatalog.loadTable(tableIdent).currentSnapshot().snapshotId();
+    validationCatalog.loadTable(tableIdent).manageSnapshots().createBranch("b1").commit();
+
+    // create a second snapshot
+    sql("INSERT INTO %s VALUES (4, 'd', 4.0), (5, 'e', 5.0)", tableName);
+
+    // using branch_b1 in the table identifier and VERSION AS OF
+    Assertions.assertThatThrownBy(
+            () -> sql("SELECT * FROM %s.branch_b1 VERSION AS OF %s", tableName, snapshotId))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Cannot do time-travel based on both table identifier and AS OF");
+  }
+
   @Test
   public void testSpecifySnapshotAndTimestamp() {
     // get the snapshot ID of the last write

diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkCatalog.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkCatalog.java
@@ -170,7 +170,7 @@ public Table loadTable(Identifier ident, String version) throws NoSuchTableExcep
       SparkTable sparkTable = (SparkTable) table;
 
       Preconditions.checkArgument(
-          sparkTable.snapshotId() == null,
+          sparkTable.snapshotId() == null && sparkTable.branch() == null,
           "Cannot do time-travel based on both table identifier and AS OF");
 
       try {

diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java
@@ -173,6 +173,10 @@ public Long snapshotId() {
     return snapshotId;
   }
 
+  public String branch() {
+    return branch;
+  }
+
   public SparkTable copyWithSnapshotId(long newSnapshotId) {
     return new SparkTable(icebergTable, newSnapshotId, refreshEagerly);
   }

diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/sql/TestSelect.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/sql/TestSelect.java
@@ -421,6 +421,21 @@ public void testInvalidTimeTravelBasedOnBothAsOfAndTableIdentifier() {
         .hasMessage("Cannot do time-travel based on both table identifier and AS OF");
   }
 
+  @Test
+  public void testInvalidTimeTravelAgainstBranchIdentifierWithAsOf() {
+    long snapshotId = validationCatalog.loadTable(tableIdent).currentSnapshot().snapshotId();
+    validationCatalog.loadTable(tableIdent).manageSnapshots().createBranch("b1").commit();
+
+    // create a second snapshot
+    sql("INSERT INTO %s VALUES (4, 'd', 4.0), (5, 'e', 5.0)", tableName);
+
+    // using branch_b1 in the table identifier and VERSION AS OF
+    Assertions.assertThatThrownBy(
+            () -> sql("SELECT * FROM %s.branch_b1 VERSION AS OF %s", tableName, snapshotId))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Cannot do time-travel based on both table identifier and AS OF");
+  }
+
   @Test
   public void testSpecifySnapshotAndTimestamp() {
     // get the snapshot ID of the last write

diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkCatalog.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkCatalog.java
@@ -170,7 +170,7 @@ public Table loadTable(Identifier ident, String version) throws NoSuchTableExcep
       SparkTable sparkTable = (SparkTable) table;
 
       Preconditions.checkArgument(
-          sparkTable.snapshotId() == null,
+          sparkTable.snapshotId() == null && sparkTable.branch() == null,
           "Cannot do time-travel based on both table identifier and AS OF");
 
       try {

diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java
@@ -173,6 +173,10 @@ public Long snapshotId() {
     return snapshotId;
   }
 
+  public String branch() {
+    return branch;
+  }
+
   public SparkTable copyWithSnapshotId(long newSnapshotId) {
     return new SparkTable(icebergTable, newSnapshotId, refreshEagerly);
   }

diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/sql/TestSelect.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/sql/TestSelect.java
@@ -421,6 +421,21 @@ public void testInvalidTimeTravelBasedOnBothAsOfAndTableIdentifier() {
         .hasMessage("Cannot do time-travel based on both table identifier and AS OF");
   }
 
+  @Test
+  public void testInvalidTimeTravelAgainstBranchIdentifierWithAsOf() {
+    long snapshotId = validationCatalog.loadTable(tableIdent).currentSnapshot().snapshotId();
+    validationCatalog.loadTable(tableIdent).manageSnapshots().createBranch("b1").commit();
+
+    // create a second snapshot
+    sql("INSERT INTO %s VALUES (4, 'd', 4.0), (5, 'e', 5.0)", tableName);
+
+    // using branch_b1 in the table identifier and VERSION AS OF
+    Assertions.assertThatThrownBy(
+            () -> sql("SELECT * FROM %s.branch_b1 VERSION AS OF %s", tableName, snapshotId))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Cannot do time-travel based on both table identifier and AS OF");
+  }
+
   @Test
   public void testSpecifySnapshotAndTimestamp() {
     // get the snapshot ID of the last write