diff --git a/docs/src/main/sphinx/connector/bigquery.rst b/docs/src/main/sphinx/connector/bigquery.rst index d51ed84da754..4299cd760275 100644 --- a/docs/src/main/sphinx/connector/bigquery.rst +++ b/docs/src/main/sphinx/connector/bigquery.rst @@ -219,6 +219,7 @@ the following features: * :doc:`/sql/drop-table` * :doc:`/sql/create-schema` * :doc:`/sql/drop-schema` +* :doc:`/sql/comment` Table functions --------------- diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java index ce66521ecd80..67a24f9f05de 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java @@ -217,6 +217,18 @@ Job create(JobInfo jobInfo) return bigQuery.create(jobInfo); } + public void executeUpdate(QueryJobConfiguration job) + { + log.debug("Execute query: %s", job.getQuery()); + try { + bigQuery.query(job); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new BigQueryException(BaseHttpServiceException.UNKNOWN_CODE, format("Failed to run the query [%s]", job.getQuery()), e); + } + } + public TableResult query(String sql, boolean useQueryResultsCache, CreateDisposition createDisposition) { log.debug("Execute query: %s", sql); diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryFilterQueryBuilder.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryFilterQueryBuilder.java index 2511fdda0424..f0f80ac0408f 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryFilterQueryBuilder.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryFilterQueryBuilder.java @@ -26,14 +26,13 @@ import static com.google.common.base.Preconditions.checkState; import static com.google.common.collect.Iterables.getOnlyElement; +import static io.trino.plugin.bigquery.BigQueryUtil.quote; import static io.trino.plugin.bigquery.BigQueryUtil.toBigQueryColumnName; import static java.util.stream.Collectors.joining; import static java.util.stream.Collectors.toList; public class BigQueryFilterQueryBuilder { - private static final String QUOTE = "`"; - private static final String ESCAPED_QUOTE = "``"; private final TupleDomain tupleDomain; public static Optional buildFilter(TupleDomain tupleDomain) @@ -151,9 +150,4 @@ private Optional toPredicate(String columnName, String operator, Object } return Optional.of(quote(columnName) + " " + operator + " " + valueAsString.get()); } - - private String quote(String name) - { - return QUOTE + name.replace(QUOTE, ESCAPED_QUOTE) + QUOTE; - } } diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java index 6b0e7de7f433..64629addb815 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java @@ -17,6 +17,8 @@ import com.google.cloud.bigquery.DatasetId; import com.google.cloud.bigquery.DatasetInfo; import com.google.cloud.bigquery.Field; +import com.google.cloud.bigquery.QueryJobConfiguration; +import com.google.cloud.bigquery.QueryParameterValue; import com.google.cloud.bigquery.Schema; import com.google.cloud.bigquery.StandardTableDefinition; import com.google.cloud.bigquery.Table; @@ -94,7 +96,9 @@ import static io.trino.plugin.bigquery.BigQueryTableHandle.getPartitionType; import static io.trino.plugin.bigquery.BigQueryType.toField; import static io.trino.plugin.bigquery.BigQueryUtil.isWildcardTable; +import static io.trino.plugin.bigquery.BigQueryUtil.quote; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static java.lang.String.format; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; import static java.util.function.Function.identity; @@ -502,6 +506,44 @@ public Optional finishInsert(ConnectorSession session, return Optional.empty(); } + @Override + public void setTableComment(ConnectorSession session, ConnectorTableHandle tableHandle, Optional newComment) + { + BigQueryTableHandle table = (BigQueryTableHandle) tableHandle; + BigQueryClient client = bigQueryClientFactory.createBigQueryClient(session); + + RemoteTableName remoteTableName = table.asPlainTable().getRemoteTableName(); + String sql = format( + "ALTER TABLE %s.%s.%s SET OPTIONS (description = ?)", + quote(remoteTableName.getProjectId()), + quote(remoteTableName.getDatasetName()), + quote(remoteTableName.getTableName())); + client.executeUpdate(QueryJobConfiguration.newBuilder(sql) + .setQuery(sql) + .addPositionalParameter(QueryParameterValue.string(newComment.orElse(null))) + .build()); + } + + @Override + public void setColumnComment(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle columnHandle, Optional newComment) + { + BigQueryTableHandle table = (BigQueryTableHandle) tableHandle; + BigQueryColumnHandle column = (BigQueryColumnHandle) columnHandle; + BigQueryClient client = bigQueryClientFactory.createBigQueryClient(session); + + RemoteTableName remoteTableName = table.asPlainTable().getRemoteTableName(); + String sql = format( + "ALTER TABLE %s.%s.%s ALTER COLUMN %s SET OPTIONS (description = ?)", + quote(remoteTableName.getProjectId()), + quote(remoteTableName.getDatasetName()), + quote(remoteTableName.getTableName()), + column.getName()); + client.executeUpdate(QueryJobConfiguration.newBuilder(sql) + .setQuery(sql) + .addPositionalParameter(QueryParameterValue.string(newComment.orElse(null))) + .build()); + } + @Override public Optional> applyProjection( ConnectorSession session, diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java index e98bcc0dba01..7269662ab766 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java @@ -30,6 +30,9 @@ public final class BigQueryUtil { + private static final String QUOTE = "`"; + private static final String ESCAPED_QUOTE = "``"; + private static final Set INTERNAL_ERROR_MESSAGES = ImmutableSet.of( "HTTP/2 error code: INTERNAL_ERROR", "Connection closed with unknown cause", @@ -73,4 +76,9 @@ public static boolean isWildcardTable(TableDefinition.Type type, String tableNam { return type == TABLE && tableName.contains("*"); } + + public static String quote(String name) + { + return QUOTE + name.replace(QUOTE, ESCAPED_QUOTE) + QUOTE; + } } diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java index b2b23b3eea07..a6fb8bcea588 100644 --- a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java +++ b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java @@ -77,8 +77,6 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) case SUPPORTS_ADD_COLUMN: case SUPPORTS_DROP_COLUMN: case SUPPORTS_RENAME_COLUMN: - case SUPPORTS_COMMENT_ON_TABLE: - case SUPPORTS_COMMENT_ON_COLUMN: case SUPPORTS_NEGATIVE_DATE: case SUPPORTS_ARRAY: case SUPPORTS_ROW_TYPE: