Fix reading of specific Iceberg snapshots

Snapshot ids were incorrectly cached and when
more than one snapshot id was used in the query
only first one was used.

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

@@ -200,7 +200,7 @@ public class IcebergMetadata
     private final String trinoVersion;
     private final boolean useUniqueTableLocation;
 
-    private final Map<String, Optional<Long>> snapshotIds = new ConcurrentHashMap<>();
+    private final Map<String, Long> snapshotIds = new ConcurrentHashMap<>();
     private final Map<SchemaTableName, TableMetadata> tableMetadataCache = new ConcurrentHashMap<>();
     private final ViewReaderUtil.PrestoViewReader viewReader = new ViewReaderUtil.PrestoViewReader();
 
@@ -1024,9 +1024,12 @@ public TableStatistics getTableStatistics(ConnectorSession session, ConnectorTab
     private Optional<Long> getSnapshotId(org.apache.iceberg.Table table, Optional<Long> snapshotId)
     {
         // table.name() is an encoded version of SchemaTableName
-        return snapshotIds.computeIfAbsent(table.name(), ignored -> snapshotId
-                .map(id -> IcebergUtil.resolveSnapshotId(table, id))
-                .or(() -> Optional.ofNullable(table.currentSnapshot()).map(Snapshot::snapshotId)));
+        return snapshotId
+                .map(id ->
+                        snapshotIds.computeIfAbsent(
+                                table.name() + "@" + id,
+                                ignored -> IcebergUtil.resolveSnapshotId(table, id)))
+                .or(() -> Optional.ofNullable(table.currentSnapshot()).map(Snapshot::snapshotId));
     }
 
     org.apache.iceberg.Table getIcebergTable(ConnectorSession session, SchemaTableName schemaTableName)

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

@@ -0,0 +1,79 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.iceberg;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import io.trino.testing.AbstractTestQueryFramework;
+import io.trino.testing.QueryRunner;
+import org.testng.annotations.Test;
+
+import java.util.List;
+
+import static io.trino.plugin.iceberg.IcebergQueryRunner.createIcebergQueryRunner;
+import static io.trino.testing.sql.TestTable.randomTableSuffix;
+import static java.lang.String.format;
+import static java.util.stream.Collectors.toList;
+
+public class TestIcebergSnapshots
+        extends AbstractTestQueryFramework
+{
+    private static final int ID_FIELD = 0;
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return createIcebergQueryRunner(ImmutableMap.of(), new IcebergConfig().getFileFormat(), ImmutableList.of());
+    }
+
+    @Test
+    public void testReadingFromSpecificSnapshot()
+    {
+        String tableName = "test_reading_snapshot" + randomTableSuffix();
+        assertUpdate(format("CREATE TABLE %s (a bigint, b bigint)", tableName));
+        assertUpdate(format("INSERT INTO %s VALUES(1, 1)", tableName), 1);
+        List<Long> ids = getSnapshotsIdsByCreationOrder(tableName);
+
+        assertQuery(format("SELECT count(*) FROM \"%s@%d\"", tableName, ids.get(0)), "VALUES(0)");
+        assertQuery(format("SELECT * FROM \"%s@%d\"", tableName, ids.get(1)), "VALUES(1,1)");
+        assertUpdate(format("DROP TABLE %s", tableName));
+    }
+
+    @Test
+    public void testSelectWithMoreThanOneSnapshotOfTheSameTable()
+    {
+        String tableName = "test_reading_snapshot" + randomTableSuffix();
+        assertUpdate(format("CREATE TABLE %s (a bigint, b bigint)", tableName));
+        assertUpdate(format("INSERT INTO %s VALUES(1, 1)", tableName), 1);
+        assertUpdate(format("INSERT INTO %s VALUES(2, 2)", tableName), 1);
+        assertUpdate(format("INSERT INTO %s VALUES(3, 3)", tableName), 1);
+        List<Long> ids = getSnapshotsIdsByCreationOrder(tableName);
+
+        assertQuery(format("SELECT * FROM %s", tableName), "SELECT * FROM (VALUES(1,1), (2,2), (3,3))");
+        assertQuery(
+                format("SELECT * FROM %1$s EXCEPT (SELECT * FROM \"%1$s@%2$d\" EXCEPT SELECT * FROM \"%1$s@%3$d\")", tableName, ids.get(2), ids.get(1)),
+                "SELECT * FROM (VALUES(1,1), (3,3))");
+        assertUpdate(format("DROP TABLE %s", tableName));
+    }
+
+    private List<Long> getSnapshotsIdsByCreationOrder(String tableName)
+    {
+        return getQueryRunner().execute(
+                format("SELECT snapshot_id FROM \"%s$snapshots\" ORDER BY committed_at", tableName))
+                .getMaterializedRows().stream()
+                .map(row -> (Long) row.getField(ID_FIELD))
+                .collect(toList());
+    }
+}