From 91f4b632a990a0f5a34f2ad235afcae8de5121f4 Mon Sep 17 00:00:00 2001 From: Marius Grama Date: Wed, 31 Aug 2022 11:16:58 +0200 Subject: [PATCH] Add support for table and column comments on the memory connector --- docs/src/main/sphinx/connector/memory.rst | 1 + .../io/trino/plugin/memory/ColumnInfo.java | 20 ++++++-- .../trino/plugin/memory/MemoryMetadata.java | 46 ++++++++++++++----- .../io/trino/plugin/memory/TableInfo.java | 16 ++++++- .../memory/TestMemoryConnectorTest.java | 6 --- 5 files changed, 67 insertions(+), 22 deletions(-) diff --git a/docs/src/main/sphinx/connector/memory.rst b/docs/src/main/sphinx/connector/memory.rst index 2f646c509230..920c6c00c25d 100644 --- a/docs/src/main/sphinx/connector/memory.rst +++ b/docs/src/main/sphinx/connector/memory.rst @@ -64,6 +64,7 @@ statements, the connector supports the following features: * :doc:`/sql/drop-table` * :doc:`/sql/create-schema` * :doc:`/sql/drop-schema` +* :doc:`/sql/comment` DROP TABLE ^^^^^^^^^^ diff --git a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/ColumnInfo.java b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/ColumnInfo.java index 9e702da30f9f..e49210685ea6 100644 --- a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/ColumnInfo.java +++ b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/ColumnInfo.java @@ -17,6 +17,9 @@ import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.type.Type; +import java.util.Optional; + +import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; public class ColumnInfo @@ -25,11 +28,14 @@ public class ColumnInfo private final String name; private final Type type; - public ColumnInfo(ColumnHandle handle, String name, Type type) + private final Optional comment; + + public ColumnInfo(ColumnHandle handle, String name, Type type, Optional comment) { this.handle = requireNonNull(handle, "handle is null"); this.name = requireNonNull(name, "name is null"); this.type = requireNonNull(type, "type is null"); + this.comment = requireNonNull(comment, "comment is null"); } public ColumnHandle getHandle() @@ -44,12 +50,20 @@ public String getName() public ColumnMetadata getMetadata() { - return new ColumnMetadata(name, type); + return ColumnMetadata.builder() + .setName(name) + .setType(type) + .setComment(comment) + .build(); } @Override public String toString() { - return name + "::" + type; + return toStringHelper(this) + .add("name", name) + .add("type", type) + .add("comment", comment) + .toString(); } } diff --git a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java index 88aa70b74d9e..4c4719b010ea 100644 --- a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java +++ b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java @@ -55,19 +55,20 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.OptionalDouble; import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.atomic.AtomicLong; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; import static io.trino.spi.StandardErrorCode.NOT_FOUND; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.StandardErrorCode.SCHEMA_NOT_EMPTY; import static io.trino.spi.connector.RetryMode.NO_RETRIES; import static io.trino.spi.connector.SampleType.SYSTEM; @@ -220,7 +221,7 @@ public synchronized void renameTable(ConnectorSession session, ConnectorTableHan long tableId = handle.getId(); TableInfo oldInfo = tables.get(tableId); - tables.put(tableId, new TableInfo(tableId, newTableName.getSchemaName(), newTableName.getTableName(), oldInfo.getColumns(), oldInfo.getDataFragments())); + tables.put(tableId, new TableInfo(tableId, newTableName.getSchemaName(), newTableName.getTableName(), oldInfo.getColumns(), oldInfo.getDataFragments(), oldInfo.getComment())); tableIds.remove(oldInfo.getSchemaTableName()); tableIds.put(newTableName, tableId); @@ -236,9 +237,6 @@ public synchronized void createTable(ConnectorSession session, ConnectorTableMet @Override public synchronized MemoryOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional layout, RetryMode retryMode) { - if (tableMetadata.getComment().isPresent()) { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with table comment"); - } checkSchemaExists(tableMetadata.getTable().getSchemaName()); checkTableNotExists(tableMetadata.getTable()); long tableId = nextTableId.getAndIncrement(); @@ -248,10 +246,7 @@ public synchronized MemoryOutputTableHandle beginCreateTable(ConnectorSession se ImmutableList.Builder columns = ImmutableList.builder(); for (int i = 0; i < tableMetadata.getColumns().size(); i++) { ColumnMetadata column = tableMetadata.getColumns().get(i); - if (column.getComment() != null) { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with column comment"); - } - columns.add(new ColumnInfo(new MemoryColumnHandle(i), column.getName(), column.getType())); + columns.add(new ColumnInfo(new MemoryColumnHandle(i), column.getName(), column.getType(), Optional.ofNullable(column.getComment()))); } tableIds.put(tableMetadata.getTable(), tableId); @@ -260,7 +255,8 @@ public synchronized MemoryOutputTableHandle beginCreateTable(ConnectorSession se tableMetadata.getTable().getSchemaName(), tableMetadata.getTable().getTableName(), columns.build(), - new HashMap<>())); + new HashMap<>(), + tableMetadata.getComment())); return new MemoryOutputTableHandle(tableId, ImmutableSet.copyOf(tableIds.values())); } @@ -406,7 +402,7 @@ private void updateRowsOnHosts(long tableId, Collection fragments) dataFragments.merge(memoryDataFragment.getHostAddress(), memoryDataFragment, MemoryDataFragment::merge); } - tables.put(tableId, new TableInfo(tableId, info.getSchemaName(), info.getTableName(), info.getColumns(), dataFragments)); + tables.put(tableId, new TableInfo(tableId, info.getSchemaName(), info.getTableName(), info.getColumns(), dataFragments, info.getComment())); } @Override @@ -460,4 +456,32 @@ public Optional> applySample(Conne new MemoryTableHandle(table.getId(), table.getLimit(), OptionalDouble.of(table.getSampleRatio().orElse(1) * sampleRatio)), true)); } + + @Override + public synchronized void setTableComment(ConnectorSession session, ConnectorTableHandle tableHandle, Optional comment) + { + MemoryTableHandle table = (MemoryTableHandle) tableHandle; + TableInfo info = tables.get(table.getId()); + checkArgument(info != null, "Table not found"); + tables.put(table.getId(), new TableInfo(table.getId(), info.getSchemaName(), info.getTableName(), info.getColumns(), info.getDataFragments(), comment)); + } + + @Override + public synchronized void setColumnComment(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle columnHandle, Optional comment) + { + MemoryTableHandle table = (MemoryTableHandle) tableHandle; + TableInfo info = tables.get(table.getId()); + checkArgument(info != null, "Table not found"); + tables.put( + table.getId(), + new TableInfo( + table.getId(), + info.getSchemaName(), + info.getTableName(), + info.getColumns().stream() + .map(tableColumn -> Objects.equals(tableColumn.getHandle(), columnHandle) ? new ColumnInfo(tableColumn.getHandle(), tableColumn.getName(), tableColumn.getMetadata().getType(), comment) : tableColumn) + .collect(toImmutableList()), + info.getDataFragments(), + info.getComment())); + } } diff --git a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/TableInfo.java b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/TableInfo.java index 8439099a0ae0..5fa33d26aed9 100644 --- a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/TableInfo.java +++ b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/TableInfo.java @@ -22,8 +22,10 @@ import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.stream.Collectors; +import static java.util.Collections.emptyMap; import static java.util.Objects.requireNonNull; public class TableInfo @@ -34,13 +36,16 @@ public class TableInfo private final List columns; private final Map dataFragments; - public TableInfo(long id, String schemaName, String tableName, List columns, Map dataFragments) + private final Optional comment; + + public TableInfo(long id, String schemaName, String tableName, List columns, Map dataFragments, Optional comment) { this.id = id; this.schemaName = requireNonNull(schemaName, "schemaName is null"); this.tableName = requireNonNull(tableName, "tableName is null"); this.columns = ImmutableList.copyOf(columns); this.dataFragments = ImmutableMap.copyOf(dataFragments); + this.comment = requireNonNull(comment, "comment is null"); } public long getId() @@ -69,7 +74,9 @@ public ConnectorTableMetadata getMetadata() new SchemaTableName(schemaName, tableName), columns.stream() .map(ColumnInfo::getMetadata) - .collect(Collectors.toList())); + .collect(Collectors.toList()), + emptyMap(), + comment); } public List getColumns() @@ -89,4 +96,9 @@ public Map getDataFragments() { return dataFragments; } + + public Optional getComment() + { + return comment; + } } 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 fe8e47d63b91..40713a122443 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 @@ -94,12 +94,6 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) case SUPPORTS_RENAME_COLUMN: return false; - case SUPPORTS_CREATE_TABLE_WITH_TABLE_COMMENT: - case SUPPORTS_CREATE_TABLE_WITH_COLUMN_COMMENT: - case SUPPORTS_COMMENT_ON_TABLE: - case SUPPORTS_COMMENT_ON_COLUMN: - return false; - case SUPPORTS_RENAME_SCHEMA: return false;