SQL Database DDL & Usability Improvements

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) :
     <USE> {
         s.add(this);
     }
-    <DATABASE>
+    [ <DATABASE> ]
     databaseName = CompoundIdentifier()
     {
         return new SqlUseDatabase(

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<CalciteSchema, String> schema(
+  public static Pair<CalciteSchema, String> 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<CalciteSchema, String> schema(
     return Pair.of(checkStateNotNull(schema, "Got null sub-schema for path '%s'", path), name(id));
   }
 
-  private static @Nullable CalciteSchema childSchema(CalciteSchema rootSchema, List<String> path) {
+  public static @Nullable CalciteSchema childSchema(CalciteSchema rootSchema, List<String> path) {
+    if (path == null) {
+      return null;
+    }
     @Nullable CalciteSchema schema = rootSchema;
     for (String p : path) {
       if (schema == null) {

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<String> 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

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();