From 269e9de432f03ee85f261e2a76dfa460c7b7e781 Mon Sep 17 00:00:00 2001 From: "yuya.ebihara" Date: Sat, 10 Apr 2021 00:17:51 +0900 Subject: [PATCH] Add support for CREATE and DROP SCHEMA in BigQuery --- .../trino/plugin/bigquery/BigQueryClient.java | 10 ++++++ .../plugin/bigquery/BigQueryMetadata.java | 21 ++++++++++++ .../TestBigQueryCaseInsensitiveMapping.java | 25 ++++++++++++++ .../TestBigQueryIntegrationSmokeTest.java | 33 +++++++++++++++++++ 4 files changed, 89 insertions(+) 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 79faa13bd2075..c6c2688626553 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 @@ -207,6 +207,16 @@ Table update(TableInfo table) return bigQuery.update(table); } + public void createSchema(DatasetInfo datasetInfo) + { + bigQuery.create(datasetInfo); + } + + public void dropSchema(DatasetId datasetId) + { + bigQuery.delete(datasetId); + } + public void createTable(TableInfo tableInfo) { bigQuery.create(tableInfo); 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 f33e749110a11..4c965a7e86266 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 @@ -15,6 +15,7 @@ import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.DatasetId; +import com.google.cloud.bigquery.DatasetInfo; import com.google.cloud.bigquery.Field; import com.google.cloud.bigquery.Schema; import com.google.cloud.bigquery.StandardTableDefinition; @@ -44,12 +45,14 @@ import io.trino.spi.connector.LimitApplicationResult; import io.trino.spi.connector.ProjectionApplicationResult; import io.trino.spi.connector.RecordCursor; +import io.trino.spi.connector.SchemaNotFoundException; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.connector.SchemaTablePrefix; import io.trino.spi.connector.SystemTable; import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.expression.ConnectorExpression; import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.security.TrinoPrincipal; import io.trino.spi.type.Type; import io.trino.spi.type.VarcharType; @@ -64,6 +67,7 @@ import static com.google.cloud.bigquery.TableDefinition.Type.TABLE; import static com.google.cloud.bigquery.TableDefinition.Type.VIEW; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.trino.plugin.bigquery.BigQueryType.toField; @@ -293,6 +297,23 @@ public ConnectorTableProperties getTableProperties(ConnectorSession session, Con return new ConnectorTableProperties(); } + @Override + public void createSchema(ConnectorSession session, String schemaName, Map properties, TrinoPrincipal owner) + { + checkArgument(properties.isEmpty(), "Can't have properties for schema creation"); + DatasetInfo datasetInfo = DatasetInfo.newBuilder(schemaName).build(); + bigQueryClient.createSchema(datasetInfo); + } + + @Override + public void dropSchema(ConnectorSession session, String schemaName) + { + String remoteSchemaName = bigQueryClient.toRemoteDataset(projectId, schemaName) + .map(RemoteDatabaseObject::getOnlyRemoteName) + .orElseThrow(() -> new SchemaNotFoundException(schemaName)); + bigQueryClient.dropSchema(DatasetId.of(remoteSchemaName)); + } + @Override public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean ignoreExisting) { diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java index 89a71c4bd1fa0..6a1d1a756ae8e 100644 --- a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java +++ b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java @@ -172,6 +172,31 @@ public void testTableNameClash() } } + @Test + public void testDropSchema() + { + String schema = "Test_Drop_Case_Sensitive"; + bigQuerySqlExecutor.execute(format("DROP SCHEMA IF EXISTS `%s`", schema)); + bigQuerySqlExecutor.execute(format("CREATE SCHEMA `%s`", schema)); + + assertUpdate("DROP SCHEMA " + schema.toLowerCase(ENGLISH)); + } + + @Test + public void testDropSchemaNameClash() + { + String schema = "Test_Drop_Case_Sensitive_Clash"; + bigQuerySqlExecutor.execute(format("DROP SCHEMA IF EXISTS `%s`", schema)); + bigQuerySqlExecutor.execute(format("DROP SCHEMA IF EXISTS `%s`", schema.toLowerCase(ENGLISH))); + bigQuerySqlExecutor.execute(format("CREATE SCHEMA `%s`", schema)); + bigQuerySqlExecutor.execute(format("CREATE SCHEMA `%s`", schema.toLowerCase(ENGLISH))); + + assertQueryFails("DROP SCHEMA " + schema.toLowerCase(ENGLISH), "Found ambiguous names in BigQuery.*"); + + bigQuerySqlExecutor.execute(format("DROP SCHEMA `%s`", schema)); + bigQuerySqlExecutor.execute(format("DROP SCHEMA `%s`", schema.toLowerCase(ENGLISH))); + } + private AutoCloseable withSchema(String schemaName) { bigQuerySqlExecutor.createDataset(schemaName); diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java index 2c73bc31f1fd1..a5d13a30267a7 100644 --- a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java +++ b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java @@ -47,6 +47,39 @@ protected QueryRunner createQueryRunner() ImmutableMap.of()); } + @Test + public void testCreateSchema() + { + String schemaName = "test_create_schema"; + + assertUpdate("DROP SCHEMA IF EXISTS " + schemaName); + + assertUpdate("CREATE SCHEMA " + schemaName); + assertUpdate("CREATE SCHEMA IF NOT EXISTS " + schemaName); + + assertQueryFails( + "CREATE SCHEMA " + schemaName, + format("\\Qline 1:1: Schema 'bigquery.%s' already exists\\E", schemaName)); + + assertUpdate("DROP SCHEMA " + schemaName); + } + + @Test + public void testDropSchema() + { + String schemaName = "test_drop_schema"; + + assertUpdate("DROP SCHEMA IF EXISTS " + schemaName); + assertUpdate("CREATE SCHEMA " + schemaName); + + assertUpdate("DROP SCHEMA " + schemaName); + assertUpdate("DROP SCHEMA IF EXISTS " + schemaName); + + assertQueryFails( + "DROP SCHEMA " + schemaName, + format("\\Qline 1:1: Schema 'bigquery.%s' does not exist\\E", schemaName)); + } + @Override public void testDescribeTable() {