Add redirection awareness for ADD COLUMN task

trinodb · Mar 1, 2022 · f02dc03 · f02dc03
1 parent 2c1344c
commit f02dc03
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/execution/AddColumnTask.java b/core/trino-main/src/main/java/io/trino/execution/AddColumnTask.java
@@ -19,6 +19,7 @@
 import io.trino.execution.warnings.WarningCollector;
 import io.trino.metadata.ColumnPropertyManager;
 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;
@@ -34,7 +35,6 @@
 
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 
 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
 import static io.trino.metadata.MetadataUtil.createQualifiedObjectName;
@@ -81,20 +81,20 @@ public ListenableFuture<Void> execute(
             WarningCollector warningCollector)
     {
         Session session = stateMachine.getSession();
-        QualifiedObjectName tableName = createQualifiedObjectName(session, statement, statement.getName());
-        Optional<TableHandle> tableHandle = plannerContext.getMetadata().getTableHandle(session, tableName);
-        if (tableHandle.isEmpty()) {
+        QualifiedObjectName originalTableName = createQualifiedObjectName(session, statement, statement.getName());
+        RedirectionAwareTableHandle redirectionAwareTableHandle = plannerContext.getMetadata().getRedirectionAwareTableHandle(session, originalTableName);
+        if (redirectionAwareTableHandle.getTableHandle().isEmpty()) {
             if (!statement.isTableExists()) {
-                throw semanticException(TABLE_NOT_FOUND, statement, "Table '%s' does not exist", tableName);
+                throw semanticException(TABLE_NOT_FOUND, statement, "Table '%s' does not exist", originalTableName);
             }
             return immediateVoidFuture();
         }
+        TableHandle tableHandle = redirectionAwareTableHandle.getTableHandle().get();
+        CatalogName catalogName = getRequiredCatalogHandle(plannerContext.getMetadata(), session, statement, tableHandle.getCatalogName().getCatalogName());
 
-        CatalogName catalogName = getRequiredCatalogHandle(plannerContext.getMetadata(), session, statement, tableName.getCatalogName());
+        accessControl.checkCanAddColumns(session.toSecurityContext(), redirectionAwareTableHandle.getRedirectedTableName().orElse(originalTableName));
 
-        accessControl.checkCanAddColumns(session.toSecurityContext(), tableName);
-
-        Map<String, ColumnHandle> columnHandles = plannerContext.getMetadata().getColumnHandles(session, tableHandle.get());
+        Map<String, ColumnHandle> columnHandles = plannerContext.getMetadata().getColumnHandles(session, tableHandle);
 
         ColumnDefinition element = statement.getColumn();
         Type type;
@@ -133,7 +133,7 @@ public ListenableFuture<Void> execute(
                 .setProperties(columnProperties)
                 .build();
 
-        plannerContext.getMetadata().addColumn(session, tableHandle.get(), column);
+        plannerContext.getMetadata().addColumn(session, tableHandle, column);
 
         return immediateVoidFuture();
     }

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
@@ -321,10 +321,14 @@ public void testAlterTableAddColumn()
 
         createIcebergTable(icebergTableName, false);
 
-        //TODO restore test assertions after adding redirection awareness to the AddColumnTask
-        assertQueryFailure(() -> onTrino().executeQuery("ALTER TABLE " + hiveTableName + " ADD COLUMN some_new_column double"))
-                .hasMessageMatching("\\QQuery failed (#\\E\\S+\\Q): Cannot query Iceberg table 'default." + tableName + "'");
+        onTrino().executeQuery("ALTER TABLE " + hiveTableName + " ADD COLUMN some_new_column double");
+
+        Assertions.assertThat(onTrino().executeQuery("DESCRIBE " + icebergTableName).column(1))
+                .containsOnly("nationkey", "name", "regionkey", "comment", "some_new_column");
 
+        assertResultsEqual(
+                onTrino().executeQuery("TABLE " + icebergTableName),
+                onTrino().executeQuery("SELECT * , NULL FROM tpch.tiny.nation"));
         onTrino().executeQuery("DROP TABLE " + icebergTableName);
     }