Introduce session setup (#422)

ranaalotaibiMS · Rana Alotaibi · Rana Alotaibi · web-flow · commit 15ad6137cb05 · 2025-02-14T15:01:47.000-06:00
This PR introduces functionality for configuring a session by executing predefined statements before the benchmark execution begins. Users should specify the file containing these statements within the tag `<sessionsetupfile> </sessionsetupfile>` in the benchmarks' config files. CC: @bpkroth For example, ```xml <?xml version="1.0"?> <parameters>  <type>sqlserver</type> ...  <sessionsetupfile>config/sqlserver/session_setup_sqlserver_cmds_example.sql</sessionsetupfile> ... ``` See Also: #192 --------- Co-authored-by: Rana Alotaibi <ranaalotaibi@Ranas-MacBook-Pro.local> Co-authored-by: Rana Alotaibi <ranaalotaibi@C1320-EMC-RE.redmond.corp.microsoft.com> Co-authored-by: Brian Kroth <bpkroth@users.noreply.github.com> Co-authored-by: Brian Kroth <bpkroth@gmail.com>
diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml
@@ -689,6 +689,13 @@ jobs:
           entrypoint: /opt/mssql-tools/bin/sqlcmd
           args: -U sa -P SApassword1 -S sqlserver -b -Q "USE benchbase; CREATE USER benchuser01 FROM LOGIN benchuser01; EXEC sp_addrolemember 'db_owner', 'benchuser01';"
 
+      - name: Setup privileged access for monitoring and session start tests
+        uses: docker://mcr.microsoft.com/mssql-tools:latest
+        with:
+          entrypoint: /opt/mssql-tools/bin/sqlcmd
+          args: -U sa -P SApassword1 -S sqlserver -b -Q "USE master; GRANT ALTER SERVER STATE, VIEW SERVER PERFORMANCE STATE to benchuser01;"
+
+
       - name: Run benchmark
         # Note: user/pass should match those used in sample configs.
         run: |
diff --git a/config/sqlserver/sample_noop_config.xml b/config/sqlserver/sample_noop_config.xml
@@ -11,6 +11,10 @@
     <isolation>TRANSACTION_SERIALIZABLE</isolation>
     <batchsize>128</batchsize>
 
+    <!-- Session setup statements file -->
+    <!-- uncommented here for CI pipeline testing -->
+    <sessionsetupfile>config/sqlserver/session_setup_sqlserver_cmds_example.sql</sessionsetupfile>
+
     <!-- This parameter has no affect on this benchmark-->
     <!-- There is no data to load -->
     <scalefactor>1</scalefactor>
diff --git a/config/sqlserver/sample_tpch_config.xml b/config/sqlserver/sample_tpch_config.xml
@@ -10,6 +10,8 @@
     <reconnectOnConnectionFailure>true</reconnectOnConnectionFailure>
     <isolation>TRANSACTION_SERIALIZABLE</isolation>
     <batchsize>1024</batchsize>
+    <!-- Session setup statements file -->
+    <!-- <sessionsetupfile>config/sqlserver/session_setup_sqlserver_cmds_example.sql</sessionsetupfile> -->
 
     <scalefactor>0.1</scalefactor>
 
diff --git a/config/sqlserver/session_setup_sqlserver_cmds_example.sql b/config/sqlserver/session_setup_sqlserver_cmds_example.sql
@@ -0,0 +1,4 @@
+-- SQL Server Database Console Command statements (DBCC)
+-- NOTE: Requires "ALTER SERVER STATE" permission
+DBCC DROPCLEANBUFFERS -- clear buffers (for cold runs)
+DBCC FREEPROCCACHE -- clean plan cache
diff --git a/docker/sqlserver-latest/up.sh b/docker/sqlserver-latest/up.sh
@@ -38,3 +38,6 @@ run_sqlcmd_in_docker -Q "CREATE LOGIN benchuser01 WITH PASSWORD='P@ssw0rd';" ||
 
 # Setup access
 run_sqlcmd_in_docker -Q "USE benchbase; CREATE USER benchuser01 FROM LOGIN benchuser01; EXEC sp_addrolemember 'db_owner', 'benchuser01';" || true
+
+# Setup privileged access for monitoring and session start tests
+run_sqlcmd_in_docker -Q "USE master; GRANT ALTER SERVER STATE, VIEW SERVER PERFORMANCE STATE TO benchuser01;" || true
diff --git a/src/main/java/com/oltpbenchmark/DBWorkload.java b/src/main/java/com/oltpbenchmark/DBWorkload.java
@@ -138,6 +138,7 @@ public static void main(String[] args) throws Exception {
       wrkld.setPassword(xmlConfig.getString("password"));
       wrkld.setRandomSeed(xmlConfig.getInt("randomSeed", -1));
       wrkld.setBatchSize(xmlConfig.getInt("batchsize", 128));
+      wrkld.setSessionSetupFile(xmlConfig.getString("sessionsetupfile"));
       wrkld.setMaxRetries(xmlConfig.getInt("retries", 3));
       wrkld.setNewConnectionPerTxn(xmlConfig.getBoolean("newConnectionPerTxn", false));
       wrkld.setReconnectOnConnectionFailure(
@@ -196,6 +197,9 @@ public static void main(String[] args) throws Exception {
       initDebug.put("URL", wrkld.getUrl());
       initDebug.put("Isolation", wrkld.getIsolationString());
       initDebug.put("Batch Size", wrkld.getBatchSize());
+      initDebug.put("DDL Path", wrkld.getDDLPath());
+      initDebug.put("Loader Threads", wrkld.getLoaderThreads());
+      initDebug.put("Session Setup File", wrkld.getSessionSetupFile());
       initDebug.put("Scale Factor", wrkld.getScaleFactor());
       initDebug.put("Terminals", wrkld.getTerminals());
       initDebug.put("New Connection Per Txn", wrkld.getNewConnectionPerTxn());
@@ -247,6 +251,7 @@ public static void main(String[] args) throws Exception {
         if (xmlConfig.containsKey("afterload")) {
           bench.setAfterLoadScriptPath(xmlConfig.getString("afterload"));
         }
+        initDebug.put("After Load Script", bench.getAfterLoadScriptPath());
 
         TransactionType tmpType =
             bench.initTransactionType(
diff --git a/src/main/java/com/oltpbenchmark/WorkloadConfiguration.java b/src/main/java/com/oltpbenchmark/WorkloadConfiguration.java
@@ -17,7 +17,9 @@
 
 import com.oltpbenchmark.api.TransactionTypes;
 import com.oltpbenchmark.types.DatabaseType;
+import com.oltpbenchmark.util.FileUtil;
 import com.oltpbenchmark.util.ThreadUtil;
+import java.io.FileNotFoundException;
 import java.sql.Connection;
 import java.util.ArrayList;
 import java.util.List;
@@ -33,6 +35,7 @@ public class WorkloadConfiguration {
   private String password;
   private String driverClass;
   private int batchSize;
+  private String sessionSetupFile;
   private int maxRetries;
   private int randomSeed = -1;
   private double scaleFactor = 1.0;
@@ -120,6 +123,14 @@ public void setBatchSize(int batchSize) {
     this.batchSize = batchSize;
   }
 
+  public String getSessionSetupFile() {
+    return sessionSetupFile;
+  }
+
+  public void setSessionSetupFile(String sessionSetupFile) throws FileNotFoundException {
+    this.sessionSetupFile = FileUtil.checkPath(sessionSetupFile, "sessionsetupfile");
+  }
+
   public int getMaxRetries() {
     return maxRetries;
   }
@@ -292,8 +303,8 @@ public String getDDLPath() {
   }
 
   /** Set the path in which we can find the ddl script. */
-  public void setDDLPath(String ddlPath) {
-    this.ddlPath = ddlPath;
+  public void setDDLPath(String ddlPath) throws FileNotFoundException {
+    this.ddlPath = FileUtil.checkPath(ddlPath, "ddlpath");
   }
 
   /** A utility method that init the phaseIterator and dialectMap */
@@ -394,8 +405,16 @@ public String toString() {
         + ", driverClass='"
         + driverClass
         + '\''
+        + ", reconnectOnFailure="
+        + reconnectOnConnectionFailure
+        + ", newConnectionPerTxn="
+        + newConnectionPerTxn
         + ", batchSize="
         + batchSize
+        + ", ddlpath="
+        + ddlPath
+        + ", sessionSetupFile="
+        + sessionSetupFile
         + ", maxRetries="
         + maxRetries
         + ", scaleFactor="
diff --git a/src/main/java/com/oltpbenchmark/api/BenchmarkModule.java b/src/main/java/com/oltpbenchmark/api/BenchmarkModule.java
@@ -93,6 +93,10 @@ public final void setAfterLoadScriptPath(String scriptPath) {
     this.afterLoadScriptPath = scriptPath;
   }
 
+  public String getAfterLoadScriptPath() {
+    return this.afterLoadScriptPath;
+  }
+
   // --------------------------------------------------------------------------
   // IMPLEMENTING CLASS INTERFACE
   // --------------------------------------------------------------------------
diff --git a/src/main/java/com/oltpbenchmark/api/Worker.java b/src/main/java/com/oltpbenchmark/api/Worker.java
@@ -26,6 +26,9 @@
 import com.oltpbenchmark.types.TransactionStatus;
 import com.oltpbenchmark.util.Histogram;
 import com.oltpbenchmark.util.SQLUtil;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.sql.SQLRecoverableException;
@@ -188,6 +191,13 @@ public final void run() {
     // In case of reuse reset the measurements
     latencies = new LatencyRecord(workloadState.getTestStartNs());
 
+    // Invoke setup session
+    try {
+      this.setupSession();
+    } catch (Throwable ex) {
+      throw new RuntimeException("Unexpected error when setting up the session " + this, ex);
+    }
+
     // Invoke initialize callback
     try {
       this.initialize();
@@ -723,6 +733,32 @@ protected void initialize() {
     // The default is to do nothing
   }
 
+  /**
+   * Set up the session by running a set of statements before benchmark execution begins. The path
+   * of the file where a set of statements defined should be added in &lt;sessionsetupfile&gt;
+   * &lt;/sessionsetupfile&gt;
+   */
+  protected void setupSession() {
+    try {
+      String setupSessionFile = configuration.getSessionSetupFile();
+      if (setupSessionFile == null || setupSessionFile.isEmpty()) {
+        return;
+      }
+
+      String statements = new String(Files.readAllBytes(Paths.get(setupSessionFile)));
+      if (statements.isEmpty()) {
+        return;
+      }
+
+      try (Statement stmt = conn.createStatement()) {
+        stmt.execute(statements);
+      }
+      // conn.commit();
+    } catch (SQLException | IOException ex) {
+      throw new RuntimeException("Failed setting up session", ex);
+    }
+  }
+
   /**
    * Invoke a single transaction for the given TransactionType
    *
diff --git a/src/main/java/com/oltpbenchmark/api/collectors/monitoring/SQLServerMonitor.java b/src/main/java/com/oltpbenchmark/api/collectors/monitoring/SQLServerMonitor.java
@@ -21,7 +21,8 @@
 
 /**
  * Implementation of a monitor specific to SQLServer. Uses SQLServer's system tables to extract
- * relevant query and system information.
+ * relevant query and system information. Note: Requires "VIEW SERVER PERFORMANCE STATE"
+ * permissions.
  */
 public class SQLServerMonitor extends DatabaseMonitor {
 
diff --git a/src/main/java/com/oltpbenchmark/util/FileUtil.java b/src/main/java/com/oltpbenchmark/util/FileUtil.java
@@ -18,6 +18,7 @@
 package com.oltpbenchmark.util;
 
 import java.io.File;
+import java.io.FileNotFoundException;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.util.regex.Pattern;
@@ -83,6 +84,16 @@ public static boolean exists(String path) {
     return (new File(path).exists());
   }
 
+  public static String checkPath(String path, String name) throws FileNotFoundException {
+    if (path != null) path = path.trim();
+    if (path == null || path.isEmpty()) return null;
+
+    if (!FileUtil.exists(path)) {
+      throw new FileNotFoundException(name + " not found:" + path);
+    }
+    return path;
+  }
+
   /**
    * Create any directory in the list paths if it doesn't exist
    *
diff --git a/src/main/resources/benchmarks/noop/ddl-generic.sql b/src/main/resources/benchmarks/noop/ddl-generic.sql
@@ -1,4 +1,8 @@
 DROP TABLE IF EXISTS FAKE;
 CREATE TABLE FAKE (
     ID INT PRIMARY KEY
+);
+DROP TABLE IF EXISTS FAKE2;
+CREATE TABLE FAKE2 (
+    ID INT
 );
diff --git a/src/test/java/com/oltpbenchmark/api/AbstractTestCase.java b/src/test/java/com/oltpbenchmark/api/AbstractTestCase.java
@@ -68,6 +68,7 @@ public abstract class AbstractTestCase<T extends BenchmarkModule> {
   protected final boolean createDatabase;
   protected final boolean loadDatabase;
   protected final String ddlOverridePath;
+  protected final String sessionSetupFile;
 
   private static final AtomicInteger portCounter = new AtomicInteger(9001);
   private static final int MAX_PORT_NUMBER = 65535;
@@ -77,13 +78,27 @@ public AbstractTestCase(boolean createDatabase, boolean loadDatabase) {
     this.createDatabase = createDatabase;
     this.loadDatabase = loadDatabase;
     this.ddlOverridePath = null;
+    this.sessionSetupFile = null;
   }
 
   public AbstractTestCase(boolean createDatabase, boolean loadDatabase, String ddlOverridePath) {
     this.benchmark = null;
     this.createDatabase = createDatabase;
     this.loadDatabase = loadDatabase;
     this.ddlOverridePath = ddlOverridePath;
+    this.sessionSetupFile = null;
+  }
+
+  public AbstractTestCase(
+      boolean createDatabase,
+      boolean loadDatabase,
+      String ddlOverridePath,
+      String sessionSetupFile) {
+    this.benchmark = null;
+    this.createDatabase = createDatabase;
+    this.loadDatabase = loadDatabase;
+    this.ddlOverridePath = ddlOverridePath;
+    this.sessionSetupFile = sessionSetupFile;
   }
 
   public abstract List<Class<? extends Procedure>> procedures();
@@ -127,6 +142,7 @@ public final void setUp() throws Exception {
     this.workConf.setBenchmarkName(
         BenchmarkModule.convertBenchmarkClassToBenchmarkName(benchmarkClass()));
     this.workConf.setDDLPath(this.ddlOverridePath);
+    this.workConf.setSessionSetupFile(this.sessionSetupFile);
 
     customWorkloadConfiguration(this.workConf);
 
diff --git a/src/test/java/com/oltpbenchmark/api/AbstractTestLoader.java b/src/test/java/com/oltpbenchmark/api/AbstractTestLoader.java
@@ -17,7 +17,6 @@
 package com.oltpbenchmark.api;
 
 import static org.junit.Assert.*;
-import static org.junit.Assert.fail;
 
 import com.oltpbenchmark.catalog.Table;
 import com.oltpbenchmark.util.Histogram;
diff --git a/src/test/java/com/oltpbenchmark/api/AbstractTestWorker.java b/src/test/java/com/oltpbenchmark/api/AbstractTestWorker.java
@@ -41,6 +41,10 @@ public AbstractTestWorker(String ddlOverridePath) {
     super(true, true, ddlOverridePath);
   }
 
+  public AbstractTestWorker(String ddlOverridePath, String sessionSetupFile) {
+    super(true, true, ddlOverridePath, sessionSetupFile);
+  }
+
   @Override
   public List<String> ignorableTables() {
     return null;
@@ -67,12 +71,14 @@ public void testGetProcedure() {
     }
   }
 
-  /** testExecuteWork */
+  /* testExecuteWork
+   * Similar to Worker.run()
+   */
   @Test
   public void testExecuteWork() throws Exception {
-
     Worker<?> w = workers.get(0);
     assertNotNull(w);
+    w.setupSession();
     w.initialize();
     assertFalse(this.conn.isReadOnly());
     for (TransactionType txnType : this.workConf.getTransTypes()) {
diff --git a/src/test/java/com/oltpbenchmark/benchmarks/noop/TestNoOpLoader.java b/src/test/java/com/oltpbenchmark/benchmarks/noop/TestNoOpLoader.java
@@ -34,6 +34,6 @@ public Class<NoOpBenchmark> benchmarkClass() {
 
   @Override
   public List<String> ignorableTables() {
-    return List.of("FAKE");
+    return List.of("FAKE", "FAKE2");
   }
 }
diff --git a/src/test/java/com/oltpbenchmark/benchmarks/noop/TestNoOpWorker.java b/src/test/java/com/oltpbenchmark/benchmarks/noop/TestNoOpWorker.java
@@ -16,9 +16,21 @@
 
 package com.oltpbenchmark.benchmarks.noop;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
 import com.oltpbenchmark.api.AbstractTestWorker;
+import com.oltpbenchmark.api.BenchmarkModule;
 import com.oltpbenchmark.api.Procedure;
+import com.oltpbenchmark.api.Worker;
+import com.oltpbenchmark.catalog.Table;
+import com.oltpbenchmark.util.SQLUtil;
+import java.sql.ResultSet;
+import java.sql.Statement;
 import java.util.List;
+import org.junit.Test;
 
 public class TestNoOpWorker extends AbstractTestWorker<NoOpBenchmark> {
 
@@ -31,4 +43,33 @@ public List<Class<? extends Procedure>> procedures() {
   public Class<NoOpBenchmark> benchmarkClass() {
     return NoOpBenchmark.class;
   }
+
+  @Test
+  public void testNoSessionSetupFile() throws Exception {
+    // Check that there is no session setup file assigned to the worker's config
+    assertNull("Session setup file should be null", this.workConf.getSessionSetupFile());
+
+    List<Worker<? extends BenchmarkModule>> workers = this.benchmark.makeWorkers();
+    Worker<?> worker = workers.get(0);
+    assertNull(
+        "Session setup file should be null",
+        worker.getWorkloadConfiguration().getSessionSetupFile());
+
+    // Make sure there are no rows in the table
+    this.testExecuteWork();
+
+    Table catalog_tbl = this.catalog.getTable("FAKE2");
+    String sql = SQLUtil.getCountSQL(this.workConf.getDatabaseType(), catalog_tbl);
+    try (Statement stmt = conn.createStatement();
+        ResultSet result = stmt.executeQuery(sql); ) {
+
+      assertNotNull(result);
+
+      boolean adv = result.next();
+      assertTrue(sql, adv);
+
+      int count = result.getInt(1);
+      assertEquals(0, count);
+    }
+  }
 }
diff --git a/src/test/java/com/oltpbenchmark/benchmarks/noop/TestNoOpWorkerWithSessionSetupFile.java b/src/test/java/com/oltpbenchmark/benchmarks/noop/TestNoOpWorkerWithSessionSetupFile.java
diff --git a/src/test/resources/benchmarks/noop/sessionSetupFile-hsqldb.sql b/src/test/resources/benchmarks/noop/sessionSetupFile-hsqldb.sql
