Skip to content

Commit

Permalink
Add ability to load after scripts from local filesystem (#600)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
bpkroth authored Feb 14, 2025
1 parent 15ad613 commit 9834ead
Show file tree
Hide file tree
Showing 4 changed files with 77 additions and 7 deletions.
45 changes: 39 additions & 6 deletions src/main/java/com/oltpbenchmark/api/BenchmarkModule.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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);
}
}
}

Expand Down Expand Up @@ -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());
}

Expand Down
2 changes: 1 addition & 1 deletion src/main/java/com/oltpbenchmark/util/ScriptRunner.java
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down
29 changes: 29 additions & 0 deletions src/test/java/com/oltpbenchmark/api/AbstractTestLoader.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(),
Expand Down
8 changes: 8 additions & 0 deletions src/test/java/com/oltpbenchmark/api/after-load-external.sql
Original file line number Diff line number Diff line change
@@ -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);

0 comments on commit 9834ead

Please sign in to comment.