Add support for CREATE and DROP SCHEMA in BigQuery

trinodb · Apr 19, 2021 · 269e9de · 269e9de
1 parent f4dcf75
commit 269e9de
Show file tree

Hide file tree

Showing 4 changed files with 89 additions and 0 deletions.
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
@@ -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
@@ -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<String, Object> 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/...o-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java b/...o-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/...ino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java b/...ino-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()
  {