From 9834eada6a796c20d15a0c733c2795afc5750ef8 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Fri, 14 Feb 2025 15:46:01 -0600 Subject: [PATCH] Add ability to load after scripts from local filesystem (#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 #595 - [x] Add unit test --- .../oltpbenchmark/api/BenchmarkModule.java | 45 ++++++++++++++++--- .../com/oltpbenchmark/util/ScriptRunner.java | 2 +- .../oltpbenchmark/api/AbstractTestLoader.java | 29 ++++++++++++ .../oltpbenchmark/api/after-load-external.sql | 8 ++++ 4 files changed, 77 insertions(+), 7 deletions(-) create mode 100644 src/test/java/com/oltpbenchmark/api/after-load-external.sql diff --git a/src/main/java/com/oltpbenchmark/api/BenchmarkModule.java b/src/main/java/com/oltpbenchmark/api/BenchmarkModule.java index 193151ccb..ceb632719 100644 --- a/src/main/java/com/oltpbenchmark/api/BenchmarkModule.java +++ b/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 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()); } diff --git a/src/main/java/com/oltpbenchmark/util/ScriptRunner.java b/src/main/java/com/oltpbenchmark/util/ScriptRunner.java index 9b192aa6b..bc25b66ac 100644 --- a/src/main/java/com/oltpbenchmark/util/ScriptRunner.java +++ b/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)) { diff --git a/src/test/java/com/oltpbenchmark/api/AbstractTestLoader.java b/src/test/java/com/oltpbenchmark/api/AbstractTestLoader.java index 83f8e88a8..3adba6c50 100644 --- a/src/test/java/com/oltpbenchmark/api/AbstractTestLoader.java +++ b/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(), diff --git a/src/test/java/com/oltpbenchmark/api/after-load-external.sql b/src/test/java/com/oltpbenchmark/api/after-load-external.sql new file mode 100644 index 000000000..55ed7ec7b --- /dev/null +++ b/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);