Skip to content

Commit

Permalink
Add redirection awareness for COMMENT task
Browse files Browse the repository at this point in the history
  • Loading branch information
@findinpath @phd3
findinpath authored and phd3 committed Mar 1, 2022
1 parent f02dc03 commit 8b91718
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 17 deletions.
30 changes: 16 additions & 14 deletions core/trino-main/src/main/java/io/trino/execution/CommentTask.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import io.trino.execution.warnings.WarningCollector;
import io.trino.metadata.Metadata;
import io.trino.metadata.QualifiedObjectName;
import io.trino.metadata.RedirectionAwareTableHandle;
import io.trino.metadata.TableHandle;
import io.trino.security.AccessControl;
import io.trino.spi.connector.ColumnHandle;
Expand Down Expand Up @@ -69,37 +70,38 @@ public ListenableFuture<Void> execute(
Session session = stateMachine.getSession();

if (statement.getType() == Comment.Type.TABLE) {
QualifiedObjectName tableName = createQualifiedObjectName(session, statement, statement.getName());
Optional<TableHandle> tableHandle = metadata.getTableHandle(session, tableName);
if (tableHandle.isEmpty()) {
throw semanticException(TABLE_NOT_FOUND, statement, "Table does not exist: %s", tableName);
QualifiedObjectName originalTableName = createQualifiedObjectName(session, statement, statement.getName());
RedirectionAwareTableHandle redirectionAwareTableHandle = metadata.getRedirectionAwareTableHandle(session, originalTableName);
if (redirectionAwareTableHandle.getTableHandle().isEmpty()) {
throw semanticException(TABLE_NOT_FOUND, statement, "Table does not exist: %s", originalTableName);
}

accessControl.checkCanSetTableComment(session.toSecurityContext(), tableName);

metadata.setTableComment(session, tableHandle.get(), statement.getComment());
accessControl.checkCanSetTableComment(session.toSecurityContext(), redirectionAwareTableHandle.getRedirectedTableName().orElse(originalTableName));
TableHandle tableHandle = redirectionAwareTableHandle.getTableHandle().get();
metadata.setTableComment(session, tableHandle, statement.getComment());
}
else if (statement.getType() == Comment.Type.COLUMN) {
Optional<QualifiedName> prefix = statement.getName().getPrefix();
if (prefix.isEmpty()) {
throw semanticException(MISSING_TABLE, statement, "Table must be specified");
}

QualifiedObjectName tableName = createQualifiedObjectName(session, statement, prefix.get());
Optional<TableHandle> tableHandle = metadata.getTableHandle(session, tableName);
if (tableHandle.isEmpty()) {
throw semanticException(TABLE_NOT_FOUND, statement, "Table does not exist: " + tableName);
QualifiedObjectName originalTableName = createQualifiedObjectName(session, statement, prefix.get());
RedirectionAwareTableHandle redirectionAwareTableHandle = metadata.getRedirectionAwareTableHandle(session, originalTableName);
if (redirectionAwareTableHandle.getTableHandle().isEmpty()) {
throw semanticException(TABLE_NOT_FOUND, statement, "Table does not exist: " + originalTableName);
}
TableHandle tableHandle = redirectionAwareTableHandle.getTableHandle().get();

String columnName = statement.getName().getSuffix();
Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle.get());
Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle);
if (!columnHandles.containsKey(columnName)) {
throw semanticException(COLUMN_NOT_FOUND, statement, "Column does not exist: " + columnName);
}

accessControl.checkCanSetColumnComment(session.toSecurityContext(), tableName);
accessControl.checkCanSetColumnComment(session.toSecurityContext(), redirectionAwareTableHandle.getRedirectedTableName().orElse(originalTableName));

metadata.setColumnComment(session, tableHandle.get(), columnHandles.get(columnName), statement.getComment());
metadata.setColumnComment(session, tableHandle, columnHandles.get(columnName), statement.getComment());
}
else {
throw semanticException(NOT_SUPPORTED, statement, "Unsupported comment type: %s", statement.getType());
Expand Down
44 changes: 41 additions & 3 deletions ...product-tests/src/main/java/io/trino/tests/product/hive/TestHiveRedirectionToIceberg.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -344,9 +344,32 @@ public void testCommentTable()
assertTableComment("hive", "default", tableName).isNull();
assertTableComment("iceberg", "default", tableName).isNull();

//TODO restore test assertions after adding redirection awareness to the CommentTask
assertQueryFailure(() -> onTrino().executeQuery("COMMENT ON TABLE " + hiveTableName + " IS 'This is my table, there are many like it but this one is mine'"))
.hasMessageMatching("\\QQuery failed (#\\E\\S+\\Q): Cannot query Iceberg table 'default." + tableName + "'");
String tableComment = "This is my table, there are many like it but this one is mine";
onTrino().executeQuery(format("COMMENT ON TABLE " + hiveTableName + " IS '%s'", tableComment));

assertTableComment("hive", "default", tableName).isEqualTo(tableComment);
assertTableComment("iceberg", "default", tableName).isEqualTo(tableComment);

onTrino().executeQuery("DROP TABLE " + icebergTableName);
}

@Test(groups = {HIVE_ICEBERG_REDIRECTIONS, PROFILE_SPECIFIC_TESTS})
public void testCommentColumn()
{
String tableName = "iceberg_comment_column_" + randomTableSuffix();
String hiveTableName = "hive.default." + tableName;
String icebergTableName = "iceberg.default." + tableName;
String columnName = "nationkey";
createIcebergTable(icebergTableName, false);

assertColumnComment("hive", "default", tableName, columnName).isNull();
assertColumnComment("iceberg", "default", tableName, columnName).isNull();

String columnComment = "Internal identifier for the nation";
onTrino().executeQuery(format("COMMENT ON COLUMN %s.%s IS '%s'", hiveTableName, columnName, columnComment));

assertColumnComment("hive", "default", tableName, columnName).isEqualTo(columnComment);
assertColumnComment("iceberg", "default", tableName, columnName).isEqualTo(columnComment);

onTrino().executeQuery("DROP TABLE " + icebergTableName);
}
Expand Down Expand Up @@ -468,6 +491,21 @@ private static QueryResult readTableComment(String catalog, String schema, Strin
param(VARCHAR, tableName));
}

private static AbstractStringAssert<?> assertColumnComment(String catalog, String schema, String tableName, String columnName)
{
QueryResult queryResult = readColumnComment(catalog, schema, tableName, columnName);
return Assertions.assertThat((String) getOnlyElement(getOnlyElement(queryResult.rows())));
}

private static QueryResult readColumnComment(String catalog, String schema, String tableName, String columnName)
{
return onTrino().executeQuery(
format("SELECT comment FROM %s.information_schema.columns WHERE table_schema = ? AND table_name = ? AND column_name = ?", catalog),
param(VARCHAR, schema),
param(VARCHAR, tableName),
param(VARCHAR, columnName));
}

private static void assertResultsEqual(QueryResult first, QueryResult second)
{
assertThat(first).containsOnly(second.rows().stream()
Expand Down

0 comments on commit 8b91718

Please sign in to comment.