Return only published main branch snapshots in history metadata table

Author: maswin

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/HistoryTable.java

@@ -25,11 +25,14 @@
 import io.trino.spi.connector.SystemTable;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.type.TimeZoneKey;
+import org.apache.iceberg.HistoryEntry;
 import org.apache.iceberg.Snapshot;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.util.SnapshotUtil;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import static io.trino.spi.type.BigintType.BIGINT;
@@ -76,8 +79,13 @@ public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, Connect
 
         Set<Long> ancestorIds = ImmutableSet.copyOf(SnapshotUtil.currentAncestorIds(icebergTable));
         TimeZoneKey timeZoneKey = session.getTimeZoneKey();
+        Map<Long, Snapshot> snapshots = new HashMap<>();
         for (Snapshot snapshot : icebergTable.snapshots()) {
-            long snapshotId = snapshot.snapshotId();
+            snapshots.put(snapshot.snapshotId(), snapshot);
+        }
+        for (HistoryEntry historyEntry : icebergTable.history()) {
+            Long snapshotId = historyEntry.snapshotId();
+            Snapshot snapshot = snapshots.get(snapshotId);
 
             table.addRow(
                     packDateTimeWithZone(snapshot.timestampMillis(), timeZoneKey),

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergSystemTables.java

@@ -241,6 +241,20 @@ public void testHistoryTable()
 
         // Test the number of history entries
         assertQuery("SELECT count(*) FROM test_schema.\"test_table$history\"", "VALUES 3");
+
+        try (TestTable table = newTrinoTable("test_history", "AS SELECT 1 id")) {
+            assertQuery("SELECT count(*) FROM \"" + table.getName() + "$history\"", "VALUES 1");
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES (2)", 1);
+            assertQuery("SELECT count(*) FROM \"" + table.getName() + "$history\"", "VALUES 2");
+
+            // Perform a stage only commit
+            Table icebergTable = this.loadTable(table.getName());
+            icebergTable.newAppend().set("table_id", "1").stageOnly().commit();
+
+            // Check if history table doesn't contain the stage commit snapshot
+            assertQuery("SELECT count(*) FROM \"" + table.getName() + "$history\"", "VALUES 2");
+            assertQuery("SELECT count(*) FROM \"" + table.getName() + "$snapshots\"", "VALUES 3");
+        }
     }
 
     @Test

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergVendingRestCatalogConnectorSmokeTest.java

@@ -312,6 +312,14 @@ public void testDropTableWithNonExistentTableLocation()
                 .hasMessageMatching("Failed to load table: (.*)");
     }
 
+    @Test
+    @Override
+    public void testCreateOrReplaceTable()
+    {
+        // TODO: Enable this test once tabulario/iceberg-rest image is updated to 1.8.0+
+        // Without this change - https://github.com/apache/iceberg/pull/11779 CREATE OR REPLACE clears snapshotLog resulting in incorrect results in history metadata table
+    }
+
     @Override
     protected boolean isFileSorted(Location path, String sortColumnName)
     {