fix: batch don't execute all queries with pgJdbc driver (#70)

* feat: add java-kotlin support * feat: refactoring * feat: move javaapi scala classes * fix: change import JdbcDsl * fix: fix Dependencies and readme * fix: fix Dependencies * fix: batch don't execute all queries with pgJdbc driver * feat: add batch test for all queries * feat: add batch check queries for java and kotlin --------- Co-authored-by: a.v.muratova <[email protected]>
Tinkoff · Mar 7, 2023 · dc27145 · dc27145
1 parent 3ef4b21
commit dc27145
Show file tree

Hide file tree

Showing 10 changed files with 129 additions and 36 deletions.
diff --git a/src/main/scala/ru/tinkoff/load/jdbc/db/JDBCClient.scala b/src/main/scala/ru/tinkoff/load/jdbc/db/JDBCClient.scala
@@ -56,7 +56,6 @@ class JDBCClient(pool: HikariDataSource, blockingPool: ExecutorService) {
                         _ <- s.close
                       } yield (),
                     )
-
     } yield stmt
 
   private def statementResource: ResourceFut[StatementWrapper[Future]] =
@@ -116,19 +115,13 @@ class JDBCClient(pool: HikariDataSource, blockingPool: ExecutorService) {
   ): Unit =
     withCompletion(callableStatementResource(sqlCall, params.toMap, outParams.toMap).use(_.executeUpdate))(s, f)
 
-  def batch[U](queries: Seq[SqlWithParam], batchSize: Int = 1000)(s: Array[Int] => U, f: Throwable => U): Unit =
+  def batch[U](queries: Seq[SqlWithParam])(s: Array[Int] => U, f: Throwable => U): Unit =
     withCompletion(
       statementForBatchResource.use(stmt =>
         queries
-          .map(_.substituteParams)
-          .grouped(batchSize)
-          .map(batch =>
-            batch
-              .map(query => stmt.addBatch(query))
-              .reduce((f1, f2) => f1.flatMap(_ => f2))
-              .flatMap(_ => stmt.executeBatch),
-          )
-          .reduce((f1, f2) => f1.flatMap(a1 => f2.map(a2 => a1 ++ a2))),
+          .map(query => stmt.addBatch(query.substituteParams))
+          .reduce((f1, f2) => f1.flatMap(_ => f2))
+          .flatMap(_ => stmt.executeBatch),
       ),
     )(s, f)
 

diff --git a/src/test/java/GatlingRunner.scala b/src/test/java/GatlingRunner.scala
@@ -0,0 +1,18 @@
+import io.gatling.app.Gatling
+import io.gatling.core.config.GatlingPropertiesBuilder
+import ru.tinkoff.load.jdbc.test.JdbcDebugTest;
+
+object GatlingRunner {
+
+  def main(args: Array[String]): Unit = {
+
+    // this is where you specify the class you want to run
+    val simulationClass = classOf[JdbcDebugTest].getName
+
+    val props = new GatlingPropertiesBuilder
+    props.simulationClass(simulationClass)
+
+    Gatling.fromMap(props.build)
+  }
+
+}
diff --git a/src/test/java/ru/tinkoff/load/jdbc/test/cases/JdbcActions.java b/src/test/java/ru/tinkoff/load/jdbc/test/cases/JdbcActions.java
@@ -4,6 +4,8 @@
 import ru.tinkoff.load.javaapi.check.simpleCheckType;
 
 import java.util.Map;
