Support renaming a field with RENAME COLUMN in engine

trinodb · Jul 24, 2023 · 3004e83 · 3004e83
1 parent 7a5f728
commit 3004e83
Show file tree

Hide file tree

Showing 15 changed files with 316 additions and 25 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/execution/RenameColumnTask.java b/core/trino-main/src/main/java/io/trino/execution/RenameColumnTask.java
@@ -22,15 +22,23 @@
 import io.trino.metadata.RedirectionAwareTableHandle;
 import io.trino.metadata.TableHandle;
 import io.trino.security.AccessControl;
+import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ColumnHandle;
+import io.trino.spi.connector.ColumnMetadata;
+import io.trino.spi.type.RowType;
+import io.trino.spi.type.Type;
 import io.trino.sql.tree.Expression;
 import io.trino.sql.tree.RenameColumn;
 
 import java.util.List;
 import java.util.Map;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.Iterables.getLast;
+import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
 import static io.trino.metadata.MetadataUtil.createQualifiedObjectName;
+import static io.trino.spi.StandardErrorCode.AMBIGUOUS_NAME;
 import static io.trino.spi.StandardErrorCode.COLUMN_ALREADY_EXISTS;
 import static io.trino.spi.StandardErrorCode.COLUMN_NOT_FOUND;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
@@ -76,11 +84,10 @@ public ListenableFuture<Void> execute(
         }
         TableHandle tableHandle = redirectionAwareTableHandle.getTableHandle().get();
 
-        String source = statement.getSource().getValue().toLowerCase(ENGLISH);
+        String source = statement.getSource().getParts().get(0).toLowerCase(ENGLISH);
         String target = statement.getTarget().getValue().toLowerCase(ENGLISH);
 
         QualifiedObjectName qualifiedTableName = redirectionAwareTableHandle.getRedirectedTableName().orElse(originalTableName);
-        accessControl.checkCanRenameColumn(session.toSecurityContext(), qualifiedTableName);
 
         Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle);
         ColumnHandle columnHandle = columnHandles.get(source);
@@ -90,17 +97,72 @@ public ListenableFuture<Void> execute(
             }
             return immediateVoidFuture();
         }
-
-        if (columnHandles.containsKey(target)) {
-            throw semanticException(COLUMN_ALREADY_EXISTS, statement, "Column '%s' already exists", target);
-        }
-
         if (metadata.getColumnMetadata(session, tableHandle, columnHandle).isHidden()) {
             throw semanticException(NOT_SUPPORTED, statement, "Cannot rename hidden column");
         }
 
-        metadata.renameColumn(session, tableHandle, qualifiedTableName.asCatalogSchemaTableName(), columnHandle, target);
+        if (statement.getSource().getParts().size() == 1) {
+            accessControl.checkCanRenameColumn(session.toSecurityContext(), qualifiedTableName);
+
+            if (columnHandles.containsKey(target)) {
+                throw semanticException(COLUMN_ALREADY_EXISTS, statement, "Column '%s' already exists", target);
+            }
+
+            metadata.renameColumn(session, tableHandle, qualifiedTableName.asCatalogSchemaTableName(), columnHandle, target);
+        }
+        else {
+            accessControl.checkCanAlterColumn(session.toSecurityContext(), qualifiedTableName);
+
+            List<String> fieldPath = statement.getSource().getParts();
+
+            ColumnMetadata columnMetadata = metadata.getColumnMetadata(session, tableHandle, columnHandle);
+            Type currentType = columnMetadata.getType();
+            for (int i = 1; i < fieldPath.size() - 1; i++) {
+                String fieldName = fieldPath.get(i);
+                List<RowType.Field> candidates = getCandidates(currentType, fieldName);
+
+                if (candidates.isEmpty()) {
+                    throw semanticException(COLUMN_NOT_FOUND, statement, "Field '%s' does not exist within %s", fieldName, currentType);
+                }
+                if (candidates.size() > 1) {
+                    throw semanticException(AMBIGUOUS_NAME, statement, "Field path %s within %s is ambiguous", fieldPath, columnMetadata.getType());
+                }
+                currentType = getOnlyElement(candidates).getType();
+            }
+
+            String sourceFieldName = getLast(statement.getSource().getParts());
+            List<RowType.Field> sourceCandidates = getCandidates(currentType, sourceFieldName);
+            if (sourceCandidates.isEmpty()) {
+                if (!statement.isColumnExists()) {
+                    throw semanticException(COLUMN_NOT_FOUND, statement, "Field '%s' does not exist", source);
+                }
+                return immediateVoidFuture();
+            }
+            if (sourceCandidates.size() > 1) {
+                throw semanticException(AMBIGUOUS_NAME, statement, "Field path %s within %s is ambiguous", fieldPath, columnMetadata.getType());
+            }
+
+            List<RowType.Field> targetCandidates = getCandidates(currentType, target);
+            if (!targetCandidates.isEmpty()) {
+                throw semanticException(COLUMN_ALREADY_EXISTS, statement, "Field '%s' already exists", target);
+            }
+
+            metadata.renameField(session, tableHandle, fieldPath, target);
+        }
 
         return immediateVoidFuture();
     }
+
+    private static List<RowType.Field> getCandidates(Type type, String fieldName)
+    {
+        if (!(type instanceof RowType rowType)) {
+            throw new TrinoException(NOT_SUPPORTED, "Unsupported type: " + type);
+        }
+        List<RowType.Field> candidates = rowType.getFields().stream()
+                // case-insensitive match
+                .filter(rowField -> rowField.getName().isPresent() && rowField.getName().get().equalsIgnoreCase(fieldName))
+                .collect(toImmutableList());
+
+        return candidates;
+    }
 }
diff --git a/core/trino-main/src/main/java/io/trino/metadata/Metadata.java b/core/trino-main/src/main/java/io/trino/metadata/Metadata.java
@@ -237,6 +237,11 @@ Optional<TableExecuteHandle> getTableHandleForExecute(
      */
     void renameColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnHandle source, String target);
 
+    /**
+     * Rename the specified field.
+     */
+    void renameField(Session session, TableHandle tableHandle, List<String> fieldPath, String target);
+
     /**
      * Add the specified column to the table.
      */

diff --git a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java
@@ -757,6 +757,14 @@ public void renameColumn(Session session, TableHandle tableHandle, CatalogSchema
         }
     }
 
+    @Override
+    public void renameField(Session session, TableHandle tableHandle, List<String> fieldPath, String target)
+    {
+        CatalogHandle catalogHandle = tableHandle.getCatalogHandle();
+        ConnectorMetadata metadata = getMetadataForWrite(session, catalogHandle);
+        metadata.renameField(session.toConnectorSession(catalogHandle), tableHandle.getConnectorHandle(), fieldPath, target.toLowerCase(ENGLISH));
+    }
+
     @Override
     public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column)
     {

diff --git a/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java b/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java
@@ -477,6 +477,15 @@ public void renameColumn(ConnectorSession session, ConnectorTableHandle tableHan
         }
     }
 
+    @Override
+    public void renameField(ConnectorSession session, ConnectorTableHandle tableHandle, List<String> fieldPath, String target)
+    {
+        Span span = startSpan("renameField", tableHandle);
+        try (var ignored = scopedSpan(span)) {
+            delegate.renameField(session, tableHandle, fieldPath, target);
+        }
+    }
+
     @Override
     public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column)
     {

diff --git a/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java b/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java
@@ -440,6 +440,15 @@ public void renameColumn(Session session, TableHandle tableHandle, CatalogSchema
         }
     }
 
+    @Override
+    public void renameField(Session session, TableHandle tableHandle, List<String> fieldPath, String target)
+    {
+        Span span = startSpan("renameField", tableHandle);
+        try (var ignored = scopedSpan(span)) {
+            delegate.renameField(session, tableHandle, fieldPath, target);
+        }
+    }
+
     @Override
     public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column)
     {

diff --git a/core/trino-main/src/test/java/io/trino/connector/MockConnector.java b/core/trino-main/src/test/java/io/trino/connector/MockConnector.java
@@ -602,6 +602,12 @@ public void setTableAuthorization(ConnectorSession session, SchemaTableName tabl
         @Override
         public void renameColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle source, String target) {}
 
+        @Override
+        public void renameField(ConnectorSession session, ConnectorTableHandle tableHandle, List<String> fieldPath, String target)
+        {
+            throw new UnsupportedOperationException();
+        }
+
         @Override
         public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column) {}