KEYCLOAK-3344 Can't remove realm with mongo model

model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java

@@ -37,7 +37,9 @@ import org.keycloak.connections.jpa.updater.JpaUpdaterProvider;
 import org.keycloak.connections.jpa.util.JpaUtils;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.KeycloakSessionTask;
 import org.keycloak.models.dblock.DBLockProvider;
+import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.provider.ServerInfoAwareProviderFactory;
 import org.keycloak.models.dblock.DBLockManager;
 import org.keycloak.timer.TimerProvider;
@@ -93,8 +95,6 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
                 if (emf == null) {
                     logger.debug("Initializing JPA connections");
 
-                    Connection connection = null;
-
                     Map<String, Object> properties = new HashMap<String, Object>();
 
                     String unitName = "keycloak-default";
@@ -126,23 +126,26 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
                     }
 
 
-                    String databaseSchema = config.get("databaseSchema");
+                    String databaseSchema;
-                    if (databaseSchema == null) {
+                    String databaseSchemaConf = config.get("databaseSchema");
+                    if (databaseSchemaConf == null) {
                         throw new RuntimeException("Property 'databaseSchema' needs to be specified in the configuration");
                     }
                     
-                    if (databaseSchema.equals("development-update")) {
+                    if (databaseSchemaConf.equals("development-update")) {
                         properties.put("hibernate.hbm2ddl.auto", "update");
                         databaseSchema = null;
-                    } else if (databaseSchema.equals("development-validate")) {
+                    } else if (databaseSchemaConf.equals("development-validate")) {
                         properties.put("hibernate.hbm2ddl.auto", "validate");
                         databaseSchema = null;
+                    } else {
+                        databaseSchema = databaseSchemaConf;
                     }
 
                     properties.put("hibernate.show_sql", config.getBoolean("showSql", false));
                     properties.put("hibernate.format_sql", config.getBoolean("formatSql", true));
 
-                    connection = getConnection();
+                    Connection connection = getConnection();
                     try{ 
 	                    prepareOperationalInfo(connection);
 
@@ -161,19 +164,27 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
 
                             // Check if having DBLock before trying to initialize hibernate
                             DBLockProvider dbLock = new DBLockManager(session).getDBLock();
-                            if (!dbLock.hasLock()) {
+                            if (dbLock.hasLock()) {
-                                throw new IllegalStateException("Trying to update database, but don't have a DB lock acquired");
+                                updateOrValidateDB(databaseSchema, connection, updater, schema);
-                            }
-	
-	                        if (databaseSchema.equals("update")) {
-                                updater.update(connection, schema);
-	                        } else if (databaseSchema.equals("validate")) {
-	                            updater.validate(connection, schema);
                             } else {
-	                            throw new RuntimeException("Invalid value for databaseSchema: " + databaseSchema);
+                                logger.trace("Don't have DBLock retrieved before upgrade. Needs to acquire lock first in separate transaction");
+
+                                KeycloakModelUtils.runJobInTransaction(session.getKeycloakSessionFactory(), new KeycloakSessionTask() {
+
+                                    @Override
+                                    public void run(KeycloakSession lockSession) {
+                                        DBLockManager dbLockManager = new DBLockManager(lockSession);
+                                        DBLockProvider dbLock2 = dbLockManager.getDBLock();
+                                        dbLock2.waitForLock();
+                                        try {
+                                            updateOrValidateDB(databaseSchema, connection, updater, schema);
+                                        } finally {
+                                            dbLock2.releaseLock();
+                                        }
                                     }
 
-	                        logger.trace("Database update completed");
+                                });
+                            }
 	                    }
 
                         int globalStatsInterval = config.getInt("globalStatsInterval", -1);
@@ -272,6 +283,20 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
     }
 
 
+    // Needs to be called with acquired DBLock
+    protected void updateOrValidateDB(String databaseSchema, Connection connection, JpaUpdaterProvider updater, String schema) {
+        if (databaseSchema.equals("update")) {
+            updater.update(connection, schema);
+            logger.trace("Database update completed");
+        } else if (databaseSchema.equals("validate")) {
+            updater.validate(connection, schema);
+            logger.trace("Database validation completed");
+        } else {
+            throw new RuntimeException("Invalid value for databaseSchema: " + databaseSchema);
+        }
+    }
+
+
     @Override
     public Connection getConnection() {
         try {

model/mongo/src/main/java/org/keycloak/connections/mongo/DefaultMongoConnectionFactoryProvider.java

@@ -34,8 +34,10 @@ import org.keycloak.connections.mongo.impl.context.TransactionMongoStoreInvocati
 import org.keycloak.connections.mongo.updater.MongoUpdaterProvider;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.KeycloakSessionTask;
 import org.keycloak.models.dblock.DBLockManager;
 import org.keycloak.models.dblock.DBLockProvider;
+import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.provider.ServerInfoAwareProviderFactory;
 
 import com.mongodb.DB;
@@ -174,8 +176,26 @@ public class DefaultMongoConnectionFactoryProvider implements MongoConnectionPro
             }
 
             DBLockProvider dbLock = new DBLockManager(session).getDBLock();
-            if (!dbLock.hasLock()) {
+            if (dbLock.hasLock()) {
-                throw new IllegalStateException("Trying to update database, but don't have a DB lock acquired");
+                updateOrValidateDB(databaseSchema, session, mongoUpdater);
+            } else {
+                logger.trace("Don't have DBLock retrieved before upgrade. Needs to acquire lock first in separate transaction");
+
+                KeycloakModelUtils.runJobInTransaction(session.getKeycloakSessionFactory(), new KeycloakSessionTask() {
+
+                    @Override
+                    public void run(KeycloakSession lockSession) {
+                        DBLockManager dbLockManager = new DBLockManager(lockSession);
+                        DBLockProvider dbLock2 = dbLockManager.getDBLock();
+                        dbLock2.waitForLock();
+                        try {
+                            updateOrValidateDB(databaseSchema, session, mongoUpdater);
+                        } finally {
+                            dbLock2.releaseLock();
+                        }
+                    }
+
+                });
             }
 
             if (databaseSchema.equals("update")) {
@@ -197,6 +217,16 @@ public class DefaultMongoConnectionFactoryProvider implements MongoConnectionPro
         return entityClasses;
     }
 
+    protected void updateOrValidateDB(String databaseSchema, KeycloakSession session, MongoUpdaterProvider mongoUpdater) {
+        if (databaseSchema.equals("update")) {
+            mongoUpdater.update(session, db);
+        } else if (databaseSchema.equals("validate")) {
+            mongoUpdater.validate(session, db);
+        } else {
+            throw new RuntimeException("Invalid value for databaseSchema: " + databaseSchema);
+        }
+    }
+
     @Override
     public void close() {
         if (client != null) {