Add support for drop tag for Iceberg Table

prestodb · Sep 11, 2024 · 8f4d3fa · 8f4d3fa
1 parent 6d7b1a2
commit 8f4d3fa
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/presto-docs/src/main/sphinx/connector/iceberg.rst b/presto-docs/src/main/sphinx/connector/iceberg.rst
@@ -1113,6 +1113,8 @@ Alter table operations are supported in the Iceberg connector::
 
      ALTER TABLE iceberg.web.page_views DROP BRANCH 'branch1';
 
+     ALTER TABLE iceberg.web.page_views DROP TAG 'tag1';
+
 To add a new column as a partition column, identify the transform functions for the column.
 The table is partitioned by the transformed value of the column::
 

diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java
@@ -648,6 +648,21 @@ public void dropBranch(ConnectorSession session, ConnectorTableHandle tableHandl
         }
     }
 
+    @Override
+    public void dropTag(ConnectorSession session, ConnectorTableHandle tableHandle, Optional<String> tagName)
+    {
+        String tag = tagName.get().replace("'", "");
+        IcebergTableHandle icebergTableHandle = (IcebergTableHandle) tableHandle;
+        verify(icebergTableHandle.getIcebergTableName().getTableType() == DATA, "only the data table can have tag dropped");
+        Table icebergTable = getIcebergTable(session, icebergTableHandle.getSchemaTableName());
+        if (icebergTable.refs().containsKey(tag) && icebergTable.refs().get(tag).isTag()) {
+            icebergTable.manageSnapshots().removeTag(tag).commit();
+        }
+        else {
+            throw new PrestoException(NOT_FOUND, format("Tag %s doesn't exist in table %s", tagName.get(), icebergTableHandle.getSchemaTableName().getTableName()));
+        }
+    }
+
     @Override
     public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column)
     {

diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java
@@ -1649,6 +1649,33 @@ public void testDropBranch()
         assertQuerySucceeds("DROP TABLE test_table_branch");
     }
 
+    @Test
+    public void testDropTag()
+    {
+        assertUpdate("CREATE TABLE test_table_tag (id1 BIGINT, id2 BIGINT)");
+        assertUpdate("INSERT INTO test_table_tag VALUES (0, 00), (1, 10)", 2);
+
+        Table icebergTable = loadTable("test_table_tag");
+        icebergTable.manageSnapshots().createTag("testTag1", icebergTable.currentSnapshot().snapshotId()).commit();
+
+        assertUpdate("INSERT INTO test_table_tag VALUES (2, 30), (3, 30)", 2);
+        icebergTable.manageSnapshots().createTag("testTag2", icebergTable.currentSnapshot().snapshotId()).commit();
+        assertUpdate("INSERT INTO test_table_tag VALUES (4, 40), (5, 50)", 2);
+        assertEquals(icebergTable.refs().size(), 3);
+
+        assertQuery("SELECT count(*) FROM test_table_tag FOR SYSTEM_VERSION AS OF 'testTag1'", "VALUES 2");
+        assertQuery("SELECT count(*) FROM test_table_tag FOR SYSTEM_VERSION AS OF 'testTag2'", "VALUES 4");
+        assertQuery("SELECT count(*) FROM test_table_tag FOR SYSTEM_VERSION AS OF 'main'", "VALUES 6");
+
+        assertQuerySucceeds("ALTER TABLE test_table_tag DROP TAG 'testTag1'");
+        icebergTable = loadTable("test_table_tag");
+        assertEquals(icebergTable.refs().size(), 2);
+        assertQueryFails("ALTER TABLE test_table_tag DROP TAG 'testTagNotExist'", "Tag 'testTagNotExist' doesn't exist in table test_table_tag");
+        assertQuerySucceeds("ALTER TABLE test_table_tag DROP TAG IF EXISTS 'testTag2'");
+        assertQueryFails("ALTER TABLE test_table_tag DROP TAG IF EXISTS 'testTagNotExist'", "Tag 'testTagNotExist' doesn't exist in table test_table_tag");
+        assertQuerySucceeds("DROP TABLE test_table_tag");
+    }
+
     @Test
     public void testRefsTable()
     {