Support for multiple drivers in tpcc (#150)

Add support for multiple drivers, can be changed based on DB type.

Co-authored-by: Sonal Agarwal <[email protected]>

config/workload_all.xml

@@ -1,6 +1,6 @@
 <?xml version="1.0"?>
 <parameters>
-    <dbtype>postgres</dbtype>
+    <dbtype>yugabyte</dbtype>
     <driver>com.yugabyte.Driver</driver>
     <port>5433</port>
     <username>yugabyte</username>

config/workload_all_pg.xml

@@ -0,0 +1,56 @@
+<?xml version="1.0"?>
+<parameters>
+    <dbtype>postgres</dbtype>
+    <driver>org.postgresql.Driver</driver>
+    <port>5432</port>
+    <username>postgres</username>
+    <DBName>postgres</DBName>
+    <password>Password321</password>
+    <isolation>TRANSACTION_REPEATABLE_READ</isolation>
+
+    <sslCert></sslCert>
+    <sslKey></sslKey>
+    <jdbcURL></jdbcURL>
+
+    <batchSize>128</batchSize>
+    <useKeyingTime>true</useKeyingTime>
+    <useThinkTime>true</useThinkTime>
+    <enableForeignKeysAfterLoad>true</enableForeignKeysAfterLoad>
+    <hikariConnectionTimeoutMs>180000</hikariConnectionTimeoutMs>
+    <useStoredProcedures>true</useStoredProcedures>
+    <displayEnhancedLatencyMetrics>false</displayEnhancedLatencyMetrics>
+    <trackPerSQLStmtLatencies>false</trackPerSQLStmtLatencies>
+
+   	<transactiontypes>
+        <transaction>
+            <name>NewOrder</name>
+            <weight>45</weight>
+        </transaction>
+        <transaction>
+            <name>Payment</name>
+            <weight>43</weight>
+        </transaction>
+        <transaction>
+            <name>OrderStatus</name>
+            <weight>4</weight>
+        </transaction>
+        <transaction>
+            <name>Delivery</name>
+            <weight>4</weight>
+        </transaction>
+        <transaction>
+            <name>StockLevel</name>
+            <weight>4</weight>
+        </transaction>
+   	</transactiontypes>
+
+    <runtime>1800</runtime>
+    <rate>10000</rate>
+    <!--
+      Set the number of retries to 0 as retrying when the number of warehouses is
+      high is pointless as it just leads to more failures.
+    -->
+    <maxRetriesPerTransaction>2</maxRetriesPerTransaction>
+    <maxLoaderRetries>2</maxLoaderRetries>
+
+</parameters>

ivy.xml

@@ -22,7 +22,8 @@
       <!-- JDBC Drivers -->
 		<dependency org="com.yugabyte" name="jdbc-yugabytedb" rev="42.3.5-yb-3" force="true" conf="compile->compile(*),master(*);runtime->runtime(*)"/>
 		<dependency org="org.hsqldb" name="hsqldb" rev="2.4.1" force="true" conf="compile->compile(*),master(*);runtime->runtime(*)"/>
-
+		<!-- https://mvnrepository.com/artifact/org.postgresql/postgresql -->
+		<dependency org="org.postgresql" name="postgresql" rev="42.4.5"/>
 		<dependency org="com.zaxxer" name="HikariCP" rev="3.4.5" force="true" conf="compile->compile(*),master(*);runtime->runtime(*)"/>
 
 		<!-- Core Libraries -->

pom.xml

@@ -77,11 +77,17 @@
             <artifactId>httpmime</artifactId>
             <version>4.3.1</version>
         </dependency>
-          <dependency>
+        <dependency>
             <groupId>com.yugabyte</groupId>
             <artifactId>jdbc-yugabytedb</artifactId>
-              <version>42.3.5-yb-3</version>
-          </dependency>
+            <version>42.3.5-yb-3</version>
+        </dependency>
+        <!-- https://mvnrepository.com/artifact/org.postgresql/postgresql -->
+        <dependency>
+            <groupId>org.postgresql</groupId>
+            <artifactId>postgresql</artifactId>
+            <version>42.4.5</version>
+        </dependency>
         <dependency>
             <groupId>org.hsqldb</groupId>
             <artifactId>hsqldb-j5</artifactId>

src/com/oltpbenchmark/ConfigFileOptions.java

@@ -9,6 +9,10 @@ public ConfigFileOptions(String filePath) throws ConfigurationException {
         super(filePath);
     }
 
+    public String getDbType() {
+        return xmlConfig.getString("dbtype");
+    }
+
     public String getDbDriver() {
         return xmlConfig.getString("driver");
     }

src/com/oltpbenchmark/DBWorkload.java

@@ -52,6 +52,7 @@ public class DBWorkload {
   private static final Map<Integer, String> transactionTypes = new HashMap<>();
   private static JsonMetricsHelper jsonMetricsHelper = new JsonMetricsHelper();
 
+  public static String dbtype = "";
   /**
    * Returns true if asserts are enabled. This assumes that
    * we're always using the default system ClassLoader
@@ -159,6 +160,8 @@ public static void main(String[] args) throws Exception {
       wrkld.setNodes(nodes);
 
       wrkld.setDBName(configOptions.getDbName());
+      wrkld.setDBType(configOptions.getDbType());
+      DBWorkload.dbtype = configOptions.getDbType();
       wrkld.setDBUsername(configOptions.getDbUsername());
       wrkld.setDBPassword(configOptions.getDbPassword());
 

src/com/oltpbenchmark/WorkloadConfiguration.java

@@ -41,6 +41,7 @@ public void setBenchmarkName(String benchmarkName) {
   }
 
   private List<String> nodes;
+  private String db_type;
   private String db_name;
   private String db_username;
   private String db_password;
@@ -117,6 +118,10 @@ public void setDBName(String dbname) {
     this.db_name = dbname;
   }
 
+  public void setDBType(String dbtype) {
+    this.db_type = dbtype;
+  }
+
   public void setLoaderThreads(int loaderThreads) {
     this.loaderThreads = loaderThreads;
   }
@@ -140,6 +145,9 @@ public String getDBName() {
     return db_name;
   }
 
+  public String getDBType() {
+    return db_type;
+  }
   public void setDBUsername(String username) {
     this.db_username = username;
   }

src/com/oltpbenchmark/api/BenchmarkModule.java

@@ -90,9 +90,13 @@ public void createDataSource() {
                 ThreadUtil.sleep(5000);
             }
             Properties props = new Properties();
-            props.setProperty("dataSourceClassName", "com.yugabyte.ysql.YBClusterAwareDataSource");
-	    //props.setProperty("dataSource.serverNames", ip);
-	    props.setProperty("dataSource.serverName", ip);
+            if(workConf.getDBType().equals("yugabyte")){
+                props.setProperty("dataSourceClassName", "com.yugabyte.ysql.YBClusterAwareDataSource");
+            } else {
+                props.setProperty("dataSourceClassName", "org.postgresql.ds.PGSimpleDataSource");
+            }
+	        //props.setProperty("dataSource.serverNames", ip);
+	        props.setProperty("dataSource.serverName", ip);
             props.setProperty("dataSource.portNumber", Integer.toString(workConf.getPort()));
             props.setProperty("dataSource.user", workConf.getDBUsername());
             props.setProperty("dataSource.password", workConf.getDBPassword());
@@ -119,30 +123,57 @@ public void createDataSource() {
     }
 
     public final Connection makeConnection() throws SQLException {
-        YBClusterAwareDataSource ds = new YBClusterAwareDataSource();
-        ds.setProperty("user", workConf.getDBUsername());
-        ds.setProperty("password", workConf.getDBPassword());
-        ds.setProperty("reWriteBatchedInserts", "true");
-
-        if (workConf.getSslCert() != null && workConf.getSslCert().length() > 0) {
-          assert(workConf.getSslKey().length() > 0) : "The SSL key is empty.";
-          ds.setProperty("sslmode", "require");
-          ds.setProperty("sslcert", workConf.getSslCert());
-          ds.setProperty("sslkey", workConf.getSslKey());
+        Connection conn =  null;
+        if(workConf.getDBType().equals("yugabyte")) {
+            YBClusterAwareDataSource ds = new YBClusterAwareDataSource();
+            ds.setProperty("user", workConf.getDBUsername());
+            ds.setProperty("password", workConf.getDBPassword());
+            ds.setProperty("reWriteBatchedInserts", "true");
+
+            if (workConf.getSslCert() != null && workConf.getSslCert().length() > 0) {
+                assert (workConf.getSslKey().length() > 0) : "The SSL key is empty.";
+                ds.setProperty("sslmode", "require");
+                ds.setProperty("sslcert", workConf.getSslCert());
+                ds.setProperty("sslkey", workConf.getSslKey());
+            }
+            int r = dataSourceCounter.getAndIncrement() % workConf.getNodes().size();
+            String connectStr;
+            if (workConf.getJdbcURL() != null && workConf.getJdbcURL().length() > 0) {
+                connectStr = workConf.getJdbcURL();
+            } else {
+                connectStr = String.format("jdbc:yugabytedb://%s:%d/%s",
+                        workConf.getNodes().get(r),
+                        workConf.getPort(),
+                        workConf.getDBName());
+            }
+            ds.setUrl(connectStr);
+            conn = ds.getConnection();
         }
+        else {
+            java.util.Properties props = new java.util.Properties();
+            props.put("user", workConf.getDBUsername());
+            props.put("password", workConf.getDBPassword());
+            props.put("reWriteBatchedInserts", "true");
 
-        int r = dataSourceCounter.getAndIncrement() % workConf.getNodes().size();
-        String connectStr;
-        if (workConf.getJdbcURL() != null && workConf.getJdbcURL().length()>0) {
-            connectStr=workConf.getJdbcURL();
-        } else {
-            connectStr = String.format("jdbc:yugabytedb://%s:%d/%s",
-                workConf.getNodes().get(r),
-                workConf.getPort(),
-                workConf.getDBName());
+            if (workConf.getSslCert() != null && workConf.getSslCert().length() > 0) {
+                assert (workConf.getSslKey().length() > 0) : "The SSL key is empty.";
+                props.put("sslmode", "require");
+                props.put("sslcert", workConf.getSslCert());
+                props.put("sslkey", workConf.getSslKey());
+            }
+            int r = dataSourceCounter.getAndIncrement() % workConf.getNodes().size();
+            String connectStr;
+            if (workConf.getJdbcURL() != null && workConf.getJdbcURL().length() > 0) {
+                connectStr = workConf.getJdbcURL();
+            } else {
+                connectStr = String.format("jdbc:postgresql://%s:%d/%s",
+                        workConf.getNodes().get(r),
+                        workConf.getPort(),
+                        workConf.getDBName());
+            }
+            conn = DriverManager.getConnection(connectStr, props);
         }
-        ds.setUrl(connectStr);
-        return ds.getConnection();
+        return conn;
     }
 
     private static final AtomicInteger dataSourceCounter = new AtomicInteger(0);
@@ -199,7 +230,7 @@ public final void createDatabase() {
         try {
             Connection conn = makeConnection();
             SchemaManager schemaManager = SchemaManagerFactory.getSchemaManager(workConf, conn);
-            schemaManager.create();
+            schemaManager.create(workConf.getDBType());
             conn.close();
         } catch (SQLException ex) {
             throw new RuntimeException(String.format("Unexpected error when trying to create the %s database", this.benchmarkName), ex);

src/com/oltpbenchmark/api/Worker.java

@@ -482,7 +482,8 @@ protected final ArrayList<Pair<TransactionExecutionState, TransactionStatus>> do
 
             conn = dataSource.getConnection();
             try {
-                conn.createStatement().execute("SET yb_enable_expression_pushdown to on");
+                if(wrkld.getDBType().equals("yugabyte"))
+                    conn.createStatement().execute("SET yb_enable_expression_pushdown to on");
                 if (next.getProcedureClass() != StockLevel.class) {
                     // In accordance with 2.8.2.3 of the TPCC spec, StockLevel should execute each query in its own Snapshot
                     // Isolation.

src/com/oltpbenchmark/schema/SchemaManager.java

@@ -15,13 +15,12 @@ public SchemaManager(Connection db_connection) {
         this.db_connection = db_connection;
     }
 
-    public abstract void create() throws SQLException;
-
+    public abstract void create(String dbType) throws SQLException;
     public abstract void enableForeignKeyConstraints() throws SQLException;
 
     public abstract void dropForeignKeyConstraints() throws SQLException;
 
-    protected abstract void createIndexes() throws SQLException;
+    protected abstract void createIndexes(String dbType) throws SQLException;
 
     public static Set<String> getTableNames() {
         return TPCCTableSchemas.tables.keySet();

src/com/oltpbenchmark/schema/SchemaManagerFactory.java

@@ -10,6 +10,6 @@ public class SchemaManagerFactory {
     public static SchemaManager getSchemaManager(WorkloadConfiguration workConf, Connection conn) {
         return workConf.getGeoPartitioningEnabled()
                 ? new GeoPartitionedSchemaManager(workConf.getGeoPartitioningPolicy(), conn)
-                : new DefaultSchemaManager(conn);
+                : new DefaultSchemaManager(conn, workConf.getDBType());
     }
 }