azkaban-aplcache

refactor DatabaseSetup class for Quartz Test being developed

9/13/2017 8:21:51 PM

Details

diff --git a/azkaban-common/src/test/java/azkaban/test/Utils.java b/azkaban-common/src/test/java/azkaban/test/Utils.java
index d5ab9ae..dc915dc 100644
--- a/azkaban-common/src/test/java/azkaban/test/Utils.java
+++ b/azkaban-common/src/test/java/azkaban/test/Utils.java
@@ -6,7 +6,6 @@ import azkaban.db.AzDBTestUtility.EmbeddedH2BasicDataSource;
 import azkaban.db.AzkabanDataSource;
 import azkaban.db.DatabaseOperator;
 import azkaban.db.DatabaseSetup;
-import azkaban.utils.Props;
 import com.google.inject.AbstractModule;
 import com.google.inject.Guice;
 import com.google.inject.Injector;
@@ -33,11 +32,9 @@ public class Utils {
     final AzkabanDataSource dataSource = new EmbeddedH2BasicDataSource();
 
     final String sqlScriptsDir = new File("../azkaban-db/src/main/sql/").getCanonicalPath();
-    final Props props = new Props();
-    props.put("database.sql.scripts.dir", sqlScriptsDir);
-
     final DatabaseSetup setup = new DatabaseSetup(dataSource, sqlScriptsDir);
     setup.updateDatabase();
+
     return new DatabaseOperator(new QueryRunner(dataSource));
   }
 }
diff --git a/azkaban-db/src/main/java/azkaban/db/DatabaseSetup.java b/azkaban-db/src/main/java/azkaban/db/DatabaseSetup.java
index 477c7a1..435f528 100644
--- a/azkaban-db/src/main/java/azkaban/db/DatabaseSetup.java
+++ b/azkaban-db/src/main/java/azkaban/db/DatabaseSetup.java
@@ -23,47 +23,30 @@ import java.io.FileInputStream;
 import java.io.IOException;
 import java.sql.Connection;
 import java.sql.SQLException;
-import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Map;
 import java.util.Set;
 import org.apache.commons.dbutils.QueryRunner;
 import org.apache.commons.io.IOUtils;
 import org.apache.log4j.Logger;
 
 /**
- * TODO kunkun-tang: this class is quite messy now, needs to be refactored.
+ * This class is used for creating DB tables by specifying where the create scripts are. The
+ * input to this class is a folder path, which includes all create table sql scripts. The script's
+ * name should follow: create.[table_name].sql in order to be identified. This class is used for
+ * unit test only for now.
+ *
+ * Todo kunkun-tang: We need to fix some reliability issues if we rely on this class to
+ * create tables when launching AZ in future.
  */
 public class DatabaseSetup {
 
-  public static final String DATABASE_SQL_SCRIPT_DIR =
-      "database.sql.scripts.dir";
-  private static final Logger logger = Logger
-      .getLogger(DatabaseSetup.class);
-  private static final String DEFAULT_SCRIPT_PATH = "sql";
+  private static final Logger logger = Logger .getLogger(DatabaseSetup.class);
   private static final String CREATE_SCRIPT_PREFIX = "create.";
   private static final String SQL_SCRIPT_SUFFIX = ".sql";
 
-  private static final String INSERT_DB_PROPERTY =
-      "INSERT INTO properties (name, type, value, modified_time) values (?,?,?,?)";
-  private static final String UPDATE_DB_PROPERTY =
-      "UPDATE properties SET value=?,modified_time=? WHERE name=? AND type=?";
-
   private final AzkabanDataSource dataSource;
-  private final Set<String> missingTables = new HashSet<>();
-  private final Map<String, String> installedVersions = new HashMap<>();
-  private String version;
-  private boolean needsUpdating;
-
   private String scriptPath = null;
 
-  public DatabaseSetup(final AzkabanDataSource ds) {
-    this.dataSource = ds;
-    if (this.scriptPath == null) {
-      this.scriptPath = DEFAULT_SCRIPT_PATH;
-    }
-  }
-
   public DatabaseSetup(final AzkabanDataSource ds, final String path) {
     this.dataSource = ds;
     this.scriptPath = path;
@@ -71,11 +54,12 @@ public class DatabaseSetup {
 
   public void updateDatabase()
       throws SQLException, IOException {
-    findMissingTables();
-    createNewTables();
+    final Set<String> tables = collectAllTables();
+    createTables(tables);
   }
 
-  private void findMissingTables() {
+  private Set<String> collectAllTables() {
+    final Set<String> tables = new HashSet<>();
     final File directory = new File(this.scriptPath);
     final File[] createScripts =
         directory.listFiles(new PrefixSuffixFileFilter(
@@ -85,84 +69,39 @@ public class DatabaseSetup {
         final String name = script.getName();
         final String[] nameSplit = name.split("\\.");
         final String tableName = nameSplit[1];
-        this.missingTables.add(tableName);
+        tables.add(tableName);
       }
     }
+    return tables;
   }
 
-  private void createNewTables() throws SQLException, IOException {
+  private void createTables(final Set<String> tables) throws SQLException, IOException {
     final Connection conn = this.dataSource.getConnection();
     conn.setAutoCommit(false);
     try {
-      // Make sure that properties table is created first.
-      if (this.missingTables.contains("properties")) {
-        runTableScripts(conn, "properties", this.version, this.dataSource.getDBType(),
-            false);
-      }
-      for (final String table : this.missingTables) {
-        if (!table.equals("properties")) {
-          runTableScripts(conn, table, this.version, this.dataSource.getDBType(), false);
-          // update version as we have create a new table
-          this.installedVersions.put(table, this.version);
-        }
+      for (final String table : tables) {
+        runTableScripts(conn, table);
       }
     } finally {
       conn.close();
     }
   }
 
+  private void runTableScripts(final Connection conn, final String table)
+      throws IOException, SQLException {
+    logger.info("Creating new table " + table);
 
-  private void runTableScripts(final Connection conn, final String table, final String version,
-      final String dbType, final boolean update) throws IOException, SQLException {
-    String scriptName = "";
-    if (update) {
-      scriptName = "update." + table + "." + version;
-      logger.info("Update table " + table + " to version " + version);
-    } else {
-      scriptName = "create." + table;
-      logger.info("Creating new table " + table + " version " + version);
-    }
-
-    final String dbSpecificScript = scriptName + "." + dbType + ".sql";
-
-    File script = new File(this.scriptPath, dbSpecificScript);
-    if (!script.exists()) {
-      final String dbScript = scriptName + ".sql";
-      script = new File(this.scriptPath, dbScript);
-
-      if (!script.exists()) {
-        throw new IOException("Creation files do not exist for table " + table);
-      }
-    }
-
+    final String dbSpecificScript = "create." + table + ".sql";
+    final File script = new File(this.scriptPath, dbSpecificScript);
     BufferedInputStream buff = null;
     try {
       buff = new BufferedInputStream(new FileInputStream(script));
       final String queryStr = IOUtils.toString(buff);
-
       final String[] splitQuery = queryStr.split(";\\s*\n");
-
       final QueryRunner runner = new QueryRunner();
-
       for (final String query : splitQuery) {
         runner.update(conn, query);
       }
-
-      // If it's properties, then we want to commit the table before we update
-      // it
-      if (table.equals("properties")) {
-        conn.commit();
-      }
-
-      final String propertyName = table + ".version";
-      if (!this.installedVersions.containsKey(table)) {
-        runner.update(conn, INSERT_DB_PROPERTY, propertyName,
-            1, version,
-            System.currentTimeMillis());
-      } else {
-        runner.update(conn, UPDATE_DB_PROPERTY, version,
-            System.currentTimeMillis(), propertyName, 1);
-      }
       conn.commit();
     } finally {
       IOUtils.closeQuietly(buff);
@@ -177,7 +116,7 @@ public class DatabaseSetup {
     private final String prefix;
     private final String suffix;
 
-    public PrefixSuffixFileFilter(final String prefix, final String suffix) {
+    PrefixSuffixFileFilter(final String prefix, final String suffix) {
       this.prefix = prefix;
       this.suffix = suffix;
     }
@@ -190,11 +129,8 @@ public class DatabaseSetup {
 
       final String name = pathname.getName();
       final int length = name.length();
-      if (this.suffix.length() > length || this.prefix.length() > length) {
-        return false;
-      }
-
-      return name.startsWith(this.prefix) && name.endsWith(this.suffix);
+      return this.suffix.length() <= length && this.prefix.length() <= length && name
+          .startsWith(this.prefix) && name.endsWith(this.suffix);
     }
   }
 }