From 9253930994b826b9100799a82f939109212ff85a Mon Sep 17 00:00:00 2001 From: shandeep Date: Sun, 30 Jan 2022 19:18:23 +0530 Subject: [PATCH] Add comments at column level for views --- .../java/io/trino/execution/CommentTask.java | 39 +++++++++++----- .../main/java/io/trino/metadata/Metadata.java | 5 +++ .../io/trino/metadata/MetadataManager.java | 10 +++++ .../java/io/trino/metadata/ViewColumn.java | 24 ++++++++-- .../io/trino/metadata/ViewDefinition.java | 2 +- .../trino/metadata/AbstractMockMetadata.java | 6 +++ .../spi/connector/ConnectorMetadata.java | 8 ++++ .../connector/ConnectorViewDefinition.java | 18 +++++++- .../ClassLoaderSafeConnectorMetadata.java | 8 ++++ .../io/trino/plugin/hive/HiveMetadata.java | 44 +++++++++++++++++++ 10 files changed, 148 insertions(+), 16 deletions(-) diff --git a/core/trino-main/src/main/java/io/trino/execution/CommentTask.java b/core/trino-main/src/main/java/io/trino/execution/CommentTask.java index 1e482f350b91..eecc0de56b2d 100644 --- a/core/trino-main/src/main/java/io/trino/execution/CommentTask.java +++ b/core/trino-main/src/main/java/io/trino/execution/CommentTask.java @@ -19,6 +19,8 @@ import io.trino.metadata.Metadata; import io.trino.metadata.QualifiedObjectName; import io.trino.metadata.TableHandle; +import io.trino.metadata.ViewColumn; +import io.trino.metadata.ViewDefinition; import io.trino.security.AccessControl; import io.trino.spi.connector.ColumnHandle; import io.trino.sql.tree.Comment; @@ -86,20 +88,37 @@ else if (statement.getType() == Comment.Type.COLUMN) { } QualifiedObjectName tableName = createQualifiedObjectName(session, statement, prefix.get()); - Optional tableHandle = metadata.getTableHandle(session, tableName); - if (tableHandle.isEmpty()) { - throw semanticException(TABLE_NOT_FOUND, statement, "Table does not exist: " + tableName); - } - String columnName = statement.getName().getSuffix(); - Map columnHandles = metadata.getColumnHandles(session, tableHandle.get()); - if (!columnHandles.containsKey(columnName)) { - throw semanticException(COLUMN_NOT_FOUND, statement, "Column does not exist: " + columnName); + Optional optionalView = metadata.getView(session, tableName); + if (optionalView.isPresent()) { + ViewDefinition viewDefinition = optionalView.get(); + List viewColumns = viewDefinition.getColumns(); + Optional optionalViewColumn = viewColumns + .stream() + .filter(viewColumn -> viewColumn.getName().equals(columnName)) + .findFirst(); + + if (optionalViewColumn.isEmpty()) { + throw semanticException(COLUMN_NOT_FOUND, statement, "Column does not exist: " + columnName); + } + + accessControl.checkCanSetColumnComment(session.toSecurityContext(), tableName); } + else { + Optional tableHandle = metadata.getTableHandle(session, tableName); + if (tableHandle.isEmpty()) { + throw semanticException(TABLE_NOT_FOUND, statement, "Table does not exist: " + tableName); + } + + Map columnHandles = metadata.getColumnHandles(session, tableHandle.get()); + if (!columnHandles.containsKey(columnName)) { + throw semanticException(COLUMN_NOT_FOUND, statement, "Column does not exist: " + columnName); + } - accessControl.checkCanSetColumnComment(session.toSecurityContext(), tableName); + accessControl.checkCanSetColumnComment(session.toSecurityContext(), tableName); - metadata.setColumnComment(session, tableHandle.get(), columnHandles.get(columnName), statement.getComment()); + metadata.setColumnComment(session, tableHandle.get(), columnHandles.get(columnName), statement.getComment()); + } } else { throw semanticException(NOT_SUPPORTED, statement, "Unsupported comment type: %s", statement.getType()); 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 index 79d3be582414..2aa3e684d354 100644 --- a/core/trino-main/src/main/java/io/trino/metadata/Metadata.java +++ b/core/trino-main/src/main/java/io/trino/metadata/Metadata.java @@ -713,4 +713,9 @@ default boolean isMaterializedView(Session session, QualifiedObjectName viewName * Returns a table handle for the specified table name with a specified version */ Optional getTableHandle(Session session, QualifiedObjectName tableName, Optional startVersion, Optional endVersion); + + /** + * Comments to the specified view column + */ + void setViewColumnComment(Session session, QualifiedObjectName viewName, ViewDefinition viewDefinition, ViewColumn viewColumn, Optional comment); } 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 index be7c01bb9d5e..9e300194e673 100644 --- a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java +++ b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java @@ -2539,6 +2539,16 @@ private Optional toConnectorVersion(Optional comment) + { + CatalogMetadata catalogMetadata = getCatalogMetadataForWrite(session, viewName.getCatalogName()); + CatalogName catalogName = catalogMetadata.getCatalogName(); + ConnectorMetadata metadata = catalogMetadata.getMetadata(); + + metadata.setViewColumnComment(session.toConnectorSession(catalogName), viewName.asSchemaTableName(), viewDefinition.toConnectorViewDefinition(), viewColumn.getName(), comment); + } + private static class OperatorCacheKey { private final OperatorType operatorType; diff --git a/core/trino-main/src/main/java/io/trino/metadata/ViewColumn.java b/core/trino-main/src/main/java/io/trino/metadata/ViewColumn.java index 7dc2215c3b12..460dc2edf3b9 100644 --- a/core/trino-main/src/main/java/io/trino/metadata/ViewColumn.java +++ b/core/trino-main/src/main/java/io/trino/metadata/ViewColumn.java @@ -16,18 +16,27 @@ import io.trino.spi.type.TypeId; import java.util.Objects; +import java.util.Optional; +import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; public final class ViewColumn { private final String name; private final TypeId type; + private final Optional comment; public ViewColumn(String name, TypeId type) + { + this(name, type, Optional.empty()); + } + + public ViewColumn(String name, TypeId type, Optional comment) { this.name = requireNonNull(name, "name is null"); this.type = requireNonNull(type, "type is null"); + this.comment = comment; } public String getName() @@ -40,10 +49,19 @@ public TypeId getType() return type; } + public Optional getComment() + { + return comment; + } + @Override public String toString() { - return name + " " + type; + return toStringHelper(this).omitNullValues() + .add("name", name) + .add("type", type) + .add("comment", comment.orElse(null)) + .toString(); } @Override @@ -56,12 +74,12 @@ public boolean equals(Object o) return false; } ViewColumn that = (ViewColumn) o; - return Objects.equals(name, that.name) && Objects.equals(type, that.type); + return Objects.equals(name, that.name) && Objects.equals(type, that.type) && Objects.equals(comment, that.comment); } @Override public int hashCode() { - return Objects.hash(name, type); + return Objects.hash(name, type, comment); } } diff --git a/core/trino-main/src/main/java/io/trino/metadata/ViewDefinition.java b/core/trino-main/src/main/java/io/trino/metadata/ViewDefinition.java index c30914f58ad9..4899ec89b552 100644 --- a/core/trino-main/src/main/java/io/trino/metadata/ViewDefinition.java +++ b/core/trino-main/src/main/java/io/trino/metadata/ViewDefinition.java @@ -124,7 +124,7 @@ public ConnectorViewDefinition toConnectorViewDefinition() catalog, schema, columns.stream() - .map(column -> new ConnectorViewDefinition.ViewColumn(column.getName(), column.getType())) + .map(column -> new ConnectorViewDefinition.ViewColumn(column.getName(), column.getType(), column.getComment())) .collect(toImmutableList()), comment, runAsIdentity.map(Identity::getUser), diff --git a/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java b/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java index d9ccf1a7b64f..93bc59ee3f0a 100644 --- a/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java +++ b/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java @@ -886,4 +886,10 @@ public Optional getTableHandle(Session session, QualifiedObjectName { throw new UnsupportedOperationException(); } + + @Override + public void setViewColumnComment(Session session, QualifiedObjectName viewName, ViewDefinition viewDefinition, ViewColumn viewColumn, Optional comment) + { + throw new UnsupportedOperationException(); + } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java index a6528290442c..604040e10026 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java +++ b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java @@ -1340,4 +1340,12 @@ default boolean isSupportedVersionType(ConnectorSession session, SchemaTableName { throw new TrinoException(NOT_SUPPORTED, "This connector does not support versioned tables"); } + + /** + * Comments to the specified view column + */ + default void setViewColumnComment(ConnectorSession session, SchemaTableName viewName, ConnectorViewDefinition viewDefinition, String columnName, Optional comment) + { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting view column comments"); + } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorViewDefinition.java b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorViewDefinition.java index db5e1a20b012..aba7935f96c7 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorViewDefinition.java +++ b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorViewDefinition.java @@ -133,14 +133,22 @@ public static final class ViewColumn { private final String name; private final TypeId type; + private final Optional comment; + + public ViewColumn(String name, TypeId type) + { + this(name, type, Optional.empty()); + } @JsonCreator public ViewColumn( @JsonProperty("name") String name, - @JsonProperty("type") TypeId type) + @JsonProperty("type") TypeId type, + @JsonProperty("comment") Optional comment) { this.name = requireNonNull(name, "name is null"); this.type = requireNonNull(type, "type is null"); + this.comment = comment; } @JsonProperty @@ -155,10 +163,16 @@ public TypeId getType() return type; } + @JsonProperty + public Optional getComment() + { + return comment; + } + @Override public String toString() { - return name + " " + type; + return name + " " + type + " " + comment.orElse(null); } } } diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java index 18bd38a0e9fe..d7ba4bc035ab 100644 --- a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java +++ b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java @@ -1038,4 +1038,12 @@ public boolean isSupportedVersionType(ConnectorSession session, SchemaTableName return delegate.isSupportedVersionType(session, tableName, pointerType, versioning); } } + + @Override + public void setViewColumnComment(ConnectorSession session, SchemaTableName viewName, ConnectorViewDefinition viewDefinition, String columnName, Optional comment) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + delegate.setViewColumnComment(session, viewName, viewDefinition, columnName, comment); + } + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java index 2026d70a0011..3c885fe933d5 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java @@ -3482,6 +3482,50 @@ private static TableNameSplitResult splitTableName(String tableName) new TableNameSplitResult(tableName.substring(0, metadataMarkerIndex), Optional.of(tableName.substring(metadataMarkerIndex))); } + @Override + public void setViewColumnComment(ConnectorSession session, SchemaTableName viewName, ConnectorViewDefinition viewDefinition, String columnName, Optional comment) + { + Table table = metastore.getTable(new HiveIdentity(session), viewName.getSchemaName(), viewName.getTableName()).get(); + + List viewColumns = new ArrayList<>(); + for (ConnectorViewDefinition.ViewColumn viewColumn : viewDefinition.getColumns()) { + if (columnName.equals(viewColumn.getName())) { + ConnectorViewDefinition.ViewColumn updatedColumn = new ConnectorViewDefinition.ViewColumn(viewColumn.getName(), viewColumn.getType(), comment); + viewColumns.add(updatedColumn); + } + else { + viewColumns.add(viewColumn); + } + } + + ConnectorViewDefinition definition = new ConnectorViewDefinition(viewDefinition.getOriginalSql(), + viewDefinition.getCatalog(), + viewDefinition.getSchema(), + viewColumns, + viewDefinition.getComment(), + viewDefinition.getOwner(), + viewDefinition.isRunAsInvoker()); + + Table.Builder tableBuilder = Table.builder() + .setDatabaseName(table.getDatabaseName()) + .setTableName(table.getTableName()) + .setOwner(table.getOwner()) + .setTableType(table.getTableType()) + .setDataColumns(table.getDataColumns()) + .setPartitionColumns(table.getPartitionColumns()) + .setParameters(table.getParameters()) + .setViewOriginalText(Optional.of(encodeViewData(definition))) + .setViewExpandedText(table.getViewExpandedText()); + + tableBuilder.getStorageBuilder() + .setStorageFormat(table.getStorage().getStorageFormat()) + .setLocation(table.getStorage().getLocation()); + Table newTable = tableBuilder.build(); + PrincipalPrivileges principalPrivileges = accessControlMetadata.isUsingSystemSecurity() ? NO_PRIVILEGES : buildInitialPrivilegeSet(session.getUser()); + + metastore.replaceTable(new HiveIdentity(session), viewName.getSchemaName(), viewName.getTableName(), newTable, principalPrivileges); + } + private static class TableNameSplitResult { private final String baseTableName;