Skip to content

Commit

Permalink
Support FIRST and AFTER clause when adding a new column in engine
Browse files Browse the repository at this point in the history
  • Loading branch information
@ebyhr
ebyhr committed Apr 21, 2024
1 parent 8e2d236 commit 1542fde
Show file tree
Hide file tree
Showing 19 changed files with 375 additions and 49 deletions.
3 changes: 2 additions & 1 deletion core/trino-grammar/src/main/antlr4/io/trino/grammar/sql/SqlBase.g4
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,8 @@ statement
| ALTER TABLE (IF EXISTS)? from=qualifiedName
RENAME TO to=qualifiedName #renameTable
| ALTER TABLE (IF EXISTS)? tableName=qualifiedName
ADD COLUMN (IF NOT EXISTS)? column=columnDefinition #addColumn
ADD COLUMN (IF NOT EXISTS)? column=columnDefinition
(FIRST | AFTER after=identifier)? #addColumn
| ALTER TABLE (IF EXISTS)? tableName=qualifiedName
RENAME COLUMN (IF EXISTS)? from=qualifiedName TO to=identifier #renameColumn
| ALTER TABLE (IF EXISTS)? tableName=qualifiedName
Expand Down
41 changes: 39 additions & 2 deletions core/trino-main/src/main/java/io/trino/execution/AddColumnTask.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,15 @@
import io.trino.sql.PlannerContext;
import io.trino.sql.tree.AddColumn;
import io.trino.sql.tree.ColumnDefinition;
import io.trino.sql.tree.ColumnPosition;
import io.trino.sql.tree.Expression;
import io.trino.sql.tree.Identifier;

import java.util.List;
import java.util.Map;
import java.util.Optional;

import static com.google.common.base.Verify.verify;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.Iterables.getLast;
import static com.google.common.collect.Iterables.getOnlyElement;
Expand All @@ -54,6 +57,7 @@
import static io.trino.spi.connector.ConnectorCapabilities.NOT_NULL_COLUMN_CONSTRAINT;
import static io.trino.sql.analyzer.SemanticExceptions.semanticException;
import static io.trino.sql.analyzer.TypeSignatureTranslator.toTypeSignature;
import static io.trino.sql.tree.ColumnPosition.AFTER;
import static io.trino.type.UnknownType.UNKNOWN;
import static java.util.Locale.ENGLISH;
import static java.util.Objects.requireNonNull;
Expand Down Expand Up @@ -104,13 +108,18 @@ public ListenableFuture<Void> execute(

ColumnDefinition element = statement.getColumn();
Identifier columnName = element.getName().getOriginalParts().get(0);
ColumnPosition position = statement.getPosition();
Optional<Identifier> afterColumnName = statement.getAfter();
Type type;
try {
type = plannerContext.getTypeManager().getType(toTypeSignature(element.getType()));
}
catch (TypeNotFoundException e) {
throw semanticException(TYPE_NOT_FOUND, element, "Unknown type '%s' for column '%s'", element.getType(), columnName);
}
if (afterColumnName.isPresent()) {
verify(position == AFTER, "afterColumnName should be empty when position is not AFTER");
}

if (element.getName().getParts().size() == 1) {
accessControl.checkCanAddColumns(session.toSecurityContext(), qualifiedTableName);
Expand All @@ -127,6 +136,9 @@ public ListenableFuture<Void> execute(
if (!element.isNullable() && !plannerContext.getMetadata().getConnectorCapabilities(session, catalogHandle).contains(NOT_NULL_COLUMN_CONSTRAINT)) {
throw semanticException(NOT_SUPPORTED, element, "Catalog '%s' does not support NOT NULL for column '%s'", catalogHandle, columnName);
}
if (position == AFTER && !columnHandles.containsKey(afterColumnName.orElseThrow().getValue().toLowerCase(ENGLISH))) {
throw semanticException(COLUMN_NOT_FOUND, statement, "Column '%s' does not", afterColumnName.orElseThrow().getValue());
}

Map<String, Object> columnProperties = columnPropertyManager.getProperties(
catalogHandle.getCatalogName().toString(),
Expand All @@ -146,7 +158,13 @@ public ListenableFuture<Void> execute(
.setProperties(columnProperties)
.build();

plannerContext.getMetadata().addColumn(session, tableHandle, qualifiedTableName.asCatalogSchemaTableName(), column);
plannerContext.getMetadata().addColumn(
session,
tableHandle,
qualifiedTableName.asCatalogSchemaTableName(),
column,
toConnectorColumnPosition(position),
afterColumnName.map(identifier -> identifier.getValue().toLowerCase(ENGLISH)));
}
else {
accessControl.checkCanAlterColumn(session.toSecurityContext(), qualifiedTableName);
Expand Down Expand Up @@ -186,7 +204,17 @@ public ListenableFuture<Void> execute(
}
throw semanticException(COLUMN_ALREADY_EXISTS, statement, "Field '%s' already exists", fieldName);
}
plannerContext.getMetadata().addField(session, tableHandle, parentPath, fieldName, type, statement.isColumnNotExists());
if (position == AFTER && getCandidates(currentType, afterColumnName.orElseThrow().getValue()).isEmpty()) {
throw semanticException(COLUMN_NOT_FOUND, statement, "Field '%s' does not exist within %s", afterColumnName.orElseThrow().getValue().toLowerCase(ENGLISH), currentType);
}
plannerContext.getMetadata().addField(
session,
tableHandle,
parentPath,
fieldName,
type,
toConnectorColumnPosition(position),
afterColumnName.map(identifier -> identifier.getValue().toLowerCase(ENGLISH)), statement.isColumnNotExists());
}

return immediateVoidFuture();
Expand All @@ -204,4 +232,13 @@ private static List<RowType.Field> getCandidates(Type type, String fieldName)

return candidates;
}

private static io.trino.spi.connector.ColumnPosition toConnectorColumnPosition(ColumnPosition columnPosition)
{
return switch (columnPosition) {
case FIRST -> io.trino.spi.connector.ColumnPosition.FIRST;
case AFTER -> io.trino.spi.connector.ColumnPosition.AFTER;
case LAST -> io.trino.spi.connector.ColumnPosition.LAST;
};
}
}
5 changes: 3 additions & 2 deletions core/trino-main/src/main/java/io/trino/metadata/Metadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,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.ConnectorCapabilities;
import io.trino.spi.connector.ConnectorOutputMetadata;
import io.trino.spi.connector.ConnectorTableMetadata;
Expand Down Expand Up @@ -270,12 +271,12 @@ Optional<TableExecuteHandle> getTableHandleForExecute(
/**
* Add the specified column to the table.
*/
void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column);
void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column, ColumnPosition position, Optional<String> afterColumnName);

