Skip to content

Commit

Permalink
Add support for table and column comments on the memory connector
Browse files Browse the repository at this point in the history
  • Loading branch information
@findinpath @ebyhr
findinpath authored and ebyhr committed Sep 1, 2022
1 parent cb9d67d commit 0c8296b
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 22 deletions.
1 change: 1 addition & 0 deletions docs/src/main/sphinx/connector/memory.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
^^^^^^^^^^
Expand Down
20 changes: 17 additions & 3 deletions plugin/trino-memory/src/main/java/io/trino/plugin/memory/ColumnInfo.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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<String> comment;

public ColumnInfo(ColumnHandle handle, String name, Type type, Optional<String> 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()
Expand All @@ -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();
}
}
46 changes: 35 additions & 11 deletions plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -236,9 +237,6 @@ public synchronized void createTable(ConnectorSession session, ConnectorTableMet
@Override
public synchronized MemoryOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional<ConnectorTableLayout> 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();
Expand All @@ -248,10 +246,7 @@ public synchronized MemoryOutputTableHandle beginCreateTable(ConnectorSession se
ImmutableList.Builder<ColumnInfo> 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);
Expand All @@ -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()));
}
Expand Down Expand Up @@ -406,7 +402,7 @@ private void updateRowsOnHosts(long tableId, Collection<Slice> 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
Expand Down Expand Up @@ -460,4 +456,32 @@ public Optional<SampleApplicationResult<ConnectorTableHandle>> 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<String> 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<String> 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()));
}
}
16 changes: 14 additions & 2 deletions plugin/trino-memory/src/main/java/io/trino/plugin/memory/TableInfo.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -34,13 +36,16 @@ public class TableInfo
private final List<ColumnInfo> columns;
private final Map<HostAddress, MemoryDataFragment> dataFragments;

public TableInfo(long id, String schemaName, String tableName, List<ColumnInfo> columns, Map<HostAddress, MemoryDataFragment> dataFragments)
private final Optional<String> comment;

public TableInfo(long id, String schemaName, String tableName, List<ColumnInfo> columns, Map<HostAddress, MemoryDataFragment> dataFragments, Optional<String> 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()
Expand Down Expand Up @@ -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<ColumnInfo> getColumns()
Expand All @@ -89,4 +96,9 @@ public Map<HostAddress, MemoryDataFragment> getDataFragments()
{
return dataFragments;
}

public Optional<String> getComment()
{
return comment;
}
}
6 changes: 0 additions & 6 deletions plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,12 +93,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;

Expand Down

0 comments on commit 0c8296b

Please sign in to comment.