Make UPDATE_TIME unique for MIGRATION_MODEL table

Closes #40088

Signed-off-by: Martin Bartoš <mabartos@redhat.com>
Signed-off-by: Alexander Schwartz <aschwart@redhat.com>
Co-authored-by: Alexander Schwartz <aschwart@redhat.com>

model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/custom/JpaUpdate26_2_6_RemoveDuplicateMigrationModelTime.java

@@ -0,0 +1,84 @@
+package org.keycloak.connections.jpa.updater.liquibase.custom;
+
+import liquibase.exception.CustomChangeException;
+import liquibase.statement.core.DeleteStatement;
+import liquibase.structure.core.Column;
+import org.keycloak.migration.ModelVersion;
+
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+/**
+ * Cleanup script for removing duplicated migration model update time in the MIGRATION_MODEL table
+ * See: <a href="https://github.com/keycloak/keycloak/issues/40088">keycloak#40088</a>
+ */
+public class JpaUpdate26_2_6_RemoveDuplicateMigrationModelTime extends CustomKeycloakTask {
+    private final static String MIGRATION_MODEL_TABLE = "MIGRATION_MODEL";
+
+    @Override
+    protected String getTaskId() {
+        return "Delete duplicated records for DB update time in MIGRATION_MODEL table";
+    }
+
+    @Override
+    protected void generateStatementsImpl() throws CustomChangeException {
+        final Map<String, ModelVersion> itemsToDelete = new HashMap<>();
+
+        final String tableName = getTableName(MIGRATION_MODEL_TABLE);
+        final String colId = database.correctObjectName("ID", Column.class);
+        final String colVersion = database.correctObjectName("VERSION", Column.class);
+        final String colUpdateTime = database.correctObjectName("UPDATE_TIME", Column.class);
+
+        final String GET_DUPLICATED_RECORDS = """
+                SELECT m1.%s, m1.%s
+                FROM %s m1
+                WHERE EXISTS (
+                    SELECT m2.%s
+                    FROM %s m2
+                    WHERE m2.%s = m1.%s AND m2.%s <> m1.%s
+                )
+                """.formatted(
+                colId, colVersion,  // SELECT m1.%s, m1.%s  =>  SELECT m1.ID, m1.VERSION
+                tableName,          // FROM %s m1           =>  FROM MIGRATION_MODEL m1
+                colId,              // SELECT m2.%s         =>  SELECT m2.ID
+                tableName,          // FROM %s m2           =>  FROM MIGRATION_MODEL m2
+                // WHERE m2.%s = m1.%s AND m2.%s <> m1.%s   =>  WHERE m2.UPDATE_TIME = m1.UPDATE_TIME AND m2.ID <> m1.ID
+                colUpdateTime, colUpdateTime, colId, colId
+        );
+
+        //noinspection SqlSourceToSinkFlow
+        try (PreparedStatement ps = connection.prepareStatement(GET_DUPLICATED_RECORDS)) {
+            ResultSet resultSet = ps.executeQuery();
+            while (resultSet.next()) {
+                String id = resultSet.getString(1);
+                ModelVersion version = new ModelVersion(resultSet.getString(2));
+                itemsToDelete.put(id, version);
+            }
+        } catch (Exception e) {
+            throw new CustomChangeException(getTaskId() + ": Failed to detect duplicate MIGRATION_MODEL rows", e);
+        }
+
+        // Get ID of the highest Keycloak version with the same update time
+        var highestVersionId = itemsToDelete.entrySet()
+                .stream()
+                .reduce((e1, e2) -> e1.getValue().lessThan(e2.getValue()) ? e2 : e1)
+                .map(Map.Entry::getKey)
+                .orElse(null);
+
+        AtomicInteger i = new AtomicInteger();
+        itemsToDelete.keySet().stream()
+                .filter(f -> !f.equals(highestVersionId))
+                .collect(Collectors.groupingByConcurrent(id -> i.getAndIncrement() / 20, Collectors.toList())) // Split into chunks of at most 20 items
+                .values().stream()
+                .map(ids -> new DeleteStatement(null, null, MIGRATION_MODEL_TABLE)
+                        .setWhere(":name IN (" + ids.stream().map(id -> "?").collect(Collectors.joining(",")) + ")")
+                        .addWhereColumnName(colId)
+                        .addWhereParameters(ids.toArray())
+                )
+                .forEach(statements::add);
+    }
+}

model/jpa/src/main/resources/META-INF/jpa-changelog-26.2.6.xml

@@ -25,4 +25,12 @@
         <addUniqueConstraint tableName="MIGRATION_MODEL" columnNames="VERSION" constraintName="UK_MIGRATION_VERSION"/>
     </changeSet>
 
+    <changeSet author="keycloak" id="26.2.6-40088-duplicate">
+        <customChange class="org.keycloak.connections.jpa.updater.liquibase.custom.JpaUpdate26_2_6_RemoveDuplicateMigrationModelTime"/>
+    </changeSet>
+
+    <changeSet author="keycloak" id="26.2.6-40088-uk">
+        <addUniqueConstraint tableName="MIGRATION_MODEL" columnNames="UPDATE_TIME" constraintName="UK_MIGRATION_UPDATE_TIME"/>
+    </changeSet>
+
 </databaseChangeLog>

testsuite/model/src/test/java/org/keycloak/testsuite/model/MigrationModelTest.java

@@ -159,6 +159,47 @@ public class MigrationModelTest extends KeycloakModelTest {
         });
     }
 
+    @Test
+    public void duplicatedUpdateTime() {
+        inComittedTransaction(1, (session, i) -> {
+            String currentVersion = new ModelVersion(Version.VERSION).toString();
+            EntityManager em = session.getProvider(JpaConnectionProvider.class).getEntityManager();
+
+            List<MigrationModelEntity> entities = getMigrationEntities(em);
+            assertThat(entities.size(), is(1));
+            assertMigrationModelEntity(entities.get(0), currentVersion);
+
+            MigrationModel m = session.getProvider(DeploymentStateProvider.class).getMigrationModel();
+            assertThat(m.getStoredVersion(), is(currentVersion));
+            assertThat(entities.get(0).getId(), is(m.getResourcesTag()));
+
+            try {
+                MigrationModelEntity mm1 = new MigrationModelEntity();
+                mm1.setId("a");
+                mm1.setUpdatedTime(0);
+                mm1.setVersion("26.0.0");
+                em.persist(mm1);
+
+                em.flush();
+
+                // Same time, everything different - testing for the constraint to be present
+                MigrationModelEntity mm2 = new MigrationModelEntity();
+                mm2.setId("b");
+                mm2.setUpdatedTime(0);
+                mm2.setVersion("26.0.1");
+                em.persist(mm2);
+
+                // added at the same time - exception thrown by the unique constraint
+                assertThrows(ModelDuplicateException.class, em::flush);
+
+            } finally {
+                em.remove(em.find(MigrationModelEntity.class, "a"));
+            }
+
+            return null;
+        });
+    }
+
     private void assertMigrationModelEntity(MigrationModelEntity model, String expectedVersion) {
         assertThat(model, notNullValue());
         assertTrue(model.getId().matches("[\\da-z]{5}"));