+import java.util.UUID;
+
 import static ru.tinkoff.load.javaapi.JdbcDsl.*;
 
 public class JdbcActions {
@@ -18,10 +20,11 @@ public static RawSqlActionBuilder createprocedure(){
         return jdbc("Procedure create")
                 .rawSql(
                         """
-                                CREATE OR REPLACE PROCEDURE TEST_PROCEDURE(p1 varchar(10), p2 integer) 
-                                LANGUAGE SQL 
-                                AS $$ 
-                                    select 1 as result;
+                                CREATE ALIAS TEST_PROCEDURE AS $$
+                                String testProcedure(String p1, Long p2) {
+                                    String suf = p1 + "test";
+                                    return p2.toString() + suf;
+                                }
                                 $$;
                                 """
                 );
@@ -45,7 +48,8 @@ public static BatchActionBuilder batchTest(){
                         .values(Map.of("ID", 20, "NAME", "Test 12", "FLAG", true)),
                 insetInto("TEST_TABLE", "ID", "NAME")
                         .values(Map.of("ID", 40, "NAME", "Test 34")),
-                update("TEST_TABLE").set(Map.of("NAME", "Test5")).where("ID = 2")
+                update("TEST_TABLE").set(Map.of("NAME", "Test5")).where("ID = 2"),
+                insetInto("TT", "ID", "NAME").values(Map.of("ID", UUID.randomUUID(), "NAME", "OOO342ff"))
         );
     }
 
@@ -58,7 +62,7 @@ public static QueryActionBuilder selectTT(){
     public static QueryActionBuilder selectTest(){
         return jdbc("SELECT TEST")
                 .queryP("SELECT * FROM TEST_TABLE WHERE ID = {id}")
-                .params(Map.of("id", 1))
+                .params(Map.of("id", 20))
                 .check(simpleCheck(simpleCheckType.NonEmpty),
                         allResults().saveAs("R"));
     }
@@ -68,4 +72,11 @@ public static QueryActionBuilder selectAfterBatch(){
                 .query("SELECT * FROM TEST_TABLE")
                 .check(allResults().saveAs("RR"));
     }
+
+    public static QueryActionBuilder checkTestTableAfterBatch(){
+        return jdbc("Check TEST_TABLE")
+                .query("SELECT * FROM TEST_TABLE WHERE EXISTS(SELECT NAME FROM TEST_TABLE WHERE ID = 2 AND NAME = 'Test5')" +
+                        "AND (SELECT COUNT(ID) FROM TEST_TABLE WHERE ID IN (20, 40, 2)) = 3")
+                .check(simpleCheck(simpleCheckType.NonEmpty));
+    }
 }
diff --git a/src/test/java/ru/tinkoff/load/jdbc/test/scenarios/JdbcBasicSimulation.java b/src/test/java/ru/tinkoff/load/jdbc/test/scenarios/JdbcBasicSimulation.java
@@ -2,9 +2,18 @@
 
 import io.gatling.javaapi.core.ScenarioBuilder;
 import static io.gatling.javaapi.core.CoreDsl.scenario;
-import static ru.tinkoff.load.jdbc.test.cases.JdbcActions.createTable;
+import static ru.tinkoff.load.jdbc.test.cases.JdbcActions.*;
 
 public class JdbcBasicSimulation {
     public static ScenarioBuilder scn = scenario("JDBC scenario")
-            .exec(createTable());
+            .exec(createTable())
+            .exec(createprocedure())
+            .exec(insertTest())
+            .exec(callTest())
+            .exec(batchTest())
+            .exec(selectTT())
+            .exec(selectTest())
+            .exec(selectAfterBatch())
+            .exec(checkTestTableAfterBatch())
+            ;
 }
diff --git a/src/test/kotlin/GatlingRunner.scala b/src/test/kotlin/GatlingRunner.scala
@@ -0,0 +1,18 @@
+import io.gatling.app.Gatling
+import io.gatling.core.config.GatlingPropertiesBuilder
+import ru.tinkoff.load.jdbc.test.KtJdbcDebugTest;
+
+object GatlingRunner {
+
+  def main(args: Array[String]): Unit = {
+
+    // this is where you specify the class you want to run
+    val simulationClass = classOf[KtJdbcDebugTest].getName
+
+    val props = new GatlingPropertiesBuilder
+    props.simulationClass(simulationClass)
+
+    Gatling.fromMap(props.build)
+  }
+
+}
diff --git a/src/test/kotlin/ru/tinkoff/load/jdbc/test/KtJdbcDebugTest.kt b/src/test/kotlin/ru/tinkoff/load/jdbc/test/KtJdbcDebugTest.kt
@@ -1,13 +1,13 @@
 package ru.tinkoff.load.jdbc.test
 
-import io.gatling.javaapi.core.CoreDsl.*
-import io.gatling.javaapi.core.*
+import io.gatling.javaapi.core.OpenInjectionStep.atOnceUsers
+import io.gatling.javaapi.core.Simulation
 import ru.tinkoff.load.jdbc.test.scenarios.KtJdbcBasicSimulation.scn
 
-class KtJdbcDebugTest : Simulation()  {
-    init{
-            setUp(
-                    scn.injectOpen(atOnceUsers(1))
-            ).protocols(KtJdbcProtocol.dataBase)
+class KtJdbcDebugTest : Simulation() {
+    init {
+        setUp(
+            scn.injectOpen(atOnceUsers(1))
+        ).protocols(KtJdbcProtocol.dataBase)
     }
-}
+}
diff --git a/src/test/kotlin/ru/tinkoff/load/jdbc/test/cases/KtJdbcActions.kt b/src/test/kotlin/ru/tinkoff/load/jdbc/test/cases/KtJdbcActions.kt
@@ -1,9 +1,10 @@
 package ru.tinkoff.load.jdbc.test.cases
 
 import ru.tinkoff.load.javaapi.JdbcDsl.*
-import ru.tinkoff.load.javaapi.check.simpleCheckType.*
 import ru.tinkoff.load.javaapi.actions.*
 import ru.tinkoff.load.javaapi.check.simpleCheckType
+import java.util.*
+import java.util.Map
 
 object KtJdbcActions {
     fun createTable(): RawSqlActionBuilder {
@@ -15,11 +16,12 @@ object KtJdbcActions {
     fun createprocedure(): RawSqlActionBuilder {
         return jdbc("Procedure create")
                 .rawSql("""
-                    CREATE OR REPLACE PROCEDURE TEST_PROCEDURE(p1 varchar(10), p2 integer)
-                    LANGUAGE SQL
-                    AS $$
-                        select 1 as result;
-                    $$;
+                     CREATE ALIAS TEST_PROCEDURE AS ${'$'}${'$'}
+                                String testProcedure(String p1, Long p2) {
+                                    String suf = p1 + "test";
+                                    return p2.toString() + suf;
+                                }
+                                ${'$'}${'$'};
                 """
                 )
     }
@@ -42,7 +44,8 @@ object KtJdbcActions {
                         .values(mapOf("ID" to 20, "NAME" to "Test 12", "FLAG" to true)),
                 insetInto("TEST_TABLE", "ID", "NAME")
                         .values(mapOf("ID" to 40, "NAME" to "Test 34")),
-                update("TEST_TABLE").set(mapOf("NAME" to "Test5")).where("ID = 2")
+                update("TEST_TABLE").set(mapOf("NAME" to "Test5")).where("ID = 2"),
+                insetInto("TT", "ID", "NAME").values(mapOf("ID" to UUID.randomUUID(), "NAME" to "OOO342ff"))
         )
     }
 
@@ -55,7 +58,7 @@ object KtJdbcActions {
     fun selectTest(): QueryActionBuilder {
         return jdbc("SELECT TEST")
                 .queryP("SELECT * FROM TEST_TABLE WHERE ID = {id}")
-                .params(mapOf("id" to 1))
+                .params(mapOf("id" to 20))
                 .check(simpleCheck(simpleCheckType.NonEmpty),
                         allResults().saveAs("R"))
     }
@@ -65,4 +68,13 @@ object KtJdbcActions {
                 .query("SELECT * FROM TEST_TABLE")
                 .check(allResults().saveAs("RR"))
     }
+
+    fun checkTestTableAfterBatch(): QueryActionBuilder {
+        return jdbc("Check TEST_TABLE")
+            .query(
+                "SELECT * FROM TEST_TABLE WHERE EXISTS(SELECT NAME FROM TEST_TABLE WHERE ID = 2 AND NAME = 'Test5')" +
+                        "AND (SELECT COUNT(ID) FROM TEST_TABLE WHERE ID IN (20, 40, 2)) = 3"
+            )
+            .check(simpleCheck(simpleCheckType.NonEmpty))
+    }
 }
diff --git a/src/test/kotlin/ru/tinkoff/load/jdbc/test/scenarios/KtJdbcBasicSimulation.kt b/src/test/kotlin/ru/tinkoff/load/jdbc/test/scenarios/KtJdbcBasicSimulation.kt
@@ -4,7 +4,18 @@ import io.gatling.javaapi.core.CoreDsl.scenario
 import io.gatling.javaapi.core.ScenarioBuilder
 import ru.tinkoff.load.jdbc.test.cases.KtJdbcActions
 
+
+
 object KtJdbcBasicSimulation {
     var scn: ScenarioBuilder = scenario("JDBC scenario")
-            .exec(KtJdbcActions.createTable())
+        .exec(KtJdbcActions.createTable())
+        .exec(KtJdbcActions.createprocedure())
+        .exec(KtJdbcActions.insertTest())
+        .exec(KtJdbcActions.callTest())
+        .exec(KtJdbcActions.batchTest())
+        .exec(KtJdbcActions.selectTest())
+        .exec(KtJdbcActions.selectTT())
+        .exec(KtJdbcActions.selectAfterBatch())
+        .exec(KtJdbcActions.checkTestTableAfterBatch())
+
 }
diff --git a/src/test/scala/ru/tinkoff/load/jdbc/test/cases/Actions.scala b/src/test/scala/ru/tinkoff/load/jdbc/test/cases/Actions.scala
@@ -64,4 +64,23 @@ object Actions {
         allResults.saveAs("RR"),
       )
 
+  def checkBatchTestTable: QueryActionBuilder =
+    jdbc("SELECT TEST_TABLE AFTER BATCH")
+      .query("""SELECT * FROM TEST_TABLE
+          |WHERE ID IN (20, 30, 40, 2)
+          |AND EXISTS(SELECT NAME FROM TEST_TABLE
+          |WHERE ID=2 AND NAME = 'bird')
+          |""".stripMargin)
+      .check(
+        simpleCheck(x => x.length == 4),
+      )
+
+  def checkBatchTT: QueryActionBuilder =
+    jdbc("SELECT TT AFTER BATCH")
+      .query("""SELECT * FROM TT
+          |WHERE NAME = 'OOO342ff'
+          |""".stripMargin)
+      .check(
+        simpleCheck(x => x.length == 1),
+      )
 }
diff --git a/src/test/scala/ru/tinkoff/load/jdbc/test/scenarios/BasicSimulation.scala b/src/test/scala/ru/tinkoff/load/jdbc/test/scenarios/BasicSimulation.scala
@@ -34,5 +34,7 @@ class BasicSimulation {
       println(s("RR").as[List[Map[String, Any]]])
       s
     }
+    .exec(Actions.checkBatchTestTable)
+    .exec(Actions.checkBatchTT)
 
 }