Add transaction tests to model tests

Closes: #15890

model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapStorageProviderFactory.java

@@ -21,6 +21,7 @@ import static org.keycloak.models.map.storage.jpa.updater.MapJpaUpdaterProvider.
 import java.sql.Connection;
 import java.sql.DatabaseMetaData;
 import java.sql.DriverManager;
+import java.sql.PreparedStatement;
 import java.sql.SQLException;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -165,6 +166,7 @@ public class JpaMapStorageProviderFactory implements
     private Config.Scope config;
     private final String sessionProviderKey;
     private final String sessionTxKey;
+    private String databaseShortName;
 
     // Object instances for each single JpaMapStorageProviderFactory instance per model type.
     // Used to synchronize on when validating the model type area.
@@ -278,7 +280,23 @@ public class JpaMapStorageProviderFactory implements
     }
 
     protected EntityManager getEntityManager() {
-        return emf.createEntityManager();
+        EntityManager em = emf.createEntityManager();
+
+        // This is a workaround for Hibernate not supporting javax.persistence.lock.timeout
+        //   config option for Postgresql/CockroachDB - https://hibernate.atlassian.net/browse/HHH-16071
+        if ("postgresql".equals(databaseShortName) || "cockroachdb".equals(databaseShortName)) {
+            Long lockTimeout = config.getLong("lockTimeout");
+            if (lockTimeout != null) {
+                em.unwrap(SessionImpl.class)
+                        .doWork(connection -> {
+                            PreparedStatement preparedStatement = connection.prepareStatement("SET LOCAL lock_timeout = '" + lockTimeout + "';");
+                            preparedStatement.execute();
+                        });
+            } else {
+                logger.warnf("Database %s used without lockTimeout option configured. This can result in deadlock where one connection waits for a pessimistic write lock forever.", databaseShortName);
+            }
+        }
+        return em;
     }
 
     @Override
@@ -370,6 +388,11 @@ public class JpaMapStorageProviderFactory implements
         properties.put("hibernate.show_sql", config.getBoolean("showSql", false));
         properties.put("hibernate.format_sql", config.getBoolean("formatSql", true));
         properties.put("hibernate.dialect", config.get("driverDialect"));
+        Integer lockTimeout = config.getInt("lockTimeout");
+        if (lockTimeout != null) {
+            // This property does not work for PostgreSQL/CockroachDB - https://hibernate.atlassian.net/browse/HHH-16071
+            properties.put("javax.persistence.lock.timeout", lockTimeout);
+        }
 
         logger.trace("Creating EntityManagerFactory");
         ParsedPersistenceXmlDescriptor descriptor = PersistenceXmlParser.locateIndividualPersistenceUnit(
@@ -416,6 +439,7 @@ public class JpaMapStorageProviderFactory implements
 
                             MapJpaUpdaterProvider updater = session.getProvider(MapJpaUpdaterProvider.class);
                             MapJpaUpdaterProvider.Status status = updater.validate(modelType, connection, config.get("schema"));
+                            databaseShortName = updater.getDatabaseShortName();
 
                             if (!status.equals(VALID)) {
                                 update(modelType, connection, session);

model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/liquibase/updater/MapJpaLiquibaseUpdaterProvider.java

@@ -45,6 +45,7 @@ public class MapJpaLiquibaseUpdaterProvider implements MapJpaUpdaterProvider {
     private static final Logger logger = Logger.getLogger(MapJpaLiquibaseUpdaterProvider.class);
 
     private final KeycloakSession session;
+    private String databaseShortName;
 
     public MapJpaLiquibaseUpdaterProvider(KeycloakSession session) {
         this.session = session;
@@ -187,6 +188,7 @@ public class MapJpaLiquibaseUpdaterProvider implements MapJpaUpdaterProvider {
         try (Database database = DatabaseFactory.getInstance().findCorrectDatabaseImplementation(new JdbcConnectionFromPool(connection.unwrap(Connection.class)))) {
             // if the database is cockroachdb, use the aggregate changelog (see GHI #11230).
             String changelog = database instanceof CockroachDatabase ? "META-INF/jpa-aggregate-changelog.xml" : "META-INF/jpa-" + modelName + "-changelog.xml";
+            databaseShortName = database.getShortName();
             return liquibaseProvider.getLiquibaseForCustomUpdate(connection, defaultSchema, changelog, this.getClass().getClassLoader(), "databasechangelog");
         } catch (SQLException e) {
             throw new LiquibaseException(e);
@@ -197,4 +199,8 @@ public class MapJpaLiquibaseUpdaterProvider implements MapJpaUpdaterProvider {
     public void close() {
     }
 
+    @Override
+    public String getDatabaseShortName() {
+        return databaseShortName;
+    }
 }

model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/updater/MapJpaUpdaterProvider.java

@@ -68,4 +68,8 @@ public interface MapJpaUpdaterProvider extends Provider {
      */
     void export(Class<?> modelType, Connection connection, String defaultSchema, File file);
 
+    /**
+     * Returns an all-lower-case short name of the used database.
+     */
+    String getDatabaseShortName();
 }

testsuite/model/src/test/java/org/keycloak/testsuite/model/parameters/JpaMapStorage.java

@@ -95,7 +95,8 @@ public class JpaMapStorage extends KeycloakModelParameters {
                 .config("user", POSTGRES_DB_USER)
                 .config("password", POSTGRES_DB_PASSWORD)
                 .config("driver", "org.postgresql.Driver")
-                .config("driverDialect", "org.keycloak.models.map.storage.jpa.hibernate.dialect.JsonbPostgreSQL95Dialect");
+                .config("driverDialect", "org.keycloak.models.map.storage.jpa.hibernate.dialect.JsonbPostgreSQL95Dialect")
+                .config("lockTimeout", "1000");
 
         cf.spi(AuthenticationSessionSpi.PROVIDER_ID).provider(MapRootAuthenticationSessionProviderFactory.PROVIDER_ID)  .config(STORAGE_CONFIG, JpaMapStorageProviderFactory.PROVIDER_ID)
           .spi("client").provider(MapClientProviderFactory.PROVIDER_ID)                                                 .config(STORAGE_CONFIG, JpaMapStorageProviderFactory.PROVIDER_ID)

testsuite/model/src/test/java/org/keycloak/testsuite/model/parameters/JpaMapStorageCockroachdb.java

@@ -96,7 +96,8 @@ public class JpaMapStorageCockroachdb extends KeycloakModelParameters {
                 .config("user", COCKROACHDB_DB_USER)
                 .config("password", COCKROACHDB_DB_PASSWORD)
                 .config("driver", "org.postgresql.Driver")
-                .config("driverDialect", "org.keycloak.models.map.storage.jpa.hibernate.dialect.JsonbPostgreSQL95Dialect");
+                .config("driverDialect", "org.keycloak.models.map.storage.jpa.hibernate.dialect.JsonbPostgreSQL95Dialect")
+                .config("lockTimeout", "1000");
 
         cf.spi(AuthenticationSessionSpi.PROVIDER_ID).provider(MapRootAuthenticationSessionProviderFactory.PROVIDER_ID)  .config(STORAGE_CONFIG, JpaMapStorageProviderFactory.PROVIDER_ID)
           .spi("client").provider(MapClientProviderFactory.PROVIDER_ID)                                                 .config(STORAGE_CONFIG, JpaMapStorageProviderFactory.PROVIDER_ID)

testsuite/model/src/test/java/org/keycloak/testsuite/model/transaction/StorageTransactionTest.java

@@ -0,0 +1,185 @@
+/*
+ * Copyright 2023 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.testsuite.model.transaction;
+
+import org.junit.Test;
+import org.keycloak.models.Constants;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.ModelException;
+import org.keycloak.models.RealmModel;
+import org.keycloak.models.RealmProvider;
+import org.keycloak.models.map.storage.MapStorageProvider;
+import org.keycloak.models.map.storage.jpa.JpaMapStorageProviderFactory;
+import org.keycloak.testsuite.model.KeycloakModelTest;
+import org.keycloak.testsuite.model.RequireProvider;
+import org.keycloak.testsuite.model.util.TransactionController;
+import org.keycloak.utils.LockObjectsForModification;
+
+import javax.persistence.OptimisticLockException;
+import javax.persistence.PessimisticLockException;
+
+import static org.hamcrest.CoreMatchers.allOf;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.internal.matchers.ThrowableCauseMatcher.hasCause;
+import static org.keycloak.testsuite.model.util.KeycloakAssertions.assertException;
+
+@RequireProvider(RealmProvider.class)
+public class StorageTransactionTest extends KeycloakModelTest {
+
+    private String realmId;
+
+    @Override
+    protected void createEnvironment(KeycloakSession s) {
+        RealmModel r = s.realms().createRealm("1");
+        r.setDefaultRole(s.roles().addRealmRole(r, Constants.DEFAULT_ROLES_ROLE_PREFIX + "-" + r.getName()));
+        r.setAttribute("k1", "v1");
+        realmId = r.getId();
+    }
+
+    @Override
+    protected void cleanEnvironment(KeycloakSession s) {
+        s.realms().removeRealm(realmId);
+    }
+
+    @Override
+    protected boolean isUseSameKeycloakSessionFactoryForAllThreads() {
+        return true;
+    }
+
+    @Test
+    public void testTwoTransactionsSequentially() throws Exception {
+        TransactionController tx1 = new TransactionController(getFactory());
+        TransactionController tx2 = new TransactionController(getFactory());
+
+        tx1.begin();
+        assertThat(
+                tx1.runStep(session -> {
+                    session.realms().getRealm(realmId).setAttribute("k2", "v1");
+                    return session.realms().getRealm(realmId).getAttribute("k2");
+                }), equalTo("v1"));
+        tx1.commit();
+
+        tx2.begin();
+        assertThat(
+                tx2.runStep(session -> session.realms().getRealm(realmId).getAttribute("k2")),
+                equalTo("v1"));
+        tx2.commit();
+
+    }
+
+    @Test
+    public void testRepeatableRead() {
+        TransactionController tx1 = new TransactionController(getFactory());
+        TransactionController tx2 = new TransactionController(getFactory());
+        TransactionController tx3 = new TransactionController(getFactory());
+
+        tx1.begin();
+        tx2.begin();
+        tx3.begin();
+
+        // Read original value in tx1
+        assertThat(
+                tx1.runStep(session -> session.realms().getRealm(realmId).getAttribute("k1")),
+                equalTo("v1"));
+
+        // change value to new in tx2
+        tx2.runStep(session -> {
+            session.realms().getRealm(realmId).setAttribute("k1", "v2");
+            return null;
+        });
+        tx2.commit();
+
+        // tx1 should still return the value that already read
+        assertThat(
+                tx1.runStep(session -> session.realms().getRealm(realmId).getAttribute("k1")),
+                equalTo("v1"));
+
+        // tx3 should return the new value
+        assertThat(
+                tx3.runStep(session -> session.realms().getRealm(realmId).getAttribute("k1")),
+                equalTo("v2"));
+        tx1.commit();
+        tx3.commit();
+    }
+
+    @Test
+    // LockObjectForModification is currently used only in map-jpa
+    @RequireProvider(value = MapStorageProvider.class, only = JpaMapStorageProviderFactory.PROVIDER_ID)
+    public void testLockObjectForModification() {
+        TransactionController tx1 = new TransactionController(getFactory());
+        TransactionController tx2 = new TransactionController(getFactory());
+        TransactionController tx3 = new TransactionController(getFactory());
+
+        tx1.begin();
+        tx2.begin();
+
+        // tx1 acquires lock
+        tx1.runStep(session -> LockObjectsForModification.lockRealmsForModification(session, () -> session.realms().getRealm(realmId)));
+
+        // tx2 should fail as tx1 locked the realm
+        assertException(() -> tx2.runStep(session -> LockObjectsForModification.lockRealmsForModification(session, () -> session.realms().getRealm(realmId))),
+                allOf(instanceOf(ModelException.class),
+                        hasCause(instanceOf(PessimisticLockException.class))));
+
+        // end both transactions
+        tx2.rollback();
+        tx1.commit();
+
+        // start new transaction and read again, it should be successful
+        tx3.begin();
+        tx3.runStep(session -> LockObjectsForModification.lockRealmsForModification(session, () -> session.realms().getRealm(realmId)));
+        tx3.commit();
+    }
+
+    @Test
+    // Optimistic locking works only with map-jpa
+    @RequireProvider(value = MapStorageProvider.class, only = JpaMapStorageProviderFactory.PROVIDER_ID)
+    public void testOptimisticLockingException() {
+        withRealm(realmId, (session, realm) -> {
+            realm.setDisplayName("displayName1");
+            return null;
+        });
+
+        TransactionController tx1 = new TransactionController(getFactory());
+        TransactionController tx2 = new TransactionController(getFactory());
+
+        // tx1 acquires lock
+        tx1.begin();
+        tx2.begin();
+
+        // both transactions touch the same entity
+        tx1.runStep(session -> {
+            session.realms().getRealm(realmId).setDisplayName("displayName2");
+            return null;
+        });
+        tx2.runStep(session -> {
+            session.realms().getRealm(realmId).setDisplayName("displayName3");
+            return null;
+        });
+
+        // tx1 transaction should be successful
+        tx1.commit();
+
+        // tx2 should fail as tx1 already changed the value
+        assertException(tx2::commit,
+                allOf(instanceOf(ModelException.class),
+                        hasCause(instanceOf(OptimisticLockException.class))));
+    }
+}

testsuite/model/src/test/java/org/keycloak/testsuite/model/util/KeycloakAssertions.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2023 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.testsuite.model.util;
+
+import org.hamcrest.Matcher;
+
+import static org.hamcrest.CoreMatchers.allOf;
+import static org.hamcrest.CoreMatchers.notNullValue;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+public class KeycloakAssertions {
+
+    /**
+     * Runs {@code task} and checks whether the execution resulted in
+     * an exception that matches {@code matcher}. The method fails also
+     * when no exception is thrown.
+     *
+     * @param task task
+     * @param matcher matcher that the exception should match
+     */
+    public static void assertException(Runnable task, Matcher<? super Throwable> matcher) {
+        Throwable ex = catchException(task);
+        assertThat(ex, allOf(notNullValue(), matcher));
+    }
+
+    /**
+     * Runs the {@code task} and returns any throwable that is thrown.
+     * If not exception is thrown, the method returns {@code null}
+     *
+     * @param task task
+     */
+    public static Throwable catchException(Runnable task) {
+        try {
+            task.run();
+            return null;
+        } catch (Throwable ex) {
+            return ex;
+        }
+    }
+}

testsuite/model/src/test/java/org/keycloak/testsuite/model/util/TransactionController.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2023 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.testsuite.model.util;
+
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.KeycloakTransactionManager;
+
+import java.util.function.Function;
+
+public class TransactionController {
+    private final KeycloakSession session;
+
+    public TransactionController(KeycloakSessionFactory sessionFactory) {
+        session = sessionFactory.create();
+    }
+
+    public void begin() {
+        getTransactionManager().begin();
+    }
+
+    public void commit() {
+        getTransactionManager().commit();
+    }
+
+    public void rollback() {
+        getTransactionManager().rollback();
+    }
+
+    public <R> R runStep(Function<KeycloakSession, R> task) {
+        return task.apply(session);
+    }
+
+    private KeycloakTransactionManager getTransactionManager() {
+        return session.getTransactionManager();
+    }
+}