Add support for DROP SCHEMA CASCADE in Kudu

SemionPar · Aug 18, 2023 · 5c3326f · 5c3326f
1 parent 47081fd
commit 5c3326f
Show file tree

Hide file tree

Showing 7 changed files with 40 additions and 9 deletions.
diff --git a/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduClientSession.java b/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduClientSession.java
@@ -261,9 +261,9 @@ public void createSchema(String schemaName)
         schemaEmulation.createSchema(client, schemaName);
     }
 
-    public void dropSchema(String schemaName)
+    public void dropSchema(String schemaName, boolean cascade)
     {
-        schemaEmulation.dropSchema(client, schemaName);
+        schemaEmulation.dropSchema(client, schemaName, cascade);
     }
 
     public void dropTable(SchemaTableName schemaTableName)

diff --git a/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduMetadata.java b/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduMetadata.java
@@ -242,9 +242,9 @@ public void createSchema(ConnectorSession session, String schemaName, Map<String
     }
 
     @Override
-    public void dropSchema(ConnectorSession session, String schemaName)
+    public void dropSchema(ConnectorSession session, String schemaName, boolean cascade)
     {
-        clientSession.dropSchema(schemaName);
+        clientSession.dropSchema(schemaName, cascade);
     }
 
     @Override

diff --git a/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/schema/NoSchemaEmulation.java b/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/schema/NoSchemaEmulation.java
@@ -37,7 +37,7 @@ public void createSchema(KuduClientWrapper client, String schemaName)
     }
 
     @Override
-    public void dropSchema(KuduClientWrapper client, String schemaName)
+    public void dropSchema(KuduClientWrapper client, String schemaName, boolean cascade)
     {
         if (DEFAULT_SCHEMA.equals(schemaName)) {
             throw new TrinoException(GENERIC_USER_ERROR, "Deleting default schema not allowed.");

diff --git a/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/schema/SchemaEmulation.java b/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/schema/SchemaEmulation.java
@@ -22,7 +22,7 @@ public interface SchemaEmulation
 {
     void createSchema(KuduClientWrapper client, String schemaName);
 
-    void dropSchema(KuduClientWrapper client, String schemaName);
+    void dropSchema(KuduClientWrapper client, String schemaName, boolean cascade);
 
     boolean existsSchema(KuduClientWrapper client, String schemaName);
 

diff --git a/...-kudu/src/main/java/io/trino/plugin/kudu/schema/SchemaEmulationByTableNameConvention.java b/...-kudu/src/main/java/io/trino/plugin/kudu/schema/SchemaEmulationByTableNameConvention.java
@@ -39,6 +39,7 @@
 import static io.trino.plugin.kudu.KuduClientSession.DEFAULT_SCHEMA;
 import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static io.trino.spi.StandardErrorCode.GENERIC_USER_ERROR;
+import static io.trino.spi.StandardErrorCode.SCHEMA_NOT_EMPTY;
 
 public class SchemaEmulationByTableNameConvention
         implements SchemaEmulation
@@ -81,14 +82,18 @@ public boolean existsSchema(KuduClientWrapper client, String schemaName)
     }
 
     @Override
-    public void dropSchema(KuduClientWrapper client, String schemaName)
+    public void dropSchema(KuduClientWrapper client, String schemaName, boolean cascade)
     {
         if (DEFAULT_SCHEMA.equals(schemaName)) {
             throw new TrinoException(GENERIC_USER_ERROR, "Deleting default schema not allowed.");
         }
         try (KuduOperationApplier operationApplier = KuduOperationApplier.fromKuduClientWrapper(client)) {
             String prefix = getPrefixForTablesOfSchema(schemaName);
-            for (String name : client.getTablesList(prefix).getTablesList()) {
+            List<String> tables = client.getTablesList(prefix).getTablesList();
+            if (!cascade && !tables.isEmpty()) {
+                throw new TrinoException(SCHEMA_NOT_EMPTY, "Cannot drop non-empty schema '%s'".formatted(schemaName));
+            }
+            for (String name : tables) {
                 client.deleteTable(name);
             }
 

diff --git a/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/BaseKuduConnectorSmokeTest.java b/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/BaseKuduConnectorSmokeTest.java
@@ -16,6 +16,7 @@
 import io.trino.testing.BaseConnectorSmokeTest;
 import io.trino.testing.QueryRunner;
 import io.trino.testing.TestingConnectorBehavior;
+import org.testng.SkipException;
 import org.testng.annotations.Test;
 
 import java.util.Optional;
@@ -26,6 +27,7 @@
 import static io.trino.plugin.kudu.TestingKuduServer.EARLIEST_TAG;
 import static io.trino.testing.TestingNames.randomNameSuffix;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 public abstract class BaseKuduConnectorSmokeTest
         extends BaseConnectorSmokeTest
@@ -162,4 +164,29 @@ public void testCreateTableWithTableComment()
 
         assertUpdate("DROP TABLE " + tableName);
     }
+
+    @Test
+    public void testDropSchemaCascade()
+    {
+        String schemaName = "test_drop_schema_cascade_" + randomNameSuffix();
+        String tableName = "test_table" + randomNameSuffix();
+        try {
+            if (getKuduSchemaEmulationPrefix().isEmpty()) {
+                assertThatThrownBy(() -> assertUpdate("CREATE SCHEMA " + schemaName))
+                        .hasMessageContaining("Creating schema in Kudu connector not allowed if schema emulation is disabled.");
+                throw new SkipException("Cannot test when schema emulation is disabled");
+            }
+            assertUpdate("CREATE SCHEMA " + schemaName);
+            assertUpdate("CREATE TABLE " + schemaName + "." + tableName + " AS SELECT 1 a", 1);
+
+            assertThat(computeActual("SHOW SCHEMAS").getOnlyColumnAsSet()).contains(schemaName);
+
+            assertUpdate("DROP SCHEMA " + schemaName + " CASCADE");
+            assertThat(computeActual("SHOW SCHEMAS").getOnlyColumnAsSet()).doesNotContain(schemaName);
+        }
+        finally {
+            assertUpdate("DROP TABLE IF EXISTS " + schemaName + "." + tableName);
+            assertUpdate("DROP SCHEMA IF EXISTS " + schemaName);
+        }
+    }
 }
diff --git a/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java b/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java
@@ -72,7 +72,6 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior)
                 return false;
 
             case SUPPORTS_RENAME_SCHEMA:
-            case SUPPORTS_DROP_SCHEMA_CASCADE:
                 return false;
 
             case SUPPORTS_CREATE_TABLE_WITH_COLUMN_COMMENT: