From d94987723c5b7ca28746f72a1bc400ec6d5e3da7 Mon Sep 17 00:00:00 2001 From: Yuya Ebihara Date: Fri, 22 Nov 2024 13:45:47 +0900 Subject: [PATCH] Support FIRST and AFTER clause when adding a new column in Iceberg --- .../plugin/jdbc/BaseJdbcConnectorTest.java | 3 +- .../deltalake/TestDeltaLakeConnectorTest.java | 3 +- .../plugin/hive/BaseHiveConnectorTest.java | 3 +- .../trino/plugin/iceberg/IcebergMetadata.java | 14 +++++--- .../plugin/kudu/TestKuduConnectorTest.java | 3 +- .../memory/TestMemoryConnectorTest.java | 3 +- .../mongodb/TestMongoConnectorTest.java | 3 +- .../io/trino/testing/BaseConnectorTest.java | 35 +++++++++++++++++++ .../testing/TestingConnectorBehavior.java | 1 + 9 files changed, 58 insertions(+), 10 deletions(-) diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java index 1f0726e30d0d..85972148fad0 100644 --- a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java +++ b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java @@ -148,7 +148,8 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) { return switch (connectorBehavior) { case SUPPORTS_UPDATE -> true; - case SUPPORTS_CREATE_MATERIALIZED_VIEW, + case SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_CREATE_MATERIALIZED_VIEW, SUPPORTS_CREATE_VIEW, SUPPORTS_MERGE, SUPPORTS_PREDICATE_EXPRESSION_PUSHDOWN, diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeConnectorTest.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeConnectorTest.java index c16c2c1c641c..d6f1910528ee 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeConnectorTest.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeConnectorTest.java @@ -171,7 +171,8 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) return switch (connectorBehavior) { case SUPPORTS_CREATE_OR_REPLACE_TABLE, SUPPORTS_REPORTING_WRITTEN_BYTES -> true; - case SUPPORTS_ADD_FIELD, + case SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_ADD_FIELD, SUPPORTS_AGGREGATION_PUSHDOWN, SUPPORTS_CREATE_MATERIALIZED_VIEW, SUPPORTS_DROP_FIELD, diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java index a8ae750bf204..dd2fb7e50da8 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java @@ -250,7 +250,8 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) return switch (connectorBehavior) { case SUPPORTS_MULTI_STATEMENT_WRITES, SUPPORTS_REPORTING_WRITTEN_BYTES -> true; // FIXME: Fails because only allowed with transactional tables - case SUPPORTS_ADD_FIELD, + case SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_ADD_FIELD, SUPPORTS_CREATE_MATERIALIZED_VIEW, SUPPORTS_DROP_FIELD, SUPPORTS_MERGE, diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java index 24433af32f66..ba91d180eae8 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java @@ -68,6 +68,7 @@ import io.trino.spi.connector.CatalogSchemaTableName; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ColumnPosition; import io.trino.spi.connector.ConnectorAccessControl; import io.trino.spi.connector.ConnectorAnalyzeMetadata; import io.trino.spi.connector.ConnectorInsertTableHandle; @@ -2442,7 +2443,7 @@ private static Term toIcebergTerm(Schema schema, PartitionField partitionField) } @Override - public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column) + public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column, ColumnPosition position) { // Spark doesn't support adding a NOT NULL column to Iceberg tables // Also, Spark throws an exception when reading the table if we add such columns and execute a rollback procedure @@ -2456,9 +2457,14 @@ public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle // added - instead of relying on addColumn in iceberg library to assign Ids AtomicInteger nextFieldId = new AtomicInteger(icebergTable.schema().highestFieldId() + 2); try { - icebergTable.updateSchema() - .addColumn(column.getName(), toIcebergTypeForNewColumn(column.getType(), nextFieldId), column.getComment()) - .commit(); + UpdateSchema updateSchema = icebergTable.updateSchema(); + updateSchema.addColumn(column.getName(), toIcebergTypeForNewColumn(column.getType(), nextFieldId), column.getComment()); + switch (position) { + case ColumnPosition.First _ -> updateSchema.moveFirst(column.getName()); + case ColumnPosition.After after -> updateSchema.moveAfter(column.getName(), after.columnName()); + case ColumnPosition.Last _ -> {} + } + updateSchema.commit(); } catch (RuntimeException e) { throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to add column: " + firstNonNull(e.getMessage(), e), e); diff --git a/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java b/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java index 8d449e8523e4..7b1f0b03b8b1 100644 --- a/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java +++ b/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java @@ -59,7 +59,8 @@ protected QueryRunner createQueryRunner() protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) { return switch (connectorBehavior) { - case SUPPORTS_ARRAY, + case SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_ARRAY, SUPPORTS_COMMENT_ON_COLUMN, SUPPORTS_COMMENT_ON_TABLE, SUPPORTS_CREATE_MATERIALIZED_VIEW, diff --git a/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java b/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java index f1d71b812e1e..e7f7e4bb97c3 100644 --- a/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java +++ b/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java @@ -88,7 +88,8 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) { return switch (connectorBehavior) { case SUPPORTS_TRUNCATE -> true; - case SUPPORTS_ADD_FIELD, + case SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_ADD_FIELD, SUPPORTS_AGGREGATION_PUSHDOWN, SUPPORTS_CREATE_MATERIALIZED_VIEW, SUPPORTS_DELETE, diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java index 512da391977c..8131a8c28445 100644 --- a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java +++ b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java @@ -99,7 +99,8 @@ public final void destroy() protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) { return switch (connectorBehavior) { - case SUPPORTS_ADD_FIELD, + case SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_ADD_FIELD, SUPPORTS_CREATE_MATERIALIZED_VIEW, SUPPORTS_CREATE_VIEW, SUPPORTS_DROP_FIELD, diff --git a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java index 59dbfe13f5fe..40673ce3d401 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java @@ -105,6 +105,7 @@ import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ADD_COLUMN; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ADD_COLUMN_NOT_NULL_CONSTRAINT; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ADD_COLUMN_WITH_COMMENT; +import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ADD_COLUMN_WITH_POSITION; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ADD_FIELD; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ADD_FIELD_IN_ARRAY; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ARRAY; @@ -2412,6 +2413,40 @@ protected void verifyAddNotNullColumnToNonEmptyTableFailurePermissible(Throwable throw new AssertionError("Unexpected failure when adding not null column", e); } + @Test + public void testAddColumnWithPosition() + { + skipTestUnless(hasBehavior(SUPPORTS_ADD_COLUMN)); // covered by testAddColumn + + if (!hasBehavior(SUPPORTS_ADD_COLUMN_WITH_POSITION)) { + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_add_column_", "AS SELECT 2 second, 4 fourth")) { + assertQueryFails( + "ALTER TABLE " + table.getName() + " ADD COLUMN first integer FIRST", + "This connector does not support adding columns with FIRST clause"); + assertQueryFails( + "ALTER TABLE " + table.getName() + " ADD COLUMN third integer AFTER second", + "This connector does not support adding columns with AFTER clause"); + } + return; + } + + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_add_column_", "AS SELECT 2 second, 4 fourth")) { + assertTableColumnNames(table.getName(), "second", "fourth"); + assertQuery("SELECT * FROM " + table.getName(), "VALUES (2, 4)"); + + assertUpdate("ALTER TABLE " + table.getName() + " ADD COLUMN first integer FIRST"); + assertTableColumnNames(table.getName(), "first", "second", "fourth"); + assertQuery("SELECT * FROM " + table.getName(), "VALUES (null, 2, 4)"); + + assertUpdate("ALTER TABLE " + table.getName() + " ADD COLUMN third integer AFTER second"); + assertTableColumnNames(table.getName(), "first", "second", "third", "fourth"); + assertQuery("SELECT * FROM " + table.getName(), "VALUES (null, 2, null, 4)"); + + assertUpdate("INSERT INTO " + table.getName() + " VALUES (10, 20, 30, 40)", 1); + assertQuery("SELECT * FROM " + table.getName(), "VALUES (null, 2, null, 4), (10, 20, 30, 40)"); + } + } + @Test public void testAddRowField() { diff --git a/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java b/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java index 1245d75ca2a2..fcab562f9762 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java @@ -85,6 +85,7 @@ public enum TestingConnectorBehavior SUPPORTS_ADD_COLUMN, SUPPORTS_ADD_COLUMN_WITH_COMMENT(SUPPORTS_ADD_COLUMN), + SUPPORTS_ADD_COLUMN_WITH_POSITION(SUPPORTS_ADD_COLUMN), SUPPORTS_ADD_FIELD(fallback -> fallback.test(SUPPORTS_ADD_COLUMN) && fallback.test(SUPPORTS_ROW_TYPE)), SUPPORTS_ADD_FIELD_IN_ARRAY(SUPPORTS_ADD_FIELD), SUPPORTS_DROP_COLUMN(SUPPORTS_ADD_COLUMN),