After load functionality with test (#397)

Splitting after load functionality needed for Oracle config from PR
#379.

Add an after load option for benchmark configs.

Benchmark configs can add an afterload tag specifying a path to an SQL
file in resources. This file will be executed after loading database.

Tested in `AbstractTestLoad.testLoadWithAfterLoad()`

---------

Co-authored-by: Brian Kroth <[email protected]>

Author: ngbanguyen

src/main/java/com/oltpbenchmark/DBWorkload.java

@@ -214,6 +214,11 @@ public static void main(String[] args) throws Exception {
                     postExecutionWait = xmlConfig.getLong(key + "/postExecutionWait");
                 }
 
+                // After load
+                if (xmlConfig.containsKey("afterload")) {
+                    bench.setAfterLoadScriptPath(xmlConfig.getString("afterload"));
+                }
+
                 TransactionType tmpType = bench.initTransactionType(txnName, txnId + txnIdOffset, preExecutionWait, postExecutionWait);
 
                 // Keep a reference for filtering
@@ -626,7 +631,7 @@ private static void runCreator(BenchmarkModule bench) throws SQLException, IOExc
         bench.createDatabase();
     }
 
-    private static void runLoader(BenchmarkModule bench) throws SQLException, InterruptedException {
+    private static void runLoader(BenchmarkModule bench) throws IOException, SQLException, InterruptedException {
         LOG.debug(String.format("Loading %s Database", bench));
         bench.loadDatabase();
     }

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

@@ -104,6 +104,16 @@ public final Connection makeConnection() throws SQLException {
         }
     }
 
+    private String afterLoadScriptPath = null;
+
+    public final void setAfterLoadScriptPath(String scriptPath) {
+        this.afterLoadScriptPath = scriptPath;
+    }
+
+    public String getAfterLoadScriptPath() {
+        return this.afterLoadScriptPath;
+    }
+
     // --------------------------------------------------------------------------
     // IMPLEMENTING CLASS INTERFACE
     // --------------------------------------------------------------------------
@@ -248,11 +258,19 @@ public final void createDatabase(DatabaseType dbType, Connection conn) throws SQ
             }
     }
 
+    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);
+        }
+    }
 
     /**
      * Invoke this benchmark's database loader
      */
-    public final Loader<? extends BenchmarkModule> loadDatabase() throws SQLException, InterruptedException {
+    public final Loader<? extends BenchmarkModule> loadDatabase() throws IOException, SQLException, InterruptedException {
         Loader<? extends BenchmarkModule> loader;
 
         loader = this.makeLoaderImpl();
@@ -275,6 +293,12 @@ public final Loader<? extends BenchmarkModule> loadDatabase() throws SQLExceptio
             }
         }
 
+        if (this.afterLoadScriptPath != null) {
+            LOG.debug("Running script after load for {} benchmark...", this.workConf.getBenchmarkName().toUpperCase());
+            runScript(this.afterLoadScriptPath);
+            LOG.debug("Finished running script after load for {} benchmark...", this.workConf.getBenchmarkName().toUpperCase());
+        }
+
         return loader;
     }
 

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

@@ -17,22 +17,22 @@
 
 package com.oltpbenchmark.api;
 
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
-
 import com.oltpbenchmark.catalog.Table;
 import com.oltpbenchmark.util.Histogram;
 import com.oltpbenchmark.util.SQLUtil;
 import org.junit.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.List;
 
+import static org.junit.Assert.*;
+import static org.junit.Assert.fail;
+
 public abstract class AbstractTestLoader<T extends BenchmarkModule> extends AbstractTestCase<T> {
 
     private static final Logger LOG = LoggerFactory.getLogger(AbstractTestLoader.class);
@@ -58,6 +58,27 @@ public void testLoad() throws Exception {
 
     }
 
+    /**
+     * testLoad with after load script
+     */
+    @Test
+    public void testLoadWithAfterLoad() throws Exception {
+        this.benchmark.setAfterLoadScriptPath("/after-load.sql");
+
+        this.benchmark.loadDatabase();
+
+        // A table called extra is added with after-load, with one entry zero
+        try (PreparedStatement stmt = conn.prepareStatement("SELECT * FROM extra"); ResultSet rs = stmt.executeQuery()) {
+            while (rs.next()) {
+                assertEquals("Table 'extra' from after-load.sql has value different than 0", rs.getInt(1), 0);
+            }
+        } catch (Exception e) {
+            fail("Table 'extra' from after-load.sql was not created");
+        }
+
+        validateLoad();
+    }
+
     private void validateLoad() throws SQLException {
         assertFalse("Failed to get table names for " + benchmark.getBenchmarkName().toUpperCase(), this.catalog.getTables().isEmpty());
 

src/test/resources/after-load.sql

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