diff --git a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl index 94c0161c492c..cb8eec438728 100644 --- a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl +++ b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl @@ -381,7 +381,7 @@ SqlCreate SqlCreateDatabase(Span s, boolean replace) : } /** - * USE DATABASE ( catalog_name '.' )? database_name + * USE [ DATABASE ] ( catalog_name '.' )? database_name */ SqlCall SqlUseDatabase(Span s, String scope) : { @@ -391,7 +391,7 @@ SqlCall SqlUseDatabase(Span s, String scope) : { s.add(this); } - + [ ] databaseName = CompoundIdentifier() { return new SqlUseDatabase( diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java index 6d6be5d5a127..f8d7e6f73851 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java @@ -55,7 +55,7 @@ public static SqlNode column( } /** Returns the schema in which to create an object. */ - static Pair schema( + public static Pair schema( CalcitePrepare.Context context, boolean mutable, SqlIdentifier id) { CalciteSchema rootSchema = mutable ? context.getMutableRootSchema() : context.getRootSchema(); @Nullable CalciteSchema schema = null; @@ -72,7 +72,10 @@ static Pair schema( return Pair.of(checkStateNotNull(schema, "Got null sub-schema for path '%s'", path), name(id)); } - private static @Nullable CalciteSchema childSchema(CalciteSchema rootSchema, List path) { + public static @Nullable CalciteSchema childSchema(CalciteSchema rootSchema, List path) { + if (path == null) { + return null; + } @Nullable CalciteSchema schema = rootSchema; for (String p : path) { if (schema == null) { diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/InMemoryCatalog.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/InMemoryCatalog.java index cdee6c930224..68e80c2340fa 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/InMemoryCatalog.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/InMemoryCatalog.java @@ -18,7 +18,6 @@ package org.apache.beam.sdk.extensions.sql.meta.catalog; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.util.Collection; import java.util.Collections; @@ -111,13 +110,20 @@ public Collection databases() { @Override public boolean dropDatabase(String database, boolean cascade) { - checkState(!cascade, "%s does not support CASCADE.", getClass().getSimpleName()); + MetaStore metaStore = metaStores.get(database); + if (!cascade && metaStore != null && !metaStore.getTables().isEmpty()) { + throw new IllegalStateException("Database '" + database + "' is not empty."); + } boolean removed = databases.remove(database); + if (!removed) { + return false; + } if (database.equals(currentDatabase)) { currentDatabase = null; } - return removed; + metaStores.remove(database); + return true; } @Override diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliDatabaseTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliDatabaseTest.java index 588caa78a2b7..54682911fe11 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliDatabaseTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliDatabaseTest.java @@ -83,6 +83,15 @@ public void testUseDatabase() { assertEquals("my_database2", catalogManager.currentCatalog().currentDatabase()); } + @Test + public void testUseDatabaseWithoutDatabaseKeyword() { + assertEquals(DEFAULT, catalogManager.currentCatalog().currentDatabase()); + cli.execute("CREATE DATABASE my_database"); + assertEquals(DEFAULT, catalogManager.currentCatalog().currentDatabase()); + cli.execute("USE my_database"); + assertEquals("my_database", catalogManager.currentCatalog().currentDatabase()); + } + @Test public void testUseDatabase_doesNotExist() { assertEquals(DEFAULT, catalogManager.currentCatalog().currentDatabase()); @@ -126,6 +135,49 @@ public void testDropDatabase_nonexistent() { cli.execute("DROP DATABASE my_database"); } + @Test + public void testDropDatabase_ifExists_nonexistent() { + assertFalse(catalogManager.currentCatalog().databaseExists("my_database")); + // Should not throw exception + cli.execute("DROP DATABASE IF EXISTS my_database"); + assertFalse(catalogManager.currentCatalog().databaseExists("my_database")); + } + + @Test + public void testDropDatabase_ifExists_exists() { + cli.execute("CREATE DATABASE my_database"); + assertTrue(catalogManager.currentCatalog().databaseExists("my_database")); + cli.execute("DROP DATABASE IF EXISTS my_database"); + assertFalse(catalogManager.currentCatalog().databaseExists("my_database")); + } + + @Test + public void testDropDatabase_notEmpty_restrict() { + cli.execute("CREATE DATABASE db_1"); + cli.execute("USE db_1"); + + TestTableProvider testTableProvider = new TestTableProvider(); + catalogManager.registerTableProvider(testTableProvider); + cli.execute("CREATE EXTERNAL TABLE person(id int, name varchar, age int) TYPE 'test'"); + + thrown.expect(CalciteContextException.class); + thrown.expectMessage("Database 'db_1' is not empty."); + cli.execute("DROP DATABASE db_1"); + } + + @Test + public void testDropDatabase_notEmpty_cascade() { + cli.execute("CREATE DATABASE db_1"); + cli.execute("USE db_1"); + + TestTableProvider testTableProvider = new TestTableProvider(); + catalogManager.registerTableProvider(testTableProvider); + cli.execute("CREATE EXTERNAL TABLE person(id int, name varchar, age int) TYPE 'test'"); + + cli.execute("DROP DATABASE db_1 CASCADE"); + assertFalse(catalogManager.currentCatalog().databaseExists("db_1")); + } + @Test public void testCreateInsertDropTableUsingDefaultDatabase() { Catalog catalog = catalogManager.currentCatalog();