From 29db96c8c553ff0907f42af3060d08e7d540313b Mon Sep 17 00:00:00 2001
From: Sfurti-yb <78196231+Sfurti-yb@users.noreply.github.com>
Date: Thu, 4 Apr 2024 08:52:27 +0530
Subject: [PATCH] Incorporate option to execute write YSQL sample app workloads
 indefinitely with upserts (#93)

* Incorporate option to execute write YSQL sample app workloads indefinitely with upserts

* Removing on conflict clause from SQLInsets workload
---
 src/main/java/com/yugabyte/sample/apps/SqlDataLoad.java      | 1 +
 .../com/yugabyte/sample/apps/SqlForeignKeysAndJoins.java     | 5 ++++-
 .../com/yugabyte/sample/apps/SqlGeoPartitionedTable.java     | 4 +++-
 src/main/java/com/yugabyte/sample/apps/SqlInserts.java       | 1 +
 src/main/java/com/yugabyte/sample/apps/SqlSnapshotTxns.java  | 4 +++-
 5 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/yugabyte/sample/apps/SqlDataLoad.java b/src/main/java/com/yugabyte/sample/apps/SqlDataLoad.java
index 4ef1aac..899a9a2 100644
--- a/src/main/java/com/yugabyte/sample/apps/SqlDataLoad.java
+++ b/src/main/java/com/yugabyte/sample/apps/SqlDataLoad.java
@@ -225,6 +225,7 @@ private PreparedStatement getPreparedInsert() throws Exception {
 
             close(insConnection);
             insConnection = getPostgresConnection();
+            insConnection.createStatement().execute("set yb_enable_upsert_mode = true");
             preparedInsert = insConnection.prepareStatement(sb.toString());
         }
         return preparedInsert;
diff --git a/src/main/java/com/yugabyte/sample/apps/SqlForeignKeysAndJoins.java b/src/main/java/com/yugabyte/sample/apps/SqlForeignKeysAndJoins.java
index 0ba7a8d..c461b0e 100644
--- a/src/main/java/com/yugabyte/sample/apps/SqlForeignKeysAndJoins.java
+++ b/src/main/java/com/yugabyte/sample/apps/SqlForeignKeysAndJoins.java
@@ -218,7 +218,9 @@ private PreparedStatement getPreparedInsertUser() throws Exception {
         if (preparedInsertUser == null) {
           String stmt = String.format("INSERT INTO %s (user_id, user_details) VALUES (?, ?);", getTable1Name());
           if (preparedInsertUser == null) {
-            preparedInsertUser = getPostgresConnection().prepareStatement(stmt);
+            Connection connection = getPostgresConnection();
+            connection.createStatement().execute("set yb_enable_upsert_mode = true");
+            preparedInsertUser = connection.prepareStatement(stmt);
           }
         }
         preparedInsertUserLocal = preparedInsertUser;
@@ -236,6 +238,7 @@ private PreparedStatement getPreparedInsertUserOrder() throws Exception {
                                       getTable2Name());
           if (preparedInsertOrder == null) {
             Connection connection = getPostgresConnection();
+            connection.createStatement().execute("set yb_enable_upsert_mode = true");
             connection.setTransactionIsolation(Connection.TRANSACTION_SERIALIZABLE);
             preparedInsertOrder = connection.prepareStatement(stmt);
           }
diff --git a/src/main/java/com/yugabyte/sample/apps/SqlGeoPartitionedTable.java b/src/main/java/com/yugabyte/sample/apps/SqlGeoPartitionedTable.java
index b5c935d..4e6abd7 100644
--- a/src/main/java/com/yugabyte/sample/apps/SqlGeoPartitionedTable.java
+++ b/src/main/java/com/yugabyte/sample/apps/SqlGeoPartitionedTable.java
@@ -202,7 +202,9 @@ public long doRead() {
 
   private PreparedStatement getPreparedInsert() throws Exception {
     if (preparedInsert == null) {
-      preparedInsert = getPostgresConnection().prepareStatement(
+      Connection connection = getPostgresConnection();
+      connection.createStatement().execute("set yb_enable_upsert_mode = true");
+      preparedInsert = connection.prepareStatement(
           String.format("INSERT INTO %s (region, k, v) VALUES (?, ?, ?)",
                         getTableName()));
     }
diff --git a/src/main/java/com/yugabyte/sample/apps/SqlInserts.java b/src/main/java/com/yugabyte/sample/apps/SqlInserts.java
index 63640d5..694ada4 100644
--- a/src/main/java/com/yugabyte/sample/apps/SqlInserts.java
+++ b/src/main/java/com/yugabyte/sample/apps/SqlInserts.java
@@ -151,6 +151,7 @@ private PreparedStatement getPreparedInsert() throws Exception {
     if (preparedInsert == null) {
       close(insConnection);
       insConnection = getPostgresConnection();
+      insConnection.createStatement().execute("set yb_enable_upsert_mode = true");
       preparedInsert = insConnection.prepareStatement(
           String.format("INSERT INTO %s (k, v) VALUES (?, ?);", getTableName()));
     }
diff --git a/src/main/java/com/yugabyte/sample/apps/SqlSnapshotTxns.java b/src/main/java/com/yugabyte/sample/apps/SqlSnapshotTxns.java
index 9c62893..8d2041f 100644
--- a/src/main/java/com/yugabyte/sample/apps/SqlSnapshotTxns.java
+++ b/src/main/java/com/yugabyte/sample/apps/SqlSnapshotTxns.java
@@ -148,7 +148,9 @@ private PreparedStatement getPreparedInsert() throws Exception {
                   String.format("INSERT INTO %s (k, v) VALUES (?, ?);", getTableName()) +
                   "COMMIT;";
     if (preparedInsert == null) {
-      preparedInsert = getPostgresConnection().prepareStatement(stmt);
+      Connection connection = getPostgresConnection();
+      connection.createStatement().execute("set yb_enable_upsert_mode = true");
+      preparedInsert = connection.prepareStatement(stmt);
     }
     return preparedInsert;
   }