Add redirection awareness for DROP COLUMN task

trinodb · Mar 7, 2022 · 06b19c5 · 06b19c5
1 parent 4e13342
commit 06b19c5
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 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
@@ -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/...product-tests/src/main/java/io/trino/tests/product/hive/TestHiveRedirectionToIceberg.java b/...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/...duct-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergRedirectionToHive.java b/...duct-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()
     {