Merge pull request #39 from flyway/improve-databricks

Improve databricks

Author: Barry-RG

flyway-database-databricks/src/main/java/org/flywaydb/community/database/databricks/DatabricksSchema.java

@@ -9,6 +9,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 public class DatabricksSchema extends Schema<DatabricksDatabase, DatabricksTable> {
     /**
@@ -20,25 +21,40 @@ public DatabricksSchema(JdbcTemplate jdbcTemplate, DatabricksDatabase database,
         super(jdbcTemplate, database, name);
     }
 
-    private List<String> fetchAllObjs(String obj) throws SQLException  {
+    private List<String> fetchAllObjs(String obj, String column) throws SQLException  {
         List<Map<String, String>> tableInfos = jdbcTemplate.queryForList(
                 "show " + obj + "s from " + database.quote(name)
         );
         List<String> tableNames = new ArrayList<String>();
         for (Map<String, String> tableInfo : tableInfos) {
-            tableNames.add(tableInfo.get("tableName"));
+            tableNames.add(tableInfo.get(column));
         }
         return tableNames;
     }
+    
+    private List<String> fetchAllSchemas() throws SQLException {
+        List<Map<String, String>> schemaInfos = jdbcTemplate.queryForList("show schemas");
+        List<String> schemaNames = new ArrayList<>();
+        for (Map<String, String> schemaInfo : schemaInfos) {
+            schemaNames.add(schemaInfo.get("databaseName"));
+        }
+        return schemaNames;
+    }
+    
+    private List<String> fetchAllTables() throws SQLException {
+        var tables = fetchAllObjs("table", "tableName");
+        var views = fetchAllObjs("view", "viewName");
+        return tables.stream().filter(t -> !views.contains(t)).toList();
+    }
 
     @Override
     protected boolean doExists() throws SQLException {
-        return fetchAllObjs("table").size() > 0;
+        return fetchAllSchemas().stream().anyMatch(schema -> Objects.equals(schema, name));
     }
 
     @Override
     protected boolean doEmpty() throws SQLException {
-        return fetchAllObjs("table").size() == 0;
+        return fetchAllTables().size() == 0;
     }
 
     @Override
@@ -53,19 +69,18 @@ protected void doDrop() throws SQLException {
 
     @Override
     protected void doClean() throws SQLException {
-        for (String statement : generateDropStatements("TABLE")) {
+        for (String statement : generateDropStatements("TABLE", fetchAllTables())) {
             jdbcTemplate.execute(statement);
         }
-        for (String statement : generateDropStatements("VIEW")) {
+        for (String statement : generateDropStatements("VIEW", fetchAllObjs("VIEW", "viewName"))) {
             jdbcTemplate.execute(statement);
         }
-        for (String statement : generateDropStatements("FUNCTION")) {
+        for (String statement : generateDropStatements("FUNCTION",fetchAllObjs("USER FUNCTION", "function"))) {
             jdbcTemplate.execute(statement);
         }
     }
 
-    private List<String> generateDropStatements(String objType) throws SQLException {
-        List<String> names = fetchAllObjs(objType);
+    private List<String> generateDropStatements(String objType, List<String> names) throws SQLException {
         List<String> statements = new ArrayList<>();
         for (String domainName : names) {
             statements.add("drop " + objType + " if exists " + database.quote(name, domainName) + ";");
@@ -76,7 +91,7 @@ private List<String> generateDropStatements(String objType) throws SQLException
 
     @Override
     protected DatabricksTable[] doAllTables() throws SQLException {
-        List<String> tableNames = fetchAllObjs("TABLE");
+        List<String> tableNames = fetchAllTables();
         DatabricksTable[] tables = new DatabricksTable[tableNames.size()];
         for (int i = 0; i < tableNames.size(); i++) {
             tables[i] = new DatabricksTable(jdbcTemplate, database, this, tableNames.get(i));