From 28dac4ecc73be0595955656936c11064279a52b7 Mon Sep 17 00:00:00 2001
From: Lukas Eder <lukas.eder@gmail.com>
Date: Mon, 19 Jun 2023 10:21:28 +0200
Subject: [PATCH] [jOOQ/jOOQ#15249] Add ExecuteContext.skipUpdateCounts to
 allow for setting the AbstractContext.skipUpdateCounts flag in an
 ExecuteListener

---
 .../main/java/org/jooq/ExecuteContext.java    | 26 +++++++++++++++++++
 .../java/org/jooq/impl/AbstractContext.java   |  2 +-
 .../org/jooq/impl/DefaultExecuteContext.java  | 12 ++++++++-
 .../org/jooq/impl/SimpleExecuteContext.java   | 10 +++++++
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/jOOQ/src/main/java/org/jooq/ExecuteContext.java b/jOOQ/src/main/java/org/jooq/ExecuteContext.java
index fc4dc5d830..23d333601b 100644
--- a/jOOQ/src/main/java/org/jooq/ExecuteContext.java
+++ b/jOOQ/src/main/java/org/jooq/ExecuteContext.java
@@ -48,6 +48,7 @@ import java.util.stream.Collector;
 import org.jooq.conf.Settings;
 import org.jooq.conf.StatementType;
 import org.jooq.exception.DataAccessException;
+import org.jooq.impl.CustomQueryPart;
 import org.jooq.impl.DSL;
 
 import org.jetbrains.annotations.NotNull;
@@ -242,6 +243,31 @@ public interface ExecuteContext extends Scope {
      */
     void sql(String sql);
 
+    /**
+     * The number of user defined update counts that are going to be skipped
+     * when a statement batch is executed.
+     * <p>
+     * This is in <em>addition</em> to any skips added during the rendering of a
+     * {@link QueryPart} by jOOQ's internals or by user-defined query parts,
+     * such as {@link CustomQueryPart}.
+     */
+    int skipUpdateCounts();
+
+    /**
+     * Override the number of update counts that are going to be skipped when a
+     * statement batch is executed.
+     * <p>
+     * This is in <em>addition</em> to any skips added during the rendering of a
+     * {@link QueryPart} by jOOQ's internals or by user-defined query parts,
+     * such as {@link CustomQueryPart}.
+     * <p>
+     * This may have no effect, if called at the wrong moment.
+     *
+     * @see ExecuteListener#renderEnd(ExecuteContext)
+     * @see ExecuteListener#prepareStart(ExecuteContext)
+     */
+    void skipUpdateCounts(int skip);
+
     /**
      * The generated SQL statements that are being executed in batch mode, or
      * empty if the query is unknown or if there was no SQL statement.
diff --git a/jOOQ/src/main/java/org/jooq/impl/AbstractContext.java b/jOOQ/src/main/java/org/jooq/impl/AbstractContext.java
index ae4d2e7fb6..534807a559 100644
--- a/jOOQ/src/main/java/org/jooq/impl/AbstractContext.java
+++ b/jOOQ/src/main/java/org/jooq/impl/AbstractContext.java
@@ -934,7 +934,7 @@ abstract class AbstractContext<C extends Context<C>> extends AbstractScope imple
 
     @Override
     public final int skipUpdateCounts() {
-        return skipUpdateCounts;
+        return skipUpdateCounts + (ctx != null ? ctx.skipUpdateCounts() : 0);
     }
 
     @Override
diff --git a/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java b/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java
index 19a88cd0a5..c90c392867 100644
--- a/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java
+++ b/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java
@@ -88,7 +88,6 @@ import org.jooq.conf.Settings;
 import org.jooq.tools.JooqLogger;
 import org.jooq.tools.jdbc.JDBCUtils;
 import org.jooq.tools.reflect.Reflect;
-import org.jooq.tools.reflect.ReflectException;
 
 import org.jetbrains.annotations.NotNull;
 
@@ -119,6 +118,7 @@ class DefaultExecuteContext implements ExecuteContext {
     private Query                                         query;
     private final Routine<?>                              routine;
     private String                                        sql;
+    private int                                           skipUpdateCounts;
 
     private final BatchMode                               batchMode;
     private Query[]                                       batchQueries;
@@ -545,6 +545,16 @@ class DefaultExecuteContext implements ExecuteContext {
         return sql;
     }
 
+    @Override
+    public final int skipUpdateCounts() {
+        return this.skipUpdateCounts;
+    }
+
+    @Override
+    public void skipUpdateCounts(int skip) {
+        this.skipUpdateCounts = skip;
+    }
+
     @Override
     public final String[] batchSQL() {
         return batchMode != BatchMode.NONE
diff --git a/jOOQ/src/main/java/org/jooq/impl/SimpleExecuteContext.java b/jOOQ/src/main/java/org/jooq/impl/SimpleExecuteContext.java
index b88416abbb..546ceaf077 100644
--- a/jOOQ/src/main/java/org/jooq/impl/SimpleExecuteContext.java
+++ b/jOOQ/src/main/java/org/jooq/impl/SimpleExecuteContext.java
@@ -150,6 +150,16 @@ final class SimpleExecuteContext extends AbstractScope implements ExecuteContext
         throw new UnsupportedOperationException("Not implemented");
     }
 
+    @Override
+    public final int skipUpdateCounts() {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
+    @Override
+    public final void skipUpdateCounts(int skip) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
     @Override
     public final String[] batchSQL() {
         throw new UnsupportedOperationException("Not implemented");