From b9a1c6ead0daa7691e063ee6b14a8f58f2aeca6a Mon Sep 17 00:00:00 2001
From: Lukas Eder <lukas.eder@gmail.com>
Date: Tue, 16 Jun 2020 16:26:01 +0200
Subject: [PATCH] [jOOQ/jOOQ#6736] Add support for PostgreSQL ON CONFLICT ..
 WHERE <index_predicate>

---
 .../java/org/jooq/InsertOnDuplicateStep.java  |  4 +--
 jOOQ/src/main/java/org/jooq/InsertQuery.java  | 11 ++++++
 .../main/java/org/jooq/impl/InsertImpl.java   | 34 +++++++++++++++----
 .../java/org/jooq/impl/InsertQueryImpl.java   | 13 +++++++
 4 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/jOOQ/src/main/java/org/jooq/InsertOnDuplicateStep.java b/jOOQ/src/main/java/org/jooq/InsertOnDuplicateStep.java
index effa3d9f53..b82b15968f 100644
--- a/jOOQ/src/main/java/org/jooq/InsertOnDuplicateStep.java
+++ b/jOOQ/src/main/java/org/jooq/InsertOnDuplicateStep.java
@@ -122,7 +122,7 @@ public interface InsertOnDuplicateStep<R extends Record> extends InsertReturning
      * <code>MERGE</code>, if table meta data is available.
      */
     @Support
-    InsertOnConflictDoUpdateStep<R> onConflict(Field<?>... keys);
+    InsertOnConflictWhereIndexPredicateStep<R> onConflict(Field<?>... keys);
 
     /**
      * Add an <code>ON CONFLICT</code> clause to this INSERT statement.
@@ -132,7 +132,7 @@ public interface InsertOnDuplicateStep<R extends Record> extends InsertReturning
      * <code>MERGE</code>, if table meta data is available.
      */
     @Support
-    InsertOnConflictDoUpdateStep<R> onConflict(Collection<? extends Field<?>> keys);
+    InsertOnConflictWhereIndexPredicateStep<R> onConflict(Collection<? extends Field<?>> keys);
 
     /**
      * Add an <code>ON CONFLICT DO NOTHING</code> clause to this INSERT statement.
diff --git a/jOOQ/src/main/java/org/jooq/InsertQuery.java b/jOOQ/src/main/java/org/jooq/InsertQuery.java
index 4d5086b924..0b0221db81 100644
--- a/jOOQ/src/main/java/org/jooq/InsertQuery.java
+++ b/jOOQ/src/main/java/org/jooq/InsertQuery.java
@@ -251,6 +251,17 @@ public interface InsertQuery<R extends Record> extends StoreQuery<R>, Insert<R>,
     @Support({ CUBRID, DERBY, H2, HSQLDB, MARIADB, MYSQL, POSTGRES, SQLITE })
     void addValuesForUpdate(Map<?, ?> map);
 
+    /**
+     * Adds a new condition the {@link #onConflict(Field...)} clause.
+     * <p>
+     * This is for use with {@link SQLDialect#POSTGRES}'s
+     * {@link #onConflict(Field...)} clause.
+     *
+     * @param condition The condition
+     */
+    @Support({ POSTGRES, SQLITE })
+    void onConflictWhere(Condition condition);
+
     /**
      * Adds a new condition to the query, connecting it to existing conditions
      * with {@link Operator#AND}.
diff --git a/jOOQ/src/main/java/org/jooq/impl/InsertImpl.java b/jOOQ/src/main/java/org/jooq/impl/InsertImpl.java
index fb187c10df..504ba39c0f 100644
--- a/jOOQ/src/main/java/org/jooq/impl/InsertImpl.java
+++ b/jOOQ/src/main/java/org/jooq/impl/InsertImpl.java
@@ -55,7 +55,7 @@ import org.jooq.Constraint;
 import org.jooq.Field;
 import org.jooq.FieldLike;
 import org.jooq.InsertOnConflictConditionStep;
-import org.jooq.InsertOnConflictDoUpdateStep;
+import org.jooq.InsertOnConflictWhereIndexPredicateStep;
 import org.jooq.InsertOnDuplicateSetMoreStep;
 import org.jooq.InsertQuery;
 import org.jooq.InsertResultStep;
@@ -122,7 +122,7 @@ import org.jooq.UniqueKey;
  * @author Lukas Eder
  */
 @SuppressWarnings({ "rawtypes", "unchecked" })
-class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20, T21, T22>
+final class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20, T21, T22>
     extends AbstractDelegatingRowCountQuery<InsertQuery<R>>
     implements
 
@@ -153,7 +153,7 @@ class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11,
     InsertSetStep<R>,
     InsertSetMoreStep<R>,
     InsertOnDuplicateSetMoreStep<R>,
-    InsertOnConflictDoUpdateStep<R>,
+    InsertOnConflictWhereIndexPredicateStep<R>,
     InsertOnConflictConditionStep<R>,
     InsertResultStep<R> {
 
@@ -167,6 +167,11 @@ class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11,
     private boolean           onDuplicateKeyUpdate;
     private boolean           returningResult;
 
+    /**
+     * Whether {@link #where(Condition)} adds conditions to the <code>DO UPDATE/code> clause.
+     */
+    private transient boolean doUpdateWhere;
+
     InsertImpl(Configuration configuration, WithImpl with, Table<R> into) {
         super(new InsertQueryImpl<>(configuration, with, into));
 
@@ -622,11 +627,13 @@ class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11,
 
     @Override
     public final InsertImpl doUpdate() {
+        doUpdateWhere = true;
         return onDuplicateKeyUpdate();
     }
 
     @Override
     public final InsertImpl doNothing() {
+        doUpdateWhere = true;
         return onDuplicateKeyIgnore();
     }
 
@@ -661,12 +668,14 @@ class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11,
 
     @Override
     public final InsertImpl onConflictDoNothing() {
+        doUpdateWhere = true;
         onConflict().doNothing();
         return this;
     }
 
     @Override
     public final InsertImpl onDuplicateKeyUpdate() {
+        doUpdateWhere = true;
         onDuplicateKeyUpdate = true;
         getDelegate().onDuplicateKeyUpdate(true);
         return this;
@@ -674,6 +683,7 @@ class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11,
 
     @Override
     public final InsertImpl onDuplicateKeyIgnore() {
+        doUpdateWhere = true;
         getDelegate().onDuplicateKeyIgnore(true);
         return this;
     }
@@ -827,19 +837,31 @@ class InsertImpl<R extends Record, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11,
 
     @Override
     public final InsertImpl where(Condition condition) {
-        getDelegate().addConditions(condition);
+        if (doUpdateWhere)
+            getDelegate().addConditions(condition);
+        else
+            getDelegate().onConflictWhere(condition);
+
         return this;
     }
 
     @Override
     public final InsertImpl where(Condition... conditions) {
-        getDelegate().addConditions(conditions);
+        if (doUpdateWhere)
+            getDelegate().addConditions(conditions);
+        else
+            getDelegate().onConflictWhere(DSL.and(conditions));
+
         return this;
     }
 
     @Override
     public final InsertImpl where(Collection<? extends Condition> conditions) {
-        getDelegate().addConditions(conditions);
+        if (doUpdateWhere)
+            getDelegate().addConditions(conditions);
+        else
+            getDelegate().onConflictWhere(DSL.and(conditions));
+
         return this;
     }
 
diff --git a/jOOQ/src/main/java/org/jooq/impl/InsertQueryImpl.java b/jOOQ/src/main/java/org/jooq/impl/InsertQueryImpl.java
index 614eb0f49a..f3a3886082 100644
--- a/jOOQ/src/main/java/org/jooq/impl/InsertQueryImpl.java
+++ b/jOOQ/src/main/java/org/jooq/impl/InsertQueryImpl.java
@@ -134,6 +134,7 @@ final class InsertQueryImpl<R extends Record> extends AbstractStoreQuery<R> impl
     private Constraint                   onConstraint;
     private UniqueKey<R>                 onConstraintUniqueKey;
     private QueryPartList<Field<?>>      onConflict;
+    private final ConditionProviderImpl  onConflictWhere;
     private final ConditionProviderImpl  condition;
 
     InsertQueryImpl(Configuration configuration, WithImpl with, Table<R> into) {
@@ -141,6 +142,7 @@ final class InsertQueryImpl<R extends Record> extends AbstractStoreQuery<R> impl
 
         this.updateMap = new FieldMapForUpdate(into, INSERT_ON_DUPLICATE_KEY_UPDATE_ASSIGNMENT);
         this.insertMaps = new FieldMapsForInsert(into);
+        this.onConflictWhere = new ConditionProviderImpl();
         this.condition = new ConditionProviderImpl();
     }
 
@@ -178,6 +180,11 @@ final class InsertQueryImpl<R extends Record> extends AbstractStoreQuery<R> impl
         this.onConflict = new QueryPartList<>(fields);
     }
 
+    @Override
+    public final void onConflictWhere(Condition conditions) {
+        onConflictWhere.addConditions(conditions);
+    }
+
     @Override
     public final void onConflictOnConstraint(Constraint constraint) {
         onConflictOnConstraint0(constraint);
@@ -392,6 +399,12 @@ final class InsertQueryImpl<R extends Record> extends AbstractStoreQuery<R> impl
                         ctx.sql(')');
                     }
 
+                    if (onConflictWhere.hasWhere())
+                        ctx.sql(' ')
+                           .visit(K_WHERE)
+                           .sql(' ')
+                           .visit(onConflictWhere.getWhere());
+
                     ctx.sql(' ')
                        .visit(K_DO_UPDATE)
                        .formatSeparator()