/**
* Add the specified field to the column.
*/
void addField(Session session, TableHandle tableHandle, List<String> parentPath, String fieldName, Type type, boolean ignoreExisting);
void addField(Session session, TableHandle tableHandle, List<String> parentPath, String fieldName, Type type, ColumnPosition position, Optional<String> afterFieldName, boolean ignoreExisting);

/**
* Set the specified type to the column.
Expand Down
9 changes: 5 additions & 4 deletions core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,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.ConnectorAnalyzeMetadata;
import io.trino.spi.connector.ConnectorCapabilities;
import io.trino.spi.connector.ConnectorInsertTableHandle;
Expand Down Expand Up @@ -902,23 +903,23 @@ public void renameField(Session session, TableHandle tableHandle, List<String> f
}

@Override
public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column)
public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column, ColumnPosition position, Optional<String> afterColumnName)
{
CatalogHandle catalogHandle = tableHandle.catalogHandle();
CatalogMetadata catalogMetadata = getCatalogMetadataForWrite(session, catalogHandle.getCatalogName().toString());
ConnectorMetadata metadata = getMetadataForWrite(session, catalogHandle);
metadata.addColumn(session.toConnectorSession(catalogHandle), tableHandle.connectorHandle(), column);
metadata.addColumn(session.toConnectorSession(catalogHandle), tableHandle.connectorHandle(), column, position, afterColumnName);
if (catalogMetadata.getSecurityManagement() == SYSTEM) {
systemSecurityMetadata.columnCreated(session, table, column.getName());
}
}

@Override
public void addField(Session session, TableHandle tableHandle, List<String> parentPath, String fieldName, Type type, boolean ignoreExisting)
public void addField(Session session, TableHandle tableHandle, List<String> parentPath, String fieldName, Type type, ColumnPosition position, Optional<String> afterFieldName, boolean ignoreExisting)
{
CatalogHandle catalogHandle = tableHandle.catalogHandle();
ConnectorMetadata metadata = getMetadataForWrite(session, catalogHandle);
metadata.addField(session.toConnectorSession(catalogHandle), tableHandle.connectorHandle(), parentPath, fieldName, type, ignoreExisting);
metadata.addField(session.toConnectorSession(catalogHandle), tableHandle.connectorHandle(), parentPath, fieldName, type, position, afterFieldName, ignoreExisting);
}

@Override
Expand Down
19 changes: 19 additions & 0 deletions core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,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.ConnectorAnalyzeMetadata;
import io.trino.spi.connector.ConnectorInsertTableHandle;
import io.trino.spi.connector.ConnectorMaterializedViewDefinition;
Expand Down Expand Up @@ -485,6 +486,15 @@ public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle
}
}

@Override
public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column, ColumnPosition position, Optional<String> afterColumnName)
{
Span span = startSpan("addColumn", tableHandle);
try (var ignored = scopedSpan(span)) {
delegate.addColumn(session, tableHandle, column, position, afterColumnName);
}
}

@Override
public void addField(ConnectorSession session, ConnectorTableHandle tableHandle, List<String> parentPath, String fieldName, Type type, boolean ignoreExisting)
{
Expand All @@ -494,6 +504,15 @@ public void addField(ConnectorSession session, ConnectorTableHandle tableHandle,
}
}

@Override
public void addField(ConnectorSession session, ConnectorTableHandle tableHandle, List<String> parentPath, String fieldName, Type type, ColumnPosition position, Optional<String> afterFieldName, boolean ignoreExisting)
{
Span span = startSpan("addField", tableHandle);
try (var ignored = scopedSpan(span)) {
delegate.addField(session, tableHandle, parentPath, fieldName, type, position, afterFieldName, ignoreExisting);
}
}

@Override
public void setColumnType(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column, Type type)
{
Expand Down
9 changes: 5 additions & 4 deletions core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,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.ConnectorCapabilities;
import io.trino.spi.connector.ConnectorOutputMetadata;
import io.trino.spi.connector.ConnectorTableMetadata;
Expand Down Expand Up @@ -483,20 +484,20 @@ public void renameField(Session session, TableHandle tableHandle, List<String> f
}

@Override
public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column)
public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column, ColumnPosition position, Optional<String> afterColumnName)
{
Span span = startSpan("addColumn", table);
try (var ignored = scopedSpan(span)) {
delegate.addColumn(session, tableHandle, table, column);
delegate.addColumn(session, tableHandle, table, column, position, afterColumnName);
}
}

@Override
public void addField(Session session, TableHandle tableHandle, List<String> parentPath, String fieldName, Type type, boolean ignoreExisting)
public void addField(Session session, TableHandle tableHandle, List<String> parentPath, String fieldName, Type type, ColumnPosition position, Optional<String> afterFieldName, boolean ignoreExisting)
{
Span span = startSpan("addField", tableHandle);
try (var ignored = scopedSpan(span)) {
delegate.addField(session, tableHandle, parentPath, fieldName, type, ignoreExisting);
delegate.addField(session, tableHandle, parentPath, fieldName, type, position, afterFieldName, ignoreExisting);
}
}

Expand Down
23 changes: 20 additions & 3 deletions core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,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.ConnectorTableMetadata;
import io.trino.spi.connector.SaveMode;
import io.trino.spi.connector.SchemaTableName;
Expand Down Expand Up @@ -368,14 +369,30 @@ public void renameTable(Session session, TableHandle tableHandle, CatalogSchemaT
}

@Override
public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column)
public void addColumn(Session session, TableHandle tableHandle, CatalogSchemaTableName table, ColumnMetadata column, ColumnPosition position, Optional<String> afterColumnName)
{
SchemaTableName tableName = table.getSchemaTableName();
ConnectorTableMetadata metadata = tables.get(tableName);

ImmutableList.Builder<ColumnMetadata> columns = ImmutableList.builderWithExpectedSize(metadata.getColumns().size() + 1);
columns.addAll(metadata.getColumns());
columns.add(column);
switch (position) {
case FIRST -> {
columns.add(column);
columns.addAll(metadata.getColumns());
}
case AFTER -> {
for (ColumnMetadata existingColumn : metadata.getColumns()) {
columns.add(existingColumn);
if (existingColumn.getName().equals(afterColumnName.orElseThrow())) {
columns.add(column);
}
}
}
case LAST -> {
columns.addAll(metadata.getColumns());
columns.add(column);
}
}
tables.put(tableName, new ConnectorTableMetadata(tableName, columns.build()));
}

Expand Down
Loading

0 comments on commit 1542fde

Please sign in to comment.