From 06b19c58dc17c4b8316fb10ac0c81185294e5296 Mon Sep 17 00:00:00 2001
From: Marius Grama <findinpath@gmail.com>
Date: Thu, 3 Mar 2022 14:17:26 +0100
Subject: [PATCH] Add redirection awareness for DROP COLUMN task

---
 .../io/trino/execution/DropColumnTask.java    | 11 +++++-----
 .../hive/TestHiveRedirectionToIceberg.java    | 20 +++++++++++++++++++
 .../iceberg/TestIcebergRedirectionToHive.java | 16 +++++++++++++++
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/core/trino-main/src/main/java/io/trino/execution/DropColumnTask.java b/core/trino-main/src/main/java/io/trino/execution/DropColumnTask.java
index 0dbe3aed38f5..d7b998e29c09 100644
--- a/core/trino-main/src/main/java/io/trino/execution/DropColumnTask.java
+++ b/core/trino-main/src/main/java/io/trino/execution/DropColumnTask.java
@@ -18,6 +18,7 @@
 import io.trino.execution.warnings.WarningCollector;
 import io.trino.metadata.Metadata;
 import io.trino.metadata.QualifiedObjectName;
+import io.trino.metadata.RedirectionAwareTableHandle;
 import io.trino.metadata.TableHandle;
 import io.trino.security.AccessControl;
 import io.trino.spi.connector.ColumnHandle;
@@ -27,7 +28,6 @@
 import javax.inject.Inject;
 
 import java.util.List;
-import java.util.Optional;
 
 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
 import static io.trino.metadata.MetadataUtil.createQualifiedObjectName;
@@ -66,19 +66,18 @@ public ListenableFuture<Void> execute(
     {
         Session session = stateMachine.getSession();
         QualifiedObjectName tableName = createQualifiedObjectName(session, statement, statement.getTable());
-        Optional<TableHandle> tableHandleOptional = metadata.getTableHandle(session, tableName);
-
-        if (tableHandleOptional.isEmpty()) {
+        RedirectionAwareTableHandle redirectionAwareTableHandle = metadata.getRedirectionAwareTableHandle(session, tableName);
+        if (redirectionAwareTableHandle.getTableHandle().isEmpty()) {
             if (!statement.isTableExists()) {
                 throw semanticException(TABLE_NOT_FOUND, statement, "Table '%s' does not exist", tableName);
             }
             return immediateVoidFuture();
         }
-        TableHandle tableHandle = tableHandleOptional.get();
+        TableHandle tableHandle = redirectionAwareTableHandle.getTableHandle().get();
 
         String column = statement.getColumn().getValue().toLowerCase(ENGLISH);
 
-        accessControl.checkCanDropColumn(session.toSecurityContext(), tableName);
+        accessControl.checkCanDropColumn(session.toSecurityContext(), redirectionAwareTableHandle.getRedirectedTableName().orElse(tableName));
 
         ColumnHandle columnHandle = metadata.getColumnHandles(session, tableHandle).get(column);
         if (columnHandle == null) {
diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveRedirectionToIceberg.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveRedirectionToIceberg.java
index 0ce30a4aa016..55b8f1778188 100644
--- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveRedirectionToIceberg.java
+++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveRedirectionToIceberg.java
@@ -332,6 +332,26 @@ public void testAlterTableAddColumn()
         onTrino().executeQuery("DROP TABLE " + icebergTableName);
     }
 
+    @Test(groups = {HIVE_ICEBERG_REDIRECTIONS, PROFILE_SPECIFIC_TESTS})
+    public void testAlterTableDropColumn()
+    {
+        String tableName = "iceberg_alter_table_drop_column_" + randomTableSuffix();
+        String hiveTableName = "hive.default." + tableName;
+        String icebergTableName = "iceberg.default." + tableName;
+
+        createIcebergTable(icebergTableName, false);
+
+        onTrino().executeQuery("ALTER TABLE " + hiveTableName + " DROP COLUMN comment");
+
+        Assertions.assertThat(onTrino().executeQuery("DESCRIBE " + icebergTableName).column(1))
+                .containsOnly("nationkey", "name", "regionkey");
+
+        assertResultsEqual(
+                onTrino().executeQuery("TABLE " + icebergTableName),
+                onTrino().executeQuery("SELECT nationkey, name, regionkey FROM tpch.tiny.nation"));
+        onTrino().executeQuery("DROP TABLE " + icebergTableName);
+    }
+
     @Test(groups = {HIVE_ICEBERG_REDIRECTIONS, PROFILE_SPECIFIC_TESTS})
     public void testCommentTable()
     {
diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergRedirectionToHive.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergRedirectionToHive.java
index f11e21ca9620..f1101d6ec110 100644
--- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergRedirectionToHive.java
+++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergRedirectionToHive.java
@@ -315,6 +315,22 @@ public void testAlterTableAddColumn()
         onTrino().executeQuery("DROP TABLE " + hiveTableName);
     }
 
+    @Test(groups = {HIVE_ICEBERG_REDIRECTIONS, PROFILE_SPECIFIC_TESTS})
+    public void testAlterTableDropColumn()
+    {
+        String tableName = "hive_alter_table_drop_column_" + randomTableSuffix();
+        String hiveTableName = "hive.default." + tableName;
+        String icebergTableName = "iceberg.default." + tableName;
+
+        createHiveTable(hiveTableName, false);
+
+        // TODO: support redirects from Iceberg to Hive
+        assertQueryFailure(() -> onTrino().executeQuery("ALTER TABLE " + icebergTableName + " DROP COLUMN comment"))
+                .hasMessageMatching("\\QQuery failed (#\\E\\S+\\Q): Not an Iceberg table: default." + tableName);
+
+        onTrino().executeQuery("DROP TABLE " + hiveTableName);
+    }
+
     @Test(groups = {HIVE_ICEBERG_REDIRECTIONS, PROFILE_SPECIFIC_TESTS})
     public void testCommentTable()
     {