Add ability to load after scripts from local filesystem (cmu-db#600)

Still fallsback to loading from resource stream when file is not found.

So far logging is only done to the debug stream.

Quick (hackish) fix to close cmu-db#595

- [x] Add unit test

Author: bpkroth

src/main/java/com/oltpbenchmark/api/BenchmarkModule.java

@@ -19,9 +19,11 @@
 import com.oltpbenchmark.catalog.AbstractCatalog;
 import com.oltpbenchmark.types.DatabaseType;
 import com.oltpbenchmark.util.ClassUtil;
+import com.oltpbenchmark.util.FileUtil;
 import com.oltpbenchmark.util.SQLUtil;
 import com.oltpbenchmark.util.ScriptRunner;
 import com.oltpbenchmark.util.ThreadUtil;
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.sql.Connection;
@@ -89,8 +91,22 @@ public final Connection makeConnection() throws SQLException {
 
   private String afterLoadScriptPath = null;
 
-  public final void setAfterLoadScriptPath(String scriptPath) {
-    this.afterLoadScriptPath = scriptPath;
+  public final void setAfterLoadScriptPath(String scriptPath) throws FileNotFoundException {
+    if (scriptPath != null) scriptPath = scriptPath.trim();
+    try {
+      this.afterLoadScriptPath = FileUtil.checkPath(scriptPath, "afterload");
+      return;
+    } catch (FileNotFoundException ex) {
+      this.afterLoadScriptPath = null;
+    }
+
+    if (this.afterLoadScriptPath == null && scriptPath != null && !scriptPath.isEmpty()) {
+      if (this.getClass().getResourceAsStream(scriptPath) == null) {
+        throw new FileNotFoundException(
+            "Couldn't find " + scriptPath + " as local file or resource.");
+      }
+      this.afterLoadScriptPath = scriptPath;
+    }
   }
 
   public String getAfterLoadScriptPath() {
@@ -241,8 +257,23 @@ public final void runScript(String scriptPath) throws SQLException, IOException
     try (Connection conn = this.makeConnection()) {
       DatabaseType dbType = this.workConf.getDatabaseType();
       ScriptRunner runner = new ScriptRunner(conn, true, true);
-      LOG.debug("Executing script [{}] for database type [{}]", scriptPath, dbType);
-      runner.runScript(scriptPath);
+      LOG.debug(
+          "Checking for script [{}] on local filesystem for database type [{}]",
+          scriptPath,
+          dbType);
+      if (FileUtil.exists(scriptPath)) {
+        LOG.debug(
+            "Executing script [{}] from local filesystem for database type [{}]",
+            scriptPath,
+            dbType);
+        runner.runExternalScript(scriptPath);
+      } else {
+        LOG.debug(
+            "Executing script [{}] from resource stream for database type [{}]",
+            scriptPath,
+            dbType);
+        runner.runScript(scriptPath);
+      }
     }
   }
 
@@ -274,11 +305,13 @@ public final Loader<? extends BenchmarkModule> loadDatabase()
 
     if (this.afterLoadScriptPath != null) {
       LOG.debug(
-          "Running script after load for {} benchmark...",
+          "Running script {} after load for {} benchmark...",
+          this.afterLoadScriptPath,
           this.workConf.getBenchmarkName().toUpperCase());
       runScript(this.afterLoadScriptPath);
       LOG.debug(
-          "Finished running script after load for {} benchmark...",
+          "Finished running script {} after load for {} benchmark...",
+          this.afterLoadScriptPath,
           this.workConf.getBenchmarkName().toUpperCase());
     }
 

src/main/java/com/oltpbenchmark/util/ScriptRunner.java

@@ -56,7 +56,7 @@ public void runExternalScript(String path) throws IOException, SQLException {
 
   public void runScript(String path) throws IOException, SQLException {
 
-    LOG.debug("trying to find file by path {}", path);
+    LOG.debug("trying to find file by resource stream path {}", path);
 
     try (InputStream in = this.getClass().getResourceAsStream(path);
         Reader reader = new InputStreamReader(in)) {

src/test/java/com/oltpbenchmark/api/AbstractTestLoader.java

@@ -21,6 +21,7 @@
 import com.oltpbenchmark.catalog.Table;
 import com.oltpbenchmark.util.Histogram;
 import com.oltpbenchmark.util.SQLUtil;
+import java.nio.file.Paths;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -73,6 +74,34 @@ public void testLoadWithAfterLoad() throws Exception {
     validateLoad();
   }
 
+  /** testLoad with external after load script */
+  @Test
+  public void testLoadWithExternalAfterLoad() throws Exception {
+    String afterLoadScriptPath =
+        Paths.get("src", "test", "java", "com", "oltpbenchmark", "api", "after-load-external.sql")
+            .toAbsolutePath()
+            .toString();
+
+    this.benchmark.setAfterLoadScriptPath(afterLoadScriptPath);
+
+    this.benchmark.loadDatabase();
+
+    // A table called extra is added with after-load, with one entry zero
+    try (PreparedStatement stmt = conn.prepareStatement("SELECT * FROM extra_external");
+        ResultSet rs = stmt.executeQuery()) {
+      while (rs.next()) {
+        assertEquals(
+            "Table 'extra_external' from " + afterLoadScriptPath + " has value different than 1",
+            rs.getInt(1),
+            1);
+      }
+    } catch (Exception e) {
+      fail("Table 'extra_external' from " + afterLoadScriptPath + " was not created");
+    }
+
+    validateLoad();
+  }
+
   private void validateLoad() throws SQLException {
     assertFalse(
         "Failed to get table names for " + benchmark.getBenchmarkName().toUpperCase(),

src/test/java/com/oltpbenchmark/api/after-load-external.sql

@@ -0,0 +1,8 @@
+DROP TABLE IF EXISTS extra_external CASCADE;
+
+CREATE TABLE extra_external (
+    extra_pk int NOT NULL,
+    PRIMARY KEY (extra_pk)
+);
+
+INSERT INTO extra_external VALUES (1);