KEYCLOAK-18824 Simplify MapStorageTransaction and move registerEntityForChanges to CHM transaction

model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionEntity.java

@@ -18,6 +18,7 @@ package org.keycloak.models.map.authSession;
 
 import org.keycloak.models.map.common.AbstractEntity;
 
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
@@ -25,7 +26,7 @@ import java.util.concurrent.ConcurrentHashMap;
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
  */
-public class MapRootAuthenticationSessionEntity<K> implements AbstractEntity<K> {
+public class MapRootAuthenticationSessionEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private K id;
     private String realmId;

model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionProvider.java

@@ -39,7 +39,6 @@ import java.util.function.Function;
 import java.util.function.Predicate;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
 /**
@@ -65,7 +64,7 @@ public class MapRootAuthenticationSessionProvider<K> implements AuthenticationSe
     private Function<MapRootAuthenticationSessionEntity<K>, RootAuthenticationSessionModel> entityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
 
-        return origEntity -> new MapRootAuthenticationSessionAdapter<K>(session, realm, registerEntityForChanges(tx, origEntity)) {
+        return origEntity -> new MapRootAuthenticationSessionAdapter<K>(session, realm, origEntity) {
             @Override
             public String getId() {
                 return sessionStore.getKeyConvertor().keyToString(entity.getId());
@@ -145,7 +144,7 @@ public class MapRootAuthenticationSessionProvider<K> implements AuthenticationSe
           .compare(SearchableFields.REALM_ID, Operator.EQ, realm.getId())
           .compare(SearchableFields.TIMESTAMP, Operator.LT, expired);
 
-        long deletedCount = tx.delete(sessionStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        long deletedCount = tx.delete(withCriteria(mcb));
 
         LOG.debugf("Removed %d expired authentication sessions for realm '%s'", deletedCount, realm.getName());
     }

model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionProviderFactory.java

@@ -30,7 +30,7 @@ public class MapRootAuthenticationSessionProviderFactory<K> extends AbstractMapP
         implements AuthenticationSessionProviderFactory {
 
     public MapRootAuthenticationSessionProviderFactory() {
-        super(MapRootAuthenticationSessionEntity.class, RootAuthenticationSessionModel.class);
+        super(RootAuthenticationSessionModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/authorization/MapAuthorizationStoreFactory.java

@@ -65,11 +65,11 @@ public class MapAuthorizationStoreFactory<K> implements AmphibianProviderFactory
         MapStorage resourceStore;
         MapStorage scopeStore;
 
-        permissionTicketStore = mapStorageProvider.getStorage(MapPermissionTicketEntity.class, PermissionTicket.class);
+        permissionTicketStore = mapStorageProvider.getStorage(PermissionTicket.class);
-        policyStore = mapStorageProvider.getStorage(MapPolicyEntity.class, Policy.class);
+        policyStore = mapStorageProvider.getStorage(Policy.class);
-        resourceServerStore = mapStorageProvider.getStorage(MapResourceServerEntity.class, ResourceServer.class);
+        resourceServerStore = mapStorageProvider.getStorage(ResourceServer.class);
-        resourceStore = mapStorageProvider.getStorage(MapResourceEntity.class, Resource.class);
+        resourceStore = mapStorageProvider.getStorage(Resource.class);
-        scopeStore = mapStorageProvider.getStorage(MapScopeEntity.class, Scope.class);
+        scopeStore = mapStorageProvider.getStorage(Scope.class);
         
         return new MapAuthorizationStore(session,
                     permissionTicketStore,

model/map/src/main/java/org/keycloak/models/map/authorization/MapPermissionTicketStore.java

@@ -43,7 +43,6 @@ import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 import static org.keycloak.utils.StreamsUtil.distinctByKey;
@@ -66,7 +65,7 @@ public class MapPermissionTicketStore<K extends Comparable<K>> implements Permis
     private PermissionTicket entityToAdapter(MapPermissionTicketEntity<K> origEntity) {
         if (origEntity == null) return null;
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return new MapPermissionTicketAdapter<K>(registerEntityForChanges(tx, origEntity), authorizationProvider.getStoreFactory()) {
+        return new MapPermissionTicketAdapter<K>(origEntity, authorizationProvider.getStoreFactory()) {
             @Override
             public String getId() {
                 return permissionTicketStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/authorization/MapPolicyStore.java

@@ -41,7 +41,6 @@ import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
 public class MapPolicyStore<K> implements PolicyStore {
@@ -61,7 +60,7 @@ public class MapPolicyStore<K> implements PolicyStore {
     private Policy entityToAdapter(MapPolicyEntity<K> origEntity) {
         if (origEntity == null) return null;
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return new MapPolicyAdapter<K>(registerEntityForChanges(tx, origEntity), authorizationProvider.getStoreFactory()) {
+        return new MapPolicyAdapter<K>(origEntity, authorizationProvider.getStoreFactory()) {
             @Override
             public String getId() {
                 return policyStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/authorization/MapResourceServerStore.java

@@ -39,7 +39,6 @@ import org.keycloak.models.map.storage.MapStorage;
 import org.keycloak.storage.StorageId;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 
 public class MapResourceServerStore<K> implements ResourceServerStore {
 
@@ -58,7 +57,7 @@ public class MapResourceServerStore<K> implements ResourceServerStore {
     private ResourceServer entityToAdapter(MapResourceServerEntity<K> origEntity) {
         if (origEntity == null) return null;
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return new MapResourceServerAdapter<K>(registerEntityForChanges(tx, origEntity), authorizationProvider.getStoreFactory()) {
+        return new MapResourceServerAdapter<K>(origEntity, authorizationProvider.getStoreFactory()) {
             @Override
             public String getId() {
                 return resourceServerStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/authorization/MapResourceStore.java

@@ -40,7 +40,6 @@ import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
 public class MapResourceStore<K extends Comparable<K>> implements ResourceStore {
@@ -60,7 +59,7 @@ public class MapResourceStore<K extends Comparable<K>> implements ResourceStore
     private Resource entityToAdapter(MapResourceEntity<K> origEntity) {
         if (origEntity == null) return null;
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return new MapResourceAdapter<K>(registerEntityForChanges(tx, origEntity), authorizationProvider.getStoreFactory()) {
+        return new MapResourceAdapter<K>(origEntity, authorizationProvider.getStoreFactory()) {
             @Override
             public String getId() {
                 return resourceStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/authorization/MapScopeStore.java

@@ -38,7 +38,6 @@ import java.util.Map;
 import java.util.stream.Collectors;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
 public class MapScopeStore<K> implements ScopeStore {
@@ -58,7 +57,7 @@ public class MapScopeStore<K> implements ScopeStore {
     private Scope entityToAdapter(MapScopeEntity<K> origEntity) {
         if (origEntity == null) return null;
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return new MapScopeAdapter<K>(registerEntityForChanges(tx, origEntity), authorizationProvider.getStoreFactory()) {
+        return new MapScopeAdapter<K>(origEntity, authorizationProvider.getStoreFactory()) {
             @Override
             public String getId() {
                 return scopeStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/authorization/entity/MapPermissionTicketEntity.java

@@ -19,10 +19,10 @@ package org.keycloak.models.map.authorization.entity;
 
 import org.keycloak.models.map.common.AbstractEntity;
 
-import java.util.Comparator;
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Objects;
 
-public class MapPermissionTicketEntity<K> implements AbstractEntity<K> {
+public class MapPermissionTicketEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private final K id;
     private String owner;

model/map/src/main/java/org/keycloak/models/map/authorization/entity/MapPolicyEntity.java

@@ -18,17 +18,17 @@
 package org.keycloak.models.map.authorization.entity;
 
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.representations.idm.authorization.DecisionStrategy;
 import org.keycloak.representations.idm.authorization.Logic;
 
-import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.Map;
 import java.util.Objects;
 
-public class MapPolicyEntity<K> implements AbstractEntity<K> {
+public class MapPolicyEntity<K> implements AbstractEntity<K>, UpdatableEntity {
     
     private final K id;
     private String name;

model/map/src/main/java/org/keycloak/models/map/authorization/entity/MapResourceEntity.java

@@ -19,7 +19,7 @@ package org.keycloak.models.map.authorization.entity;
 
 import org.keycloak.models.map.common.AbstractEntity;
 
-import java.util.Comparator;
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -27,7 +27,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 
-public class MapResourceEntity<K> implements AbstractEntity<K> {
+public class MapResourceEntity<K> implements AbstractEntity<K>, UpdatableEntity {
     
     private final K id;
     private String name;

model/map/src/main/java/org/keycloak/models/map/authorization/entity/MapResourceServerEntity.java

@@ -18,12 +18,13 @@
 package org.keycloak.models.map.authorization.entity;
 
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.representations.idm.authorization.DecisionStrategy;
 import org.keycloak.representations.idm.authorization.PolicyEnforcementMode;
 
 import java.util.Objects;
 
-public class MapResourceServerEntity<K> implements AbstractEntity<K> {
+public class MapResourceServerEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private final K id;
     private boolean updated = false;

model/map/src/main/java/org/keycloak/models/map/authorization/entity/MapScopeEntity.java

@@ -19,9 +19,10 @@ package org.keycloak.models.map.authorization.entity;
 
 import org.keycloak.models.map.common.AbstractEntity;
 
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Objects;
 
-public class MapScopeEntity<K> implements AbstractEntity<K> {
+public class MapScopeEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private final K id;
     private String name;

model/map/src/main/java/org/keycloak/models/map/client/MapClientEntity.java

@@ -18,6 +18,7 @@ package org.keycloak.models.map.client;
 
 import org.keycloak.models.ProtocolMapperModel;
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -28,7 +29,7 @@ import java.util.stream.Stream;
  *
  * @author hmlnarik
  */
-public interface MapClientEntity<K> extends AbstractEntity<K> {
+public interface MapClientEntity<K> extends AbstractEntity<K>, UpdatableEntity {
 
     void addClientScope(String id, Boolean defaultScope);
 

model/map/src/main/java/org/keycloak/models/map/client/MapClientProvider.java

@@ -44,7 +44,6 @@ import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import org.keycloak.models.map.storage.ModelCriteriaBuilder.Operator;
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
 import org.keycloak.models.ClientScopeModel;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
@@ -84,7 +83,7 @@ public class MapClientProvider<K> implements ClientProvider {
     private <T extends MapClientEntity<K>> Function<T, ClientModel> entityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
 
-        return (T origEntity) -> new MapClientAdapter<K>(session, realm, registerEntityForChanges(tx, origEntity)) {
+        return origEntity -> new MapClientAdapter<K>(session, realm, origEntity) {
             @Override
             public String getId() {
                 return clientStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/client/MapClientProviderFactory.java

@@ -42,7 +42,7 @@ public class MapClientProviderFactory<K> extends AbstractMapProviderFactory<Clie
     private Runnable onClose;
 
     public MapClientProviderFactory() {
-        super(MapClientEntity.class, ClientModel.class);
+        super(ClientModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/clientscope/MapClientScopeEntity.java

@@ -29,8 +29,9 @@ import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.keycloak.models.ProtocolMapperModel;
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 
-public class MapClientScopeEntity<K> implements AbstractEntity<K> {
+public class MapClientScopeEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private final K id;
     private final String realmId;

model/map/src/main/java/org/keycloak/models/map/clientscope/MapClientScopeProvider.java

@@ -36,7 +36,6 @@ import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import org.keycloak.models.map.storage.ModelCriteriaBuilder.Operator;
 import org.keycloak.models.utils.KeycloakModelUtils;
 
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
@@ -57,7 +56,7 @@ public class MapClientScopeProvider<K> implements ClientScopeProvider {
     private Function<MapClientScopeEntity<K>, ClientScopeModel> entityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
 
-        return origEntity -> new MapClientScopeAdapter<K>(session, realm, registerEntityForChanges(tx, origEntity)) {
+        return origEntity -> new MapClientScopeAdapter<K>(session, realm, origEntity) {
             @Override
             public String getId() {
                 return clientScopeStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/clientscope/MapClientScopeProviderFactory.java

@@ -25,7 +25,7 @@ import org.keycloak.models.map.common.AbstractMapProviderFactory;
 public class MapClientScopeProviderFactory<K> extends AbstractMapProviderFactory<ClientScopeProvider, K, MapClientScopeEntity<K>, ClientScopeModel> implements ClientScopeProviderFactory {
 
     public MapClientScopeProviderFactory() {
-        super(MapClientScopeEntity.class, ClientScopeModel.class);
+        super(ClientScopeModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/common/AbstractEntity.java

@@ -20,7 +20,7 @@ package org.keycloak.models.map.common;
  *
  * @author hmlnarik
  */
-public interface AbstractEntity<K> extends UpdatableEntity {
+public interface AbstractEntity<K> {
 
     K getId();
 

model/map/src/main/java/org/keycloak/models/map/common/AbstractMapProviderFactory.java

@@ -44,14 +44,11 @@ public abstract class AbstractMapProviderFactory<T extends Provider, K, V extend
 
     protected final Class<M> modelType;
 
-    protected final Class<V> entityType;
-
     private Scope storageConfigScope;
 
     @SuppressWarnings("unchecked")
-    protected AbstractMapProviderFactory(Class<? extends AbstractEntity> entityType, Class<M> modelType) {
+    protected AbstractMapProviderFactory(Class<M> modelType) {
         this.modelType = modelType;
-        this.entityType = (Class<V>) entityType;
     }
 
     @Override
@@ -64,7 +61,7 @@ public abstract class AbstractMapProviderFactory<T extends Provider, K, V extend
           MapStorageProvider.class, storageConfigScope, MapStorageSpi.NAME);
         final MapStorageProvider factory = storageProviderFactory.create(session);
 
-        return factory.getStorage(entityType, modelType);
+        return factory.getStorage(modelType);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/common/MapStorageUtils.java

@@ -1,42 +0,0 @@
-/*
- * Copyright 2021 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.models.map.common;
-
-import org.keycloak.models.map.storage.MapKeycloakTransaction;
-
-/**
- *
- * @author hmlnarik
- */
-public class MapStorageUtils {
-
-    /**
-     * Returns a deep clone of an entity. If the clone is already in the transaction, returns this one.
-     * <p>
-     * Usually used before giving an entity from a source back to the caller,
-     * to prevent changing it directly in the data store, but to keep transactional properties.
-     * @param <K>
-     * @param <V>
-     * @param tx Transaction that is checked for existence of the entity before
-     * @param origEntity
-     * @return
-     */
-    public static <K, V extends AbstractEntity<K>> V registerEntityForChanges(MapKeycloakTransaction<K, V, ?> tx, V origEntity) {
-        final V res = tx.read(origEntity.getId(), id -> Serialization.from(origEntity));
-        return tx.updateIfChanged(res, AbstractEntity<K>::isUpdated);
-    }
-}

model/map/src/main/java/org/keycloak/models/map/group/MapGroupEntity.java

@@ -19,6 +19,7 @@ package org.keycloak.models.map.group;
 
 import org.keycloak.models.map.common.AbstractEntity;
 
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -30,7 +31,7 @@ import java.util.Set;
  *
  * @author mhajas
  */
-public class MapGroupEntity<K> implements AbstractEntity<K> {
+public class MapGroupEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private final K id;
     private final String realmId;

model/map/src/main/java/org/keycloak/models/map/group/MapGroupProvider.java

@@ -39,7 +39,6 @@ import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
@@ -59,7 +58,7 @@ public class MapGroupProvider<K> implements GroupProvider {
 
     private Function<MapGroupEntity<K>, GroupModel> entityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return origEntity -> new MapGroupAdapter<K>(session, realm, registerEntityForChanges(tx, origEntity)) {
+        return origEntity -> new MapGroupAdapter<K>(session, realm, origEntity) {
             @Override
             public String getId() {
                 return groupStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/group/MapGroupProviderFactory.java

@@ -40,7 +40,7 @@ public class MapGroupProviderFactory<K> extends AbstractMapProviderFactory<Group
     private Runnable onClose;
 
     public MapGroupProviderFactory() {
-        super(MapGroupEntity.class, GroupModel.class);
+        super(GroupModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/loginFailure/MapUserLoginFailureEntity.java

@@ -18,12 +18,13 @@ package org.keycloak.models.map.loginFailure;
 
 import org.keycloak.models.map.common.AbstractEntity;
 
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Objects;
 
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
  */
-public class MapUserLoginFailureEntity<K> implements AbstractEntity<K> {
+public class MapUserLoginFailureEntity<K> implements AbstractEntity<K>, UpdatableEntity {
     private K id;
     private String realmId;
     private String userId;

model/map/src/main/java/org/keycloak/models/map/loginFailure/MapUserLoginFailureProvider.java

@@ -28,7 +28,6 @@ import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import java.util.function.Function;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
 /**
@@ -51,7 +50,7 @@ public class MapUserLoginFailureProvider<K> implements UserLoginFailureProvider
 
     private Function<MapUserLoginFailureEntity<K>, UserLoginFailureModel> userLoginFailureEntityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return origEntity -> new MapUserLoginFailureAdapter<K>(session, realm, registerEntityForChanges(userLoginFailureTx, origEntity)) {
+        return origEntity -> new MapUserLoginFailureAdapter<K>(session, realm, origEntity) {
             @Override
             public String getId() {
                 return userLoginFailureStore.getKeyConvertor().keyToString(entity.getId());
@@ -100,7 +99,7 @@ public class MapUserLoginFailureProvider<K> implements UserLoginFailureProvider
 
         LOG.tracef("removeUserLoginFailure(%s, %s)%s", realm, userId, getShortStackTrace());
 
-        userLoginFailureTx.delete(userLoginFailureStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        userLoginFailureTx.delete(withCriteria(mcb));
     }
 
     @Override
@@ -110,7 +109,7 @@ public class MapUserLoginFailureProvider<K> implements UserLoginFailureProvider
 
         LOG.tracef("removeAllUserLoginFailures(%s)%s", realm, getShortStackTrace());
 
-        userLoginFailureTx.delete(userLoginFailureStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        userLoginFailureTx.delete(withCriteria(mcb));
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/loginFailure/MapUserLoginFailureProviderFactory.java

@@ -37,7 +37,7 @@ public class MapUserLoginFailureProviderFactory<K> extends AbstractMapProviderFa
     private Runnable onClose;
 
     public MapUserLoginFailureProviderFactory() {
-        super(MapUserLoginFailureEntity.class, UserLoginFailureModel.class);
+        super(UserLoginFailureModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/realm/MapRealmEntity.java

@@ -30,6 +30,7 @@ import org.keycloak.common.util.Time;
 import org.keycloak.models.ModelDuplicateException;
 import org.keycloak.models.OTPPolicy;
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.models.map.realm.entity.MapAuthenticationExecutionEntity;
 import org.keycloak.models.map.realm.entity.MapAuthenticationFlowEntity;
 import org.keycloak.models.map.realm.entity.MapAuthenticatorConfigEntity;
@@ -42,7 +43,7 @@ import org.keycloak.models.map.realm.entity.MapRequiredActionProviderEntity;
 import org.keycloak.models.map.realm.entity.MapRequiredCredentialEntity;
 import org.keycloak.models.map.realm.entity.MapWebAuthnPolicyEntity;
 
-public class MapRealmEntity<K> implements AbstractEntity<K> {
+public class MapRealmEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private final K id;
     private String name;

model/map/src/main/java/org/keycloak/models/map/realm/MapRealmProvider.java

@@ -37,7 +37,6 @@ import org.keycloak.models.map.storage.MapStorage;
 import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import org.keycloak.models.map.storage.ModelCriteriaBuilder.Operator;
 import org.keycloak.models.utils.KeycloakModelUtils;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
@@ -56,7 +55,7 @@ public class MapRealmProvider<K> implements RealmProvider {
     }
 
     private RealmModel entityToAdapter(MapRealmEntity<K> entity) {
-        return new MapRealmAdapter<K>(session, registerEntityForChanges(tx, entity)) {
+        return new MapRealmAdapter<K>(session, entity) {
             @Override
             public String getId() {
                 return realmStore.getKeyConvertor().keyToString(entity.getId());
@@ -176,7 +175,6 @@ public class MapRealmProvider<K> implements RealmProvider {
                 .compare(SearchableFields.CLIENT_INITIAL_ACCESS, Operator.EXISTS);
 
         tx.read(withCriteria(mcb))
-                .map(e -> registerEntityForChanges(tx, e))
                 .forEach(MapRealmEntity<K>::removeExpiredClientInitialAccesses);
     }
 

model/map/src/main/java/org/keycloak/models/map/realm/MapRealmProviderFactory.java

@@ -25,7 +25,7 @@ import org.keycloak.models.RealmProviderFactory;
 public class MapRealmProviderFactory<K> extends AbstractMapProviderFactory<RealmProvider, K, MapRealmEntity<K>, RealmModel> implements RealmProviderFactory {
 
     public MapRealmProviderFactory() {
-        super(MapRealmEntity.class, RealmModel.class);
+        super(RealmModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/role/MapRoleEntity.java

@@ -23,8 +23,9 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 
-public class MapRoleEntity<K> implements AbstractEntity<K> {
+public class MapRoleEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private K id;
     private String realmId;

model/map/src/main/java/org/keycloak/models/map/role/MapRoleProvider.java

@@ -31,7 +31,6 @@ import java.util.function.Function;
 import java.util.stream.Stream;
 import org.keycloak.models.map.storage.MapStorage;
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
@@ -57,7 +56,7 @@ public class MapRoleProvider<K> implements RoleProvider {
 
     private Function<MapRoleEntity<K>, RoleModel> entityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return origEntity -> new MapRoleAdapter<K>(session, realm, registerEntityForChanges(tx, origEntity)) {
+        return origEntity -> new MapRoleAdapter<K>(session, realm, origEntity) {
             @Override
             public String getId() {
                 return roleStore.getKeyConvertor().keyToString(entity.getId());

model/map/src/main/java/org/keycloak/models/map/role/MapRoleProviderFactory.java

@@ -25,7 +25,7 @@ import org.keycloak.models.RoleProviderFactory;
 public class MapRoleProviderFactory<K> extends AbstractMapProviderFactory<RoleProvider, K, MapRoleEntity<K>, RoleModel> implements RoleProviderFactory {
 
     public MapRoleProviderFactory() {
-        super(MapRoleEntity.class, RoleModel.class);
+        super(RoleModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/storage/MapKeycloakTransaction.java

@@ -27,6 +27,9 @@ public interface MapKeycloakTransaction<K, V extends AbstractEntity<K>, M> exten
 
     /**
      * Instructs this transaction to add a new value into the underlying store on commit.
+     * <p>
+     * Updates to the returned instances of {@code V} would be visible in the current transaction
+     * and will propagate into the underlying store upon commit.
      *
      * @param value the value
      * @return Entity representing the {@code value} in the store. It may or may not be the same instance as {@code value}
@@ -35,29 +38,22 @@ public interface MapKeycloakTransaction<K, V extends AbstractEntity<K>, M> exten
 
     /**
      * Provides possibility to lookup for values by a {@code key} in the underlying store with respect to changes done
-     * in current transaction.
+     * in current transaction. Updates to the returned instance would be visible in the current transaction
+     * and will propagate into the underlying store upon commit.
      *
      * @param key identifier of a value
      * @return a value associated with the given {@code key}
      */
     V read(K key);
 
-    /**
-     * Looks up a value in the current transaction with corresponding key, returns {@code defaultValueFunc} when
-     * the transaction does not contain a value for the {@code key} identifier.
-     *
-     * @param key identifier of a value
-     * @param defaultValueFunc fallback function if the transaction does not contain a value that corresponds to {@code key}
-     * @return a value associated with the given {@code key}, or the result of {@code defaultValueFunc}
-     *
-     */
-    V read(K key, Function<K, V> defaultValueFunc);
-
     /**
      * Returns a stream of values from underlying storage that are updated based on the current transaction changes;
      * i.e. the result contains updates and excludes of records that have been created, updated or deleted in this
      * transaction by methods {@link MapKeycloakTransaction#create}, {@link MapKeycloakTransaction#update},
      * {@link MapKeycloakTransaction#delete}, etc.
+     * <p>
+     * Updates to the returned instances of {@code V} would be visible in the current transaction
+     * and will propagate into the underlying store upon commit.
      *
      * @param queryParameters parameters for the query like firstResult, maxResult, requested ordering, etc.
      * @return values that fulfill the given criteria, that are updated based on changes in the current transaction
@@ -73,38 +69,6 @@ public interface MapKeycloakTransaction<K, V extends AbstractEntity<K>, M> exten
      */
     long getCount(QueryParameters<M> queryParameters);
 
-    /**
-     * Instructs this transaction to force-update the {@code value} associated with the identifier {@code value.getId()} in the
-     * underlying store on commit.
-     *
-     * @param value updated version of the value
-     * @return Entity representing the {@code value} in the store. It may or may not be the same instance as {@code value}
-     */
-    V update(V value);
-
-    /**
-     * Returns an updated version of the {@code orig} object as updated in this transaction.
-     *
-     * If the underlying store handles transactions on its own, this can return {@code orig} directly.
-     *
-     * @param orig possibly stale version of some object from the underlying store
-     * @return the {@code orig} object as visible from this transaction, or {@code null} if the object has been removed.
-     */
-    default V getUpdated(V orig) {
-        return orig;
-    }
-
-    /**
-     * Instructs this transaction to update the {@code value} associated with the identifier {@code value.getId()} in the
-     * underlying store on commit, if by the time of {@code commit} the {@code shouldPut} predicate returns {@code true}
-     *
-     * @param value new version of the value. Must not alter the {@code id} of the entity
-     * @param shouldPut predicate to check in commit phase
-     * @return Entity representing the {@code value} in the store. It may or may not be the same instance as {@code value}
-     * @see AbstractEntity#getId()
-     */
-    V updateIfChanged(V value, Predicate<V> shouldPut);
-
     /**
      * Instructs this transaction to delete a value associated with the identifier {@code key} from the underlying store
      * on commit.
@@ -122,6 +86,6 @@ public interface MapKeycloakTransaction<K, V extends AbstractEntity<K>, M> exten
      * @param queryParameters parameters for the query like firstResult, maxResult, requested ordering, etc.
      * @return number of removed objects (might return {@code -1} if not supported)
      */
-    long delete(K artificialKey, QueryParameters<M> queryParameters);
+    long delete(QueryParameters<M> queryParameters);
 
 }

model/map/src/main/java/org/keycloak/models/map/storage/MapStorage.java

@@ -20,8 +20,6 @@ import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.map.common.AbstractEntity;
 import java.util.stream.Stream;
 
-import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
-
 /**
  * Implementation of this interface interacts with a persistence storage storing various entities, e.g. users, realms.
  * It contains basic object CRUD operations as well as bulk {@link #read(org.keycloak.models.map.storage.QueryParameters)}
@@ -41,10 +39,10 @@ import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 public interface MapStorage<K, V extends AbstractEntity<K>, M> {
 
     /**
-     * Creates an object in the store identified. The ID of the {@code value} should be non-{@code null}.
+     * Creates an object in the store. ID of the {@code value} may be prescribed in id of the {@code value}.
-     * If the ID is {@code null}, then the {@code value}'s ID will be returned
+     * If the id is {@code null}, then the {@code value}'s ID will be generated and returned in the id of the return value.
      * @param value Entity to create in the store
-     * @throws NullPointerException if object or its {@code key} is {@code null}
+     * @throws NullPointerException if {@code value} is {@code null}
      * @see AbstractEntity#getId()
      * @return Entity representing the {@code value} in the store. It may or may not be the same instance as {@code value}
      */

model/map/src/main/java/org/keycloak/models/map/storage/MapStorageProvider.java

@@ -27,13 +27,14 @@ import org.keycloak.provider.Provider;
 public interface MapStorageProvider extends Provider {
     
     /**
-     * Returns a key-value storage implementation for the particular types.
+     * Returns a key-value storage implementation for the given types.
      * @param <K> type of the primary key
      * @param <V> type of the value
-     * @param name Name of the storage
+     * @param <M> type of the corresponding model (e.g. {@code UserModel})
-     * @param flags
+     * @param modelType Model type
+     * @param flags Flags of the returned storage. Best effort, flags may be not honored by underlying implementation
      * @return
      * @throws IllegalArgumentException If some of the types is not supported by the underlying implementation.
      */
-    <K, V extends AbstractEntity<K>, M> MapStorage<K, V, M> getStorage(Class<V> valueType, Class<M> modelType, Flag... flags);
+    <K, V extends AbstractEntity<K>, M> MapStorage<K, V, M> getStorage(Class<M> modelType, Flag... flags);
 }

model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapKeycloakTransaction.java

@@ -17,6 +17,8 @@
 package org.keycloak.models.map.storage.chm;
 
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.Serialization;
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.Map;
@@ -33,7 +35,7 @@ import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import org.keycloak.models.map.storage.QueryParameters;
 import org.keycloak.utils.StreamsUtil;
 
-public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K>, M> implements MapKeycloakTransaction<K, V, M> {
+public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K> & UpdatableEntity, M> implements MapKeycloakTransaction<K, V, M> {
 
     private final static Logger log = Logger.getLogger(ConcurrentHashMapKeycloakTransaction.class);
 
@@ -98,16 +100,37 @@ public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K>
         tasks.merge(taskKey, task, MapTaskCompose::new);
     }
 
+    /**
+     * Returns a deep clone of an entity. If the clone is already in the transaction, returns this one.
+     * <p>
+     * Usually used before giving an entity from a source back to the caller,
+     * to prevent changing it directly in the data store, but to keep transactional properties.
+     * @param origEntity Original entity
+     * @return
+     */
+    public V registerEntityForChanges(V origEntity) {
+        final K key = origEntity.getId();
+        // If the entity is listed in the transaction already, return it directly
+        if (tasks.containsKey(key)) {
+            MapTaskWithValue current = tasks.get(key);
+            return current.getValue();
+        }
+        // Else enlist its copy in the transaction. Never return direct reference to the underlying map
+        final V res = Serialization.from(origEntity);
+        return updateIfChanged(res, e -> e.isUpdated());
+    }
+
     @Override
     public V read(K key) {
-        try {   // TODO: Consider using Optional rather than handling NPE
+        try { 
-            return read(key, map::read);
+            // TODO: Consider using Optional rather than handling NPE
+            final V entity = read(key, map::read);
+            return registerEntityForChanges(entity);
         } catch (NullPointerException ex) {
             return null;
         }
     }
 
-    @Override
     public V read(K key, Function<K, V> defaultValueFunc) {
         MapTaskWithValue current = tasks.get(key);
         // If the key exists, then it has entered the "tasks" after bulk delete that could have 
@@ -152,7 +175,8 @@ public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K>
         Stream<V> updatedAndNotRemovedObjectsStream = this.map.read(queryParameters)
           .filter(filterOutAllBulkDeletedObjects)
           .map(this::getUpdated)      // If the object has been removed, tx.get will return null, otherwise it will return me.getValue()
-          .filter(Objects::nonNull);
+          .filter(Objects::nonNull)
+          .map(this::registerEntityForChanges);
 
         // In case of created values stored in MapKeycloakTransaction, we need filter those according to the filter
         MapModelCriteriaBuilder<K, V, M> mapMcb = mcb.unwrap(MapModelCriteriaBuilder.class);
@@ -176,19 +200,11 @@ public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K>
         return read(queryParameters).count();
     }
 
-    @Override
+    private V getUpdated(V orig) {
-    public V getUpdated(V orig) {
         MapTaskWithValue current = orig == null ? null : tasks.get(orig.getId());
         return current == null ? orig : current.getValue();
     }
 
-    @Override
-    public V update(V value) {
-        K key = value.getId();
-        addTask(key, new UpdateOperation(value));
-        return value;
-    }
-
     @Override
     public V create(V value) {
         K key = value.getId();
@@ -196,7 +212,6 @@ public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K>
         return value;
     }
 
-    @Override
     public V updateIfChanged(V value, Predicate<V> shouldPut) {
         K key = value.getId();
         log.tracef("Adding operation UPDATE_IF_CHANGED for %s @ %08x", key, System.identityHashCode(value));
@@ -222,9 +237,11 @@ public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K>
 
 
     @Override
-    public long delete(K artificialKey, QueryParameters<M> queryParameters) {
+    public long delete(QueryParameters<M> queryParameters) {
         log.tracef("Adding operation DELETE_BULK");
 
+        K artificialKey = map.getKeyConvertor().yieldNewUniqueKey();
+
         // Remove all tasks that create / update / delete objects deleted by the bulk removal.
         final BulkDeleteOperation bdo = new BulkDeleteOperation(queryParameters);
         Predicate<V> filterForNonDeletedObjects = bdo.getFilterForNonDeletedObjects();
@@ -344,15 +361,6 @@ public class ConcurrentHashMapKeycloakTransaction<K, V extends AbstractEntity<K>
         @Override public MapOperation getOperation() { return MapOperation.CREATE; }
     }
 
-    private class UpdateOperation extends MapTaskWithValue {
-        public UpdateOperation(V value) {
-            super(value);
-        }
-
-        @Override public void execute() { map.update(getValue()); }
-        @Override public MapOperation getOperation() { return MapOperation.UPDATE; }
-    }
-
     private class DeleteOperation extends MapTaskWithValue {
         private final K key;
 

model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorage.java

@@ -19,6 +19,7 @@ package org.keycloak.models.map.storage.chm;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.map.storage.MapKeycloakTransaction;
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.models.map.storage.MapStorage;
 import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import org.keycloak.models.map.storage.QueryParameters;
@@ -42,7 +43,7 @@ import static org.keycloak.utils.StreamsUtil.paginatedStream;
  *
  * @author hmlnarik
  */
-public class ConcurrentHashMapStorage<K, V extends AbstractEntity<K>, M> implements MapStorage<K, V, M> {
+public class ConcurrentHashMapStorage<K, V extends AbstractEntity<K> & UpdatableEntity, M> implements MapStorage<K, V, M> {
 
     private final ConcurrentMap<K, V> store = new ConcurrentHashMap<>();
 

model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorageProvider.java

@@ -17,6 +17,7 @@
 package org.keycloak.models.map.storage.chm;
 
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.storage.MapStorage;
 import org.keycloak.models.map.storage.MapStorageProvider;
 import org.keycloak.models.map.storage.MapStorageProviderFactory.Flag;
 
@@ -37,8 +38,9 @@ public class ConcurrentHashMapStorageProvider implements MapStorageProvider {
     }
 
     @Override
-    public <K, V extends AbstractEntity<K>, M> ConcurrentHashMapStorage<K, V, M> getStorage(
+    @SuppressWarnings("unchecked")
-      Class<V> valueType, Class<M> modelType, Flag... flags) {
+    public <K, V extends AbstractEntity<K>, M> MapStorage<K, V, M> getStorage(Class<M> modelType, Flag... flags) {
-        return factory.getStorage(valueType, modelType, flags);
+        ConcurrentHashMapStorage storage = factory.getStorage(modelType, flags);
+        return (MapStorage<K, V, M>) storage;
     }
 }

model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorageProviderFactory.java

@@ -46,6 +46,7 @@ import org.keycloak.models.map.client.MapClientEntityImpl;
 import org.keycloak.models.map.clientscope.MapClientScopeEntity;
 import org.keycloak.models.map.common.AbstractEntity;
 import org.keycloak.models.map.common.Serialization;
+import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.models.map.group.MapGroupEntity;
 import org.keycloak.models.map.loginFailure.MapUserLoginFailureEntity;
 import org.keycloak.models.map.realm.MapRealmEntity;
@@ -114,7 +115,28 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
         MODEL_TO_NAME.put(Resource.class, "authz-resources");
         MODEL_TO_NAME.put(org.keycloak.authorization.model.Scope.class, "authz-scopes");
     }
+    
+    public static final Map<Class<?>, Class<? extends AbstractEntity>> MODEL_TO_VALUE_TYPE = new HashMap<>();
+    static {
+        MODEL_TO_VALUE_TYPE.put(AuthenticatedClientSessionModel.class, MapAuthenticatedClientSessionEntity.class);
+        MODEL_TO_VALUE_TYPE.put(ClientScopeModel.class, MapClientScopeEntity.class);
+        MODEL_TO_VALUE_TYPE.put(ClientModel.class, MapClientEntity.class);
+        MODEL_TO_VALUE_TYPE.put(GroupModel.class, MapGroupEntity.class);
+        MODEL_TO_VALUE_TYPE.put(RealmModel.class, MapRealmEntity.class);
+        MODEL_TO_VALUE_TYPE.put(RoleModel.class, MapRoleEntity.class);
+        MODEL_TO_VALUE_TYPE.put(RootAuthenticationSessionModel.class, MapRootAuthenticationSessionEntity.class);
+        MODEL_TO_VALUE_TYPE.put(UserLoginFailureModel.class, MapUserLoginFailureEntity.class);
+        MODEL_TO_VALUE_TYPE.put(UserModel.class, MapUserEntity.class);
+        MODEL_TO_VALUE_TYPE.put(UserSessionModel.class, MapUserSessionEntity.class);
 
+        // authz
+        MODEL_TO_VALUE_TYPE.put(PermissionTicket.class, MapPermissionTicketEntity.class);
+        MODEL_TO_VALUE_TYPE.put(Policy.class, MapPolicyEntity.class);
+        MODEL_TO_VALUE_TYPE.put(ResourceServer.class, MapResourceServerEntity.class);
+        MODEL_TO_VALUE_TYPE.put(Resource.class, MapResourceEntity.class);
+        MODEL_TO_VALUE_TYPE.put(org.keycloak.authorization.model.Scope.class, MapScopeEntity.class);
+    }
+    
     public static final Map<Class<?>, Class<?>> INTERFACE_TO_IMPL = new HashMap<>();
     static {
         INTERFACE_TO_IMPL.put(MapClientEntity.class, MapClientEntityImpl.class);
@@ -219,16 +241,16 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
         }
     }
 
-    private <K, V extends AbstractEntity<K>, M> ConcurrentHashMapStorage<K, V, M> loadMap(String mapName,
+    private <K, V extends AbstractEntity<K> & UpdatableEntity, M> ConcurrentHashMapStorage<K, V, M> loadMap(String mapName,
-      Class<V> valueType, Class<M> modelType, EnumSet<Flag> flags) {
+      Class<M> modelType, EnumSet<Flag> flags) {
         final StringKeyConvertor kc = keyConvertors.getOrDefault(mapName, defaultKeyConvertor);
-
+        Class<?> valueType = MODEL_TO_VALUE_TYPE.get(modelType);
         LOG.debugf("Initializing new map storage: %s", mapName);
 
         @SuppressWarnings("unchecked")
         ConcurrentHashMapStorage<K, V, M> store;
         if (modelType == UserSessionModel.class) {
-            ConcurrentHashMapStorage clientSessionStore = getStorage(MapAuthenticatedClientSessionEntity.class, AuthenticatedClientSessionModel.class);
+            ConcurrentHashMapStorage clientSessionStore = getStorage(AuthenticatedClientSessionModel.class);
             store = new UserSessionConcurrentHashMapStorage(clientSessionStore, kc) {
                 @Override
                 public String toString() {
@@ -269,8 +291,8 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
     }
 
     @SuppressWarnings("unchecked")
-    public <K, V extends AbstractEntity<K>, M> ConcurrentHashMapStorage<K, V, M> getStorage(
+    public <K, V extends AbstractEntity<K> & UpdatableEntity, M> ConcurrentHashMapStorage<K, V, M> getStorage(
-      Class<V> valueType, Class<M> modelType, Flag... flags) {
+      Class<M> modelType, Flag... flags) {
         EnumSet<Flag> f = flags == null || flags.length == 0 ? EnumSet.noneOf(Flag.class) : EnumSet.of(flags[0], flags);
         String name = MODEL_TO_NAME.getOrDefault(modelType, modelType.getSimpleName());
         /* From ConcurrentHashMapStorage.computeIfAbsent javadoc:
@@ -282,9 +304,9 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
          * to prepare clientSessionStore outside computeIfAbsent, otherwise deadlock occurs.
          */
         if (modelType == UserSessionModel.class) {
-            getStorage(MapAuthenticatedClientSessionEntity.class, AuthenticatedClientSessionModel.class);
+            getStorage(AuthenticatedClientSessionModel.class, flags);
         }
-        return (ConcurrentHashMapStorage<K, V, M>) storages.computeIfAbsent(name, n -> loadMap(name, valueType, modelType, f));
+        return (ConcurrentHashMapStorage<K, V, M>) storages.computeIfAbsent(name, n -> loadMap(name, modelType, f));
     }
 
     private File getFile(String fileName) {

model/map/src/main/java/org/keycloak/models/map/storage/chm/MapFieldPredicates.java

@@ -66,6 +66,7 @@ import java.util.function.Predicate;
 import java.util.stream.Stream;
 
 import org.keycloak.models.map.storage.CriterionNotSupportedException;
+import java.util.IdentityHashMap;
 import static org.keycloak.models.UserSessionModel.CORRESPONDING_SESSION_ID;
 
 /**
@@ -92,7 +93,7 @@ public class MapFieldPredicates {
 
     @SuppressWarnings("unchecked")
     private static final Map<Class<?>, Map> PREDICATES = new HashMap<>();
-    private static final Map<SearchableModelField<?>, Comparator<?>> COMPARATORS = new HashMap<>();
+    private static final Map<SearchableModelField<?>, Comparator<?>> COMPARATORS = new IdentityHashMap<>();
 
     static {
         put(REALM_PREDICATES, RealmModel.SearchableFields.NAME,                   MapRealmEntity::getName);

model/map/src/main/java/org/keycloak/models/map/storage/chm/UserSessionConcurrentHashMapStorage.java

@@ -52,19 +52,19 @@ public class UserSessionConcurrentHashMapStorage<K> extends ConcurrentHashMapSto
         }
 
         @Override
-        public long delete(K artificialKey, QueryParameters<UserSessionModel> queryParameters) {
+        public long delete(QueryParameters<UserSessionModel> queryParameters) {
             ModelCriteriaBuilder<UserSessionModel> mcb = queryParameters.getModelCriteriaBuilder();
 
             Set<K> ids = read(queryParameters).map(AbstractEntity::getId).collect(Collectors.toSet());
             ModelCriteriaBuilder<AuthenticatedClientSessionModel> csMcb = clientSessionStore.createCriteriaBuilder().compare(AuthenticatedClientSessionModel.SearchableFields.USER_SESSION_ID, Operator.IN, ids);
-            clientSessionTr.delete(artificialKey, withCriteria(csMcb));
+            clientSessionTr.delete(withCriteria(csMcb));
-            return super.delete(artificialKey, queryParameters);
+            return super.delete(queryParameters);
         }
 
         @Override
         public boolean delete(K key) {
             ModelCriteriaBuilder<AuthenticatedClientSessionModel> csMcb = clientSessionStore.createCriteriaBuilder().compare(AuthenticatedClientSessionModel.SearchableFields.USER_SESSION_ID, Operator.EQ, key);
-            clientSessionTr.delete(key, withCriteria(csMcb));
+            clientSessionTr.delete(withCriteria(csMcb));
             return super.delete(key);
         }
 

model/map/src/main/java/org/keycloak/models/map/user/MapUserEntity.java

@@ -19,6 +19,7 @@ package org.keycloak.models.map.user;
 
 import org.keycloak.models.ModelDuplicateException;
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.models.utils.KeycloakModelUtils;
 
 import java.util.Collection;
@@ -39,7 +40,7 @@ import java.util.stream.Stream;
  *
  * @author mhajas
  */
-public class MapUserEntity<K> implements AbstractEntity<K> {
+public class MapUserEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private final K id;
     private final String realmId;

model/map/src/main/java/org/keycloak/models/map/user/MapUserProvider.java

@@ -69,7 +69,6 @@ import static org.keycloak.models.UserModel.EMAIL_VERIFIED;
 import static org.keycloak.models.UserModel.FIRST_NAME;
 import static org.keycloak.models.UserModel.LAST_NAME;
 import static org.keycloak.models.UserModel.USERNAME;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 
@@ -89,7 +88,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
 
     private Function<MapUserEntity<K>, UserModel> entityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return origEntity -> new MapUserAdapter<K>(session, realm, registerEntityForChanges(tx, origEntity)) {
+        return origEntity -> new MapUserAdapter<K>(session, realm, origEntity) {
             @Override
             public String getId() {
                 return userStore.getKeyConvertor().keyToString(entity.getId());
@@ -127,9 +126,8 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
         }
     }
 
-    private MapUserEntity<K> getRegisteredEntityByIdOrThrow(RealmModel realm, String id) {
+    private MapUserEntity<K> getEntityByIdOrThrow(RealmModel realm, String id) {
         return getEntityById(realm, id)
-                .map(e -> registerEntityForChanges(tx, e))
                 .orElseThrow(this::userDoesntExistException);
     }
 
@@ -142,10 +140,6 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
         return Optional.empty();
     }
 
-    private Optional<MapUserEntity<K>> getRegisteredEntityById(RealmModel realm, String id) {
-        return getEntityById(realm, id).map(e -> registerEntityForChanges(tx, e));
-    }
-
     @Override
     public void addFederatedIdentity(RealmModel realm, UserModel user, FederatedIdentityModel socialLink) {
         if (user == null || user.getId() == null) {
@@ -153,7 +147,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
         }
         LOG.tracef("addFederatedIdentity(%s, %s, %s)%s", realm, user.getId(), socialLink.getIdentityProvider(), getShortStackTrace());
 
-        getRegisteredEntityById(realm, user.getId())
+        getEntityById(realm, user.getId())
                 .ifPresent(userEntity ->
                         userEntity.addFederatedIdentity(UserFederatedIdentityEntity.fromModel(socialLink)));
     }
@@ -161,7 +155,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
     @Override
     public boolean removeFederatedIdentity(RealmModel realm, UserModel user, String socialProvider) {
         LOG.tracef("removeFederatedIdentity(%s, %s, %s)%s", realm, user.getId(), socialProvider, getShortStackTrace());
-        return getRegisteredEntityById(realm, user.getId())
+        return getEntityById(realm, user.getId())
                 .map(entity -> entity.removeFederatedIdentity(socialProvider))
                 .orElse(false);
     }
@@ -175,14 +169,13 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
           .compare(SearchableFields.IDP_AND_USER, Operator.EQ, socialProvider);
 
         tx.read(withCriteria(mcb))
-                .map(e -> registerEntityForChanges(tx, e))
                 .forEach(userEntity -> userEntity.removeFederatedIdentity(socialProvider));
     }
 
     @Override
     public void updateFederatedIdentity(RealmModel realm, UserModel federatedUser, FederatedIdentityModel federatedIdentityModel) {
         LOG.tracef("updateFederatedIdentity(%s, %s, %s)%s", realm, federatedUser.getId(), federatedIdentityModel.getIdentityProvider(), getShortStackTrace());
-        getRegisteredEntityById(realm, federatedUser.getId())
+        getEntityById(realm, federatedUser.getId())
                 .ifPresent(entity -> entity.updateFederatedIdentity(UserFederatedIdentityEntity.fromModel(federatedIdentityModel)));
     }
 
@@ -229,7 +222,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
     public void addConsent(RealmModel realm, String userId, UserConsentModel consent) {
         LOG.tracef("addConsent(%s, %s, %s)%s", realm, userId, consent, getShortStackTrace());
 
-        getRegisteredEntityByIdOrThrow(realm, userId)
+        getEntityByIdOrThrow(realm, userId)
                 .addUserConsent(UserConsentEntity.fromModel(consent));
     }
 
@@ -255,7 +248,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
     public void updateConsent(RealmModel realm, String userId, UserConsentModel consent) {
         LOG.tracef("updateConsent(%s, %s, %s)%s", realm, userId, consent, getShortStackTrace());
 
-        MapUserEntity<K> user = getRegisteredEntityByIdOrThrow(realm, userId);
+        MapUserEntity<K> user = getEntityByIdOrThrow(realm, userId);
         UserConsentEntity userConsentEntity = user.getUserConsent(consent.getClient().getId());
         if (userConsentEntity == null) {
             throw new ModelException("Consent not found for client [" + consent.getClient().getId() + "] and user [" + userId + "]");
@@ -273,7 +266,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
     @Override
     public boolean revokeConsentForClient(RealmModel realm, String userId, String clientInternalId) {
         LOG.tracef("revokeConsentForClient(%s, %s, %s)%s", realm, userId, clientInternalId, getShortStackTrace());
-        return getRegisteredEntityById(realm, userId)
+        return getEntityById(realm, userId)
                 .map(userEntity -> userEntity.removeUserConsent(clientInternalId))
                 .orElse(false);
     }
@@ -281,7 +274,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
     @Override
     public void setNotBeforeForUser(RealmModel realm, UserModel user, int notBefore) {
         LOG.tracef("setNotBeforeForUser(%s, %s, %d)%s", realm, user.getId(), notBefore, getShortStackTrace());
-        getRegisteredEntityByIdOrThrow(realm, user.getId()).setNotBefore(notBefore);
+        getEntityByIdOrThrow(realm, user.getId()).setNotBefore(notBefore);
     }
 
     @Override
@@ -363,7 +356,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
         ModelCriteriaBuilder<UserModel> mcb = userStore.createCriteriaBuilder()
           .compare(SearchableFields.REALM_ID, Operator.EQ, realm.getId());
 
-        tx.delete(userStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        tx.delete(withCriteria(mcb));
     }
 
     @Override
@@ -373,7 +366,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
           .compare(SearchableFields.REALM_ID, Operator.EQ, realm.getId())
           .compare(SearchableFields.FEDERATION_LINK, Operator.EQ, storageProviderId);
 
-        tx.delete(userStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        tx.delete(withCriteria(mcb));
     }
 
     @Override
@@ -384,8 +377,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
           .compare(SearchableFields.FEDERATION_LINK, Operator.EQ, storageProviderId);
 
         try (Stream<MapUserEntity<K>> s = tx.read(withCriteria(mcb))) {
-            s.map(e -> registerEntityForChanges(tx, e))
+            s.forEach(userEntity -> userEntity.setFederationLink(null));
-              .forEach(userEntity -> userEntity.setFederationLink(null));
         }
     }
 
@@ -398,8 +390,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
           .compare(SearchableFields.ASSIGNED_ROLE, Operator.EQ, roleId);
 
         try (Stream<MapUserEntity<K>> s = tx.read(withCriteria(mcb))) {
-            s.map(e -> registerEntityForChanges(tx, e))
+            s.forEach(userEntity -> userEntity.removeRolesMembership(roleId));
-              .forEach(userEntity -> userEntity.removeRolesMembership(roleId));
         }
     }
 
@@ -412,8 +403,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
           .compare(SearchableFields.ASSIGNED_GROUP, Operator.EQ, groupId);
 
         try (Stream<MapUserEntity<K>> s = tx.read(withCriteria(mcb))) {
-            s.map(e -> registerEntityForChanges(tx, e))
+            s.forEach(userEntity -> userEntity.removeGroupsMembership(groupId));
-              .forEach(userEntity -> userEntity.removeGroupsMembership(groupId));
         }
     }
 
@@ -426,8 +416,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
           .compare(SearchableFields.CONSENT_FOR_CLIENT, Operator.EQ, clientId);
 
         try (Stream<MapUserEntity<K>> s = tx.read(withCriteria(mcb))) {
-            s.map(e -> registerEntityForChanges(tx, e))
+            s.forEach(userEntity -> userEntity.removeUserConsent(clientId));
-              .forEach(userEntity -> userEntity.removeUserConsent(clientId));
         }
     }
 
@@ -478,7 +467,6 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
               .collect(Collectors.toList());
 
             if (! consentClientIds.isEmpty()) {
-                userEntity = registerEntityForChanges(tx, userEntity);
                 consentClientIds.forEach(userEntity::removeUserConsent);
             }
         };
@@ -492,8 +480,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
           .compare(SearchableFields.REALM_ID, Operator.EQ, realm.getId());
 
         try (Stream<MapUserEntity<K>> s = tx.read(withCriteria(mcb))) {
-            s.map(e -> registerEntityForChanges(tx, e))
+            s.forEach(entity -> entity.addRolesMembership(roleId));
-              .forEach(entity -> entity.addRolesMembership(roleId));
         }
     }
 
@@ -535,7 +522,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
             throw new ModelDuplicateException("Multiple users with email '" + email + "' exist in Keycloak.");
         }
 
-        MapUserEntity<K> userEntity = registerEntityForChanges(tx, usersWithEmail.get(0));
+        MapUserEntity<K> userEntity = usersWithEmail.get(0);
         
         if (!realm.isDuplicateEmailsAllowed()) {
             if (userEntity.getEmail() != null && !userEntity.getEmail().equals(userEntity.getEmailConstraint())) {
@@ -545,21 +532,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
             }
         }
         
-        return new MapUserAdapter<K>(session, realm, userEntity) {
+        return entityToAdapterFunc(realm).apply(userEntity);
-            @Override
-            public String getId() {
-                return userStore.getKeyConvertor().keyToString(userEntity.getId());
-            }
-
-            @Override
-            public boolean checkEmailUniqueness(RealmModel realm, String email) {
-                return getUserByEmail(realm, email) != null;
-            }
-            @Override
-            public boolean checkUsernameUniqueness(RealmModel realm, String username) {
-                return getUserByUsername(realm, username) != null;
-            }
-        };
     }
 
     @Override
@@ -758,7 +731,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
 
     @Override
     public void updateCredential(RealmModel realm, UserModel user, CredentialModel cred) {
-        getRegisteredEntityById(realm, user.getId())
+        getEntityById(realm, user.getId())
                 .ifPresent(updateCredential(cred));
     }
     
@@ -780,7 +753,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
         LOG.tracef("createCredential(%s, %s, %s)%s", realm, user.getId(), cred.getId(), getShortStackTrace());
         UserCredentialEntity credentialEntity = UserCredentialEntity.fromModel(cred);
 
-        getRegisteredEntityByIdOrThrow(realm, user.getId())
+        getEntityByIdOrThrow(realm, user.getId())
                 .addCredential(credentialEntity);
 
         return UserCredentialEntity.toModel(credentialEntity);
@@ -789,7 +762,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
     @Override
     public boolean removeStoredCredential(RealmModel realm, UserModel user, String id) {
         LOG.tracef("removeStoredCredential(%s, %s, %s)%s", realm, user.getId(), id, getShortStackTrace());
-        return getRegisteredEntityById(realm, user.getId())
+        return getEntityById(realm, user.getId())
                 .map(mapUserEntity -> mapUserEntity.removeCredential(id))
                 .orElse(false);
     }
@@ -831,7 +804,7 @@ public class MapUserProvider<K> implements UserProvider.Streams, UserCredentialS
     public boolean moveCredentialTo(RealmModel realm, UserModel user, String id, String newPreviousCredentialId) {
         LOG.tracef("moveCredentialTo(%s, %s, %s, %s)%s", realm, user.getId(), id, newPreviousCredentialId, getShortStackTrace());
         String userId = user.getId();
-        MapUserEntity<K> userEntity = getRegisteredEntityById(realm, userId).orElse(null);
+        MapUserEntity<K> userEntity = getEntityById(realm, userId).orElse(null);
         if (userEntity == null) {
             LOG.warnf("User with id: [%s] not found", userId);
             return false;

model/map/src/main/java/org/keycloak/models/map/user/MapUserProviderFactory.java

@@ -30,7 +30,7 @@ import org.keycloak.models.map.common.AbstractMapProviderFactory;
 public class MapUserProviderFactory<K> extends AbstractMapProviderFactory<UserProvider, K, MapUserEntity<K>, UserModel> implements UserProviderFactory {
 
     public MapUserProviderFactory() {
-        super(MapUserEntity.class, UserModel.class);
+        super(UserModel.class);
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/userSession/MapAuthenticatedClientSessionEntity.java

@@ -19,6 +19,7 @@ package org.keycloak.models.map.userSession;
 import org.keycloak.common.util.Time;
 import org.keycloak.models.map.common.AbstractEntity;
 
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
@@ -26,7 +27,7 @@ import java.util.concurrent.ConcurrentHashMap;
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
  */
-public class MapAuthenticatedClientSessionEntity<K> implements AbstractEntity<K> {
+public class MapAuthenticatedClientSessionEntity<K> implements AbstractEntity<K>, UpdatableEntity {
 
     private K id;
     private String userSessionId;

model/map/src/main/java/org/keycloak/models/map/userSession/MapUserSessionEntity.java

@@ -22,6 +22,7 @@ import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.map.common.AbstractEntity;
 
+import org.keycloak.models.map.common.UpdatableEntity;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
@@ -29,7 +30,7 @@ import java.util.concurrent.ConcurrentHashMap;
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
  */
-public class MapUserSessionEntity<K> implements AbstractEntity<K> {
+public class MapUserSessionEntity<K> implements AbstractEntity<K>, UpdatableEntity {
     private K id;
 
     private String realmId;

model/map/src/main/java/org/keycloak/models/map/userSession/MapUserSessionProvider.java

@@ -46,7 +46,6 @@ import java.util.stream.Stream;
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
 import static org.keycloak.models.UserSessionModel.CORRESPONDING_SESSION_ID;
 import static org.keycloak.models.UserSessionModel.SessionPersistenceState.TRANSIENT;
-import static org.keycloak.models.map.common.MapStorageUtils.registerEntityForChanges;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 import static org.keycloak.models.map.userSession.SessionExpiration.setClientSessionExpiration;
 import static org.keycloak.models.map.userSession.SessionExpiration.setUserSessionExpiration;
@@ -90,8 +89,7 @@ public class MapUserSessionProvider<UK, CK> implements UserSessionProvider {
                 userSessionTx.delete(origEntity.getId());
                 return null;
             } else {
-                return new MapUserSessionAdapter<UK>(session, realm,
+                return new MapUserSessionAdapter<UK>(session, realm, origEntity) {
-                        Objects.equals(origEntity.getPersistenceState(), TRANSIENT) ? origEntity : registerEntityForChanges(userSessionTx, origEntity)) {
                     @Override
                     public String getId() {
                         return userSessionStore.getKeyConvertor().keyToString(entity.getId());
@@ -123,7 +121,7 @@ public class MapUserSessionProvider<UK, CK> implements UserSessionProvider {
                 clientSessionTx.delete(origEntity.getId());
                 return null;
             } else {
-                return new MapAuthenticatedClientSessionAdapter<CK>(session, realm, client, userSession, registerEntityForChanges(clientSessionTx, origEntity)) {
+                return new MapAuthenticatedClientSessionAdapter<CK>(session, realm, client, userSession, origEntity) {
                     @Override
                     public String getId() {
                         return clientSessionStore.getKeyConvertor().keyToString(entity.getId());
@@ -382,7 +380,7 @@ public class MapUserSessionProvider<UK, CK> implements UserSessionProvider {
 
         LOG.tracef("removeUserSession(%s, %s)%s", realm, session, getShortStackTrace());
 
-        userSessionTx.delete(userSessionStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        userSessionTx.delete(withCriteria(mcb));
     }
 
     @Override
@@ -393,7 +391,7 @@ public class MapUserSessionProvider<UK, CK> implements UserSessionProvider {
 
         LOG.tracef("removeUserSessions(%s, %s)%s", realm, user, getShortStackTrace());
 
-        userSessionTx.delete(userSessionStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        userSessionTx.delete(withCriteria(mcb));
     }
 
     @Override
@@ -412,7 +410,7 @@ public class MapUserSessionProvider<UK, CK> implements UserSessionProvider {
 
         LOG.tracef("removeUserSessions(%s)%s", realm, getShortStackTrace());
 
-        userSessionTx.delete(userSessionStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+        userSessionTx.delete(withCriteria(mcb));
     }
 
     @Override
@@ -469,7 +467,7 @@ public class MapUserSessionProvider<UK, CK> implements UserSessionProvider {
             UK uk = userSessionStore.getKeyConvertor().fromString(userSession.getNote(CORRESPONDING_SESSION_ID));
             mcb = realmAndOfflineCriteriaBuilder(realm, true)
                     .compare(UserSessionModel.SearchableFields.ID, ModelCriteriaBuilder.Operator.EQ, uk);
-            userSessionTx.delete(userSessionStore.getKeyConvertor().yieldNewUniqueKey(), withCriteria(mcb));
+            userSessionTx.delete(withCriteria(mcb));
             userSession.removeNote(CORRESPONDING_SESSION_ID);
         }
     }
@@ -637,7 +635,7 @@ public class MapUserSessionProvider<UK, CK> implements UserSessionProvider {
 
         if (userSessionEntity == null) {
             MapUserSessionEntity<UK> userSession = userSessionTx.read(id);
-            return userSession != null ? registerEntityForChanges(userSessionTx, userSession) : null;
+            return userSession;
         }
         return userSessionEntity;
     }

model/map/src/main/java/org/keycloak/models/map/userSession/MapUserSessionProviderFactory.java

@@ -83,12 +83,12 @@ public class MapUserSessionProviderFactory<UK, CK> implements AmphibianProviderF
         MapStorageProviderFactory storageProviderFactoryUs = (MapStorageProviderFactory) getComponentFactory(session.getKeycloakSessionFactory(),
           MapStorageProvider.class, storageConfigScopeUserSessions, MapStorageSpi.NAME);
         final MapStorageProvider factoryUs = storageProviderFactoryUs.create(session);
-        MapStorage userSessionStore = factoryUs.getStorage(MapUserSessionEntity.class, UserSessionModel.class);
+        MapStorage userSessionStore = factoryUs.getStorage(UserSessionModel.class);
 
         MapStorageProviderFactory storageProviderFactoryCs = (MapStorageProviderFactory) getComponentFactory(session.getKeycloakSessionFactory(),
           MapStorageProvider.class, storageConfigScopeClientSessions, MapStorageSpi.NAME);
         final MapStorageProvider factoryCs = storageProviderFactoryCs.create(session);
-        MapStorage clientSessionStore = factoryCs.getStorage(MapAuthenticatedClientSessionEntity.class, AuthenticatedClientSessionModel.class);
+        MapStorage clientSessionStore = factoryCs.getStorage(AuthenticatedClientSessionModel.class);
 
         return new MapUserSessionProvider<>(session, userSessionStore, clientSessionStore);
     }

server-spi/src/main/java/org/keycloak/storage/SearchableModelField.java

@@ -16,8 +16,6 @@
  */
 package org.keycloak.storage;
 
-import java.util.Objects;
-
 /**
  *
  * @author hmlnarik

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

@@ -82,9 +82,9 @@ public class MapStorageTest extends KeycloakModelTest {
         String component2Id = createMapStorageComponent("component2", "keyType", "string");
 
         Object[] ids = withRealm(realmId, (session, realm) -> {
-            MapStorage<K, MapClientEntity<K>, ClientModel> storageMain = (MapStorage<K, MapClientEntity<K>, ClientModel>) session.getProvider(MapStorageProvider.class).getStorage(MapClientEntity.class, ClientModel.class);
+            MapStorage<K, MapClientEntity<K>, ClientModel> storageMain = (MapStorage) session.getProvider(MapStorageProvider.class).getStorage(ClientModel.class);
-            MapStorage<K1, MapClientEntity<K1>, ClientModel> storage1 = (MapStorage<K1, MapClientEntity<K1>, ClientModel>) session.getComponentProvider(MapStorageProvider.class, component1Id).getStorage(MapClientEntity.class, ClientModel.class);
+            MapStorage<K1, MapClientEntity<K1>, ClientModel> storage1 = (MapStorage) session.getComponentProvider(MapStorageProvider.class, component1Id).getStorage(ClientModel.class);
-            MapStorage<K2, MapClientEntity<K2>, ClientModel> storage2 = (MapStorage<K2, MapClientEntity<K2>, ClientModel>) session.getComponentProvider(MapStorageProvider.class, component2Id).getStorage(MapClientEntity.class, ClientModel.class);
+            MapStorage<K2, MapClientEntity<K2>, ClientModel> storage2 = (MapStorage) session.getComponentProvider(MapStorageProvider.class, component2Id).getStorage(ClientModel.class);
 
             // Assert that the map storage can be used both as a standalone store and a component
             assertThat(storageMain, notNullValue());
@@ -161,11 +161,11 @@ public class MapStorageTest extends KeycloakModelTest {
         // Check that in the next transaction, the objects are still there
         withRealm(realmId, (session, realm) -> {
             @SuppressWarnings("unchecked")
-            MapStorage<K, MapClientEntity<K>, ClientModel> storageMain = (MapStorage<K, MapClientEntity<K>, ClientModel>) session.getProvider(MapStorageProvider.class).getStorage(MapClientEntity.class, ClientModel.class);
+            MapStorage<K, MapClientEntity<K>, ClientModel> storageMain = (MapStorage) session.getProvider(MapStorageProvider.class).getStorage(ClientModel.class);
             @SuppressWarnings("unchecked")
-            MapStorage<K1, MapClientEntity<K1>, ClientModel> storage1 = (MapStorage<K1, MapClientEntity<K1>, ClientModel>) session.getComponentProvider(MapStorageProvider.class, component1Id).getStorage(MapClientEntity.class, ClientModel.class);
+            MapStorage<K1, MapClientEntity<K1>, ClientModel> storage1 = (MapStorage) session.getComponentProvider(MapStorageProvider.class, component1Id).getStorage(ClientModel.class);
             @SuppressWarnings("unchecked")
-            MapStorage<K2, MapClientEntity<K2>, ClientModel> storage2 = (MapStorage<K2, MapClientEntity<K2>, ClientModel>) session.getComponentProvider(MapStorageProvider.class, component2Id).getStorage(MapClientEntity.class, ClientModel.class);
+            MapStorage<K2, MapClientEntity<K2>, ClientModel> storage2 = (MapStorage) session.getComponentProvider(MapStorageProvider.class, component2Id).getStorage(ClientModel.class);
 
             final StringKeyConvertor<K> kcMain = storageMain.getKeyConvertor();
             final StringKeyConvertor<K1> kc1 = storage1.getKeyConvertor();