Updates Datastore provider to contain full data model

Closes #15490

model/legacy-private/src/main/java/org/keycloak/storage/datastore/LegacyDatastoreProvider.java

@@ -23,9 +23,13 @@ import org.keycloak.models.GroupProvider;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmProvider;
 import org.keycloak.models.RoleProvider;
+import org.keycloak.models.SingleUseObjectProvider;
+import org.keycloak.models.UserLoginFailureProvider;
 import org.keycloak.models.UserProvider;
+import org.keycloak.models.UserSessionProvider;
 import org.keycloak.models.cache.CacheRealmProvider;
 import org.keycloak.models.cache.UserCache;
+import org.keycloak.sessions.AuthenticationSessionProvider;
 import org.keycloak.storage.ClientScopeStorageManager;
 import org.keycloak.storage.ClientStorageManager;
 import org.keycloak.storage.DatastoreProvider;
@@ -41,12 +45,16 @@ public class LegacyDatastoreProvider implements DatastoreProvider, LegacyStoreMa
     private final LegacyDatastoreProviderFactory factory;
     private final KeycloakSession session;
 
+    private AuthenticationSessionProvider authenticationSessionProvider;
     private ClientProvider clientProvider;
     private ClientScopeProvider clientScopeProvider;
     private GroupProvider groupProvider;
+    private UserLoginFailureProvider userLoginFailureProvider;
     private RealmProvider realmProvider;
     private RoleProvider roleProvider;
+    private SingleUseObjectProvider singleUseObjectProvider;
     private UserProvider userProvider;
+    private UserSessionProvider userSessionProvider;
 
     private ClientScopeStorageManager clientScopeStorageManager;
     private RoleStorageManager roleStorageManager;
@@ -170,6 +178,14 @@ public class LegacyDatastoreProvider implements DatastoreProvider, LegacyStoreMa
         }
     }
 
+    @Override
+    public AuthenticationSessionProvider authSessions() {
+        if (authenticationSessionProvider == null) {
+            authenticationSessionProvider = session.getProvider(AuthenticationSessionProvider.class);
+        }
+        return authenticationSessionProvider;
+    }
+
     @Override
     public ClientProvider clients() {
         if (clientProvider == null) {
@@ -194,6 +210,14 @@ public class LegacyDatastoreProvider implements DatastoreProvider, LegacyStoreMa
         return groupProvider;
     }
 
+    @Override
+    public UserLoginFailureProvider loginFailures() {
+        if (userLoginFailureProvider == null) {
+            userLoginFailureProvider = session.getProvider(UserLoginFailureProvider.class);
+        }
+        return userLoginFailureProvider;
+    }
+
     @Override
     public RealmProvider realms() {
         if (realmProvider == null) {
@@ -210,6 +234,14 @@ public class LegacyDatastoreProvider implements DatastoreProvider, LegacyStoreMa
         return roleProvider;
     }
 
+    @Override
+    public SingleUseObjectProvider singleUseObjects() {
+        if (singleUseObjectProvider == null) {
+            singleUseObjectProvider = session.getProvider(SingleUseObjectProvider.class);
+        }
+        return singleUseObjectProvider;
+    }
+
     @Override
     public UserProvider users() {
         if (userProvider == null) {
@@ -218,6 +250,14 @@ public class LegacyDatastoreProvider implements DatastoreProvider, LegacyStoreMa
         return userProvider;
     }
 
+    @Override
+    public UserSessionProvider userSessions() {
+        if (userSessionProvider == null) {
+            userSessionProvider = session.getProvider(UserSessionProvider.class);
+        }
+        return userSessionProvider;
+    }
+
     @Override
     public ExportImportManager getExportImportManager() {
         return new LegacyExportImportManager(session);

model/map/src/main/java/org/keycloak/models/map/datastore/ImportKeycloakSession.java

@@ -36,6 +36,7 @@ import org.keycloak.models.RealmProvider;
 import org.keycloak.models.RealmSpi;
 import org.keycloak.models.RoleProvider;
 import org.keycloak.models.RoleSpi;
+import org.keycloak.models.SingleUseObjectProvider;
 import org.keycloak.models.ThemeManager;
 import org.keycloak.models.TokenManager;
 import org.keycloak.models.UserCredentialManager;
@@ -264,6 +265,11 @@ public class ImportKeycloakSession implements KeycloakSession {
         throw new ModelException("not supported yet");
     }
 
+    @Override
+    public SingleUseObjectProvider singleUseObjects() {
+        throw new ModelException("not supported yet");
+    }
+
     @Override
     public void close() {
         session.close();

model/map/src/main/java/org/keycloak/models/map/datastore/MapDatastoreProvider.java

@@ -23,7 +23,11 @@ import org.keycloak.models.GroupProvider;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmProvider;
 import org.keycloak.models.RoleProvider;
+import org.keycloak.models.SingleUseObjectProvider;
+import org.keycloak.models.UserLoginFailureProvider;
 import org.keycloak.models.UserProvider;
+import org.keycloak.models.UserSessionProvider;
+import org.keycloak.sessions.AuthenticationSessionProvider;
 import org.keycloak.storage.DatastoreProvider;
 import org.keycloak.storage.ExportImportManager;
 
@@ -39,6 +43,11 @@ public class MapDatastoreProvider implements DatastoreProvider {
     public void close() {
     }
 
+    @Override
+    public AuthenticationSessionProvider authSessions() {
+        return session.getProvider(AuthenticationSessionProvider.class);
+    }
+
     @Override
     public ClientScopeProvider clientScopes() {
         return session.getProvider(ClientScopeProvider.class);
@@ -54,6 +63,11 @@ public class MapDatastoreProvider implements DatastoreProvider {
         return session.getProvider(GroupProvider.class);
     }
 
+    @Override
+    public UserLoginFailureProvider loginFailures() {
+        return session.getProvider(UserLoginFailureProvider.class);
+    }
+
     @Override
     public RealmProvider realms() {
         return session.getProvider(RealmProvider.class);
@@ -64,11 +78,21 @@ public class MapDatastoreProvider implements DatastoreProvider {
         return session.getProvider(RoleProvider.class);
     }
 
+    @Override
+    public SingleUseObjectProvider singleUseObjects() {
+        return session.getProvider(SingleUseObjectProvider.class);
+    }
+
     @Override
     public UserProvider users() {
         return session.getProvider(UserProvider.class);
     }
 
+    @Override
+    public UserSessionProvider userSessions() {
+        return session.getProvider(UserSessionProvider.class);
+    }
+
     @Override
     public ExportImportManager getExportImportManager() {
         return new MapExportImportManager(session);

server-spi-private/src/main/java/org/keycloak/storage/DatastoreProvider.java

@@ -22,11 +22,16 @@ import org.keycloak.models.ClientScopeProvider;
 import org.keycloak.models.GroupProvider;
 import org.keycloak.models.RealmProvider;
 import org.keycloak.models.RoleProvider;
+import org.keycloak.models.SingleUseObjectProvider;
+import org.keycloak.models.UserLoginFailureProvider;
 import org.keycloak.models.UserProvider;
+import org.keycloak.models.UserSessionProvider;
 import org.keycloak.provider.Provider;
+import org.keycloak.sessions.AuthenticationSessionProvider;
 
 
 public interface DatastoreProvider extends Provider {
+    AuthenticationSessionProvider authSessions();
 
     ClientScopeProvider clientScopes();
 
@@ -34,11 +39,17 @@ public interface DatastoreProvider extends Provider {
 
     GroupProvider groups();
 
+    UserLoginFailureProvider loginFailures();
+
     RealmProvider realms();
 
     RoleProvider roles();
+
+    SingleUseObjectProvider singleUseObjects();
     
     UserProvider users();
 
+    UserSessionProvider userSessions();
+
     ExportImportManager getExportImportManager();
 }

server-spi/src/main/java/org/keycloak/models/KeycloakSession.java

@@ -201,6 +201,8 @@ public interface KeycloakSession extends AutoCloseable {
 
     AuthenticationSessionProvider authenticationSessions();
 
+    SingleUseObjectProvider singleUseObjects();
+
 
 
     void close();

services/src/main/java/org/keycloak/authentication/authenticators/client/JWTClientAuthenticator.java

@@ -190,7 +190,7 @@ public class JWTClientAuthenticator extends AbstractClientAuthenticator {
                 throw new RuntimeException("Missing ID on the token");
             }
 
-            SingleUseObjectProvider singleUseCache = context.getSession().getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseCache = context.getSession().singleUseObjects();
             int lifespanInSecs = Math.max(token.getExpiration() - currentTime, 10);
             if (singleUseCache.putIfAbsent(token.getId(), lifespanInSecs)) {
                 logger.tracef("Added token '%s' to single-use cache. Lifespan: %d seconds, client: %s", token.getId(), lifespanInSecs, clientId);

services/src/main/java/org/keycloak/authentication/authenticators/client/JWTClientSecretAuthenticator.java

@@ -192,7 +192,7 @@ public class JWTClientSecretAuthenticator extends AbstractClientAuthenticator {
                 throw new RuntimeException("Missing ID on the token");
             }
 
-            SingleUseObjectProvider singleUseCache = context.getSession().getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseCache = context.getSession().singleUseObjects();
             int lifespanInSecs = Math.max(token.getExpiration() - currentTime, 10);
             if (singleUseCache.putIfAbsent(token.getId(), lifespanInSecs)) {
 

services/src/main/java/org/keycloak/credential/OTPCredentialProvider.java

@@ -117,7 +117,7 @@ public class OTPCredentialProvider implements CredentialProvider<OTPCredentialMo
             if (isValid) {
                 if (policy.isCodeReusable()) return true;
 
-                SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+                SingleUseObjectProvider singleUseStore = session.singleUseObjects();
                 final long validLifespan = (long) credentialData.getPeriod() * (2L * policy.getLookAheadWindow() + 1);
                 final String searchKey = credential.getId() + "." + challengeResponse;
 

services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java

@@ -1342,7 +1342,7 @@ public class TokenManager {
 
         @Override
         public boolean test(JsonWebToken token) {
-            SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseStore = session.singleUseObjects();
             return !singleUseStore.contains(token.getId() + SingleUseObjectProvider.REVOKED_KEY);
         }
     }

services/src/main/java/org/keycloak/protocol/oidc/endpoints/TokenRevocationEndpoint.java

@@ -255,7 +255,7 @@ public class TokenRevocationEndpoint {
     }
 
     private void revokeAccessToken() {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         int currentTime = Time.currentTime();
         long lifespanInSecs = Math.max(token.getExp() - currentTime, 10);
         singleUseStore.put(token.getId() + SingleUseObjectProvider.REVOKED_KEY, lifespanInSecs, Collections.emptyMap());

services/src/main/java/org/keycloak/protocol/oidc/grants/ciba/endpoints/BackchannelAuthenticationEndpoint.java

@@ -152,7 +152,7 @@ public class BackchannelAuthenticationEndpoint extends AbstractCibaEndpoint {
         // To inform "expired_token" to the client, the lifespan of the cache provider is longer than device code
         int lifespanSeconds = expiresIn + poolingInterval + 10;
 
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
 
         singleUseStore.put(deviceCode.serializeKey(), lifespanSeconds, deviceCode.toMap());
         singleUseStore.put(userCode.serializeKey(), lifespanSeconds, userCode.serializeValue());

services/src/main/java/org/keycloak/protocol/oidc/grants/device/DeviceGrantType.java

@@ -146,28 +146,28 @@ public class DeviceGrantType {
     }
 
     public static OAuth2DeviceCodeModel getDeviceByDeviceCode(KeycloakSession session, RealmModel realm, String deviceCode) {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         Map<String, String> notes = singleUseStore.get(OAuth2DeviceCodeModel.createKey(deviceCode));
         return notes != null ? OAuth2DeviceCodeModel.fromCache(realm, deviceCode, notes) : null;
     }
 
     public static void removeDeviceByDeviceCode(KeycloakSession session, String deviceCode) {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         singleUseStore.remove(OAuth2DeviceCodeModel.createKey(deviceCode));
     }
 
     public static void removeDeviceByUserCode(KeycloakSession session, RealmModel realm, String userCode) {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         singleUseStore.remove(OAuth2DeviceUserCodeModel.createKey(realm, userCode));
     }
 
     public static boolean isPollingAllowed(KeycloakSession session, OAuth2DeviceCodeModel deviceCodeModel) {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         return singleUseStore.putIfAbsent(deviceCodeModel.serializePollingKey(), deviceCodeModel.getPollingInterval());
     }
 
     public static boolean approveUserCode(KeycloakSession session, RealmModel realm, String userCode, String userSessionId, Map<String, String> additionalParams) {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         OAuth2DeviceCodeModel deviceCodeModel = DeviceEndpoint.getDeviceByUserCode(session, realm, userCode);
 
         if (deviceCodeModel != null) {
@@ -179,7 +179,7 @@ public class DeviceGrantType {
     }
 
     public static boolean denyUserCode(KeycloakSession session, RealmModel realm, String userCode) {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         OAuth2DeviceCodeModel deviceCodeModel = DeviceEndpoint.getDeviceByUserCode(session, realm, userCode);
 
         if (deviceCodeModel != null) {

services/src/main/java/org/keycloak/protocol/oidc/grants/device/endpoints/DeviceEndpoint.java

@@ -164,7 +164,7 @@ public class DeviceEndpoint extends AuthorizationEndpointBase implements RealmRe
         // To inform "expired_token" to the client, the lifespan of the cache provider is longer than device code
         int lifespanSeconds = expiresIn + interval + 10;
 
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
 
         singleUseStore.put(deviceCode.serializeKey(), lifespanSeconds, deviceCode.toMap());
         singleUseStore.put(userCode.serializeKey(), lifespanSeconds, userCode.serializeValue());
@@ -292,7 +292,7 @@ public class DeviceEndpoint extends AuthorizationEndpointBase implements RealmRe
     }
 
     public static OAuth2DeviceCodeModel getDeviceByUserCode(KeycloakSession session, RealmModel realm, String userCode) {
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         Map<String, String> notes = singleUseStore.get(OAuth2DeviceUserCodeModel.createKey(realm, userCode));
 
         if (notes != null) {

services/src/main/java/org/keycloak/protocol/oidc/par/endpoints/ParEndpoint.java

@@ -157,7 +157,7 @@ public class ParEndpoint extends AbstractParEndpoint {
             });
         params.put(PAR_CREATED_TIME, String.valueOf(System.currentTimeMillis()));
 
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         singleUseStore.put(key, expiresIn, params);
 
         ParResponse parResponse = new ParResponse(requestUri, expiresIn);

services/src/main/java/org/keycloak/protocol/oidc/par/endpoints/request/AuthzEndpointParParser.java

@@ -49,7 +49,7 @@ public class AuthzEndpointParParser extends AuthzEndpointRequestParser {
     public AuthzEndpointParParser(KeycloakSession session, ClientModel client, String requestUri) {
         this.session = session;
         this.client = client;
-        SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseStore = session.singleUseObjects();
         String key;
         try {
             key = requestUri.substring(ParEndpoint.REQUEST_URI_PREFIX_LENGTH);

services/src/main/java/org/keycloak/protocol/oidc/utils/OAuth2CodeParser.java

@@ -52,7 +52,7 @@ public class OAuth2CodeParser {
      * @return code parameter to be used in OAuth2 handshake
      */
     public static String persistCode(KeycloakSession session, AuthenticatedClientSessionModel clientSession, OAuth2Code codeData) {
-        SingleUseObjectProvider codeStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider codeStore = session.singleUseObjects();
 
         String key = codeData.getId();
         if (key == null) {
@@ -112,7 +112,7 @@ public class OAuth2CodeParser {
 
         result.clientSession = userSession.getAuthenticatedClientSessionByClient(clientUUID);
 
-        SingleUseObjectProvider codeStore = session.getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider codeStore = session.singleUseObjects();
         Map<String, String> codeData = codeStore.remove(codeUUID);
 
         // Either code not available or was already used

services/src/main/java/org/keycloak/protocol/saml/SamlProtocol.java

@@ -207,7 +207,7 @@ public class SamlProtocol implements LoginProtocol {
 
     private SingleUseObjectProvider getSingleUseStore() {
         if (singleUseStore == null) {
-            singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            singleUseStore = session.singleUseObjects();
         }
         return singleUseStore;
     }

services/src/main/java/org/keycloak/protocol/saml/SamlService.java

@@ -1126,7 +1126,7 @@ public class SamlService extends AuthorizationEndpointBase {
 
 
     private SingleUseObjectProvider getSingleUseStore() {
-        return session.getProvider(SingleUseObjectProvider.class);
+        return session.singleUseObjects();
     }
 
     /**

services/src/main/java/org/keycloak/services/DefaultKeycloakSession.java

@@ -33,6 +33,7 @@ import org.keycloak.models.KeycloakTransactionManager;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RealmProvider;
 import org.keycloak.models.RoleProvider;
+import org.keycloak.models.SingleUseObjectProvider;
 import org.keycloak.models.ThemeManager;
 import org.keycloak.models.TokenManager;
 import org.keycloak.models.UserCredentialManager;
@@ -79,9 +80,6 @@ public class DefaultKeycloakSession implements KeycloakSession {
     private DatastoreProvider datastoreProvider;
     @Deprecated
     private UserCredentialManager userCredentialStorageManager;
-    private UserSessionProvider sessionProvider;
-    private UserLoginFailureProvider userLoginFailureProvider;
-    private AuthenticationSessionProvider authenticationSessionProvider;
     private final KeycloakContext context;
     private KeyManager keyManager;
     private ThemeManager themeManager;
@@ -390,26 +388,22 @@ public class DefaultKeycloakSession implements KeycloakSession {
 
     @Override
     public UserSessionProvider sessions() {
-        if (sessionProvider == null) {
+        return getDatastoreProvider().userSessions();
-            sessionProvider = getProvider(UserSessionProvider.class);
-        }
-        return sessionProvider;
     }
 
     @Override
     public UserLoginFailureProvider loginFailures() {
-        if (userLoginFailureProvider == null) {
+        return getDatastoreProvider().loginFailures();
-            userLoginFailureProvider = getProvider(UserLoginFailureProvider.class);
-        }
-        return userLoginFailureProvider;
     }
 
     @Override
     public AuthenticationSessionProvider authenticationSessions() {
-        if (authenticationSessionProvider == null) {
+        return getDatastoreProvider().authSessions();
-            authenticationSessionProvider = getProvider(AuthenticationSessionProvider.class);
+    }
-        }
+
-        return authenticationSessionProvider;
+    @Override
+    public SingleUseObjectProvider singleUseObjects() {
+        return getDatastoreProvider().singleUseObjects();
     }
 
     @Override

services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java

@@ -1052,7 +1052,7 @@ public class AuthenticationManager {
         if (actionTokenKeyToInvalidate != null) {
             SingleUseObjectKeyModel actionTokenKey = DefaultActionTokenKey.from(actionTokenKeyToInvalidate);
             if (actionTokenKey != null) {
-                SingleUseObjectProvider singleUseObjectProvider = session.getProvider(SingleUseObjectProvider.class);
+                SingleUseObjectProvider singleUseObjectProvider = session.singleUseObjects();
                 singleUseObjectProvider.put(actionTokenKeyToInvalidate, actionTokenKey.getExpiration() - Time.currentTime(), null); // Token is invalidated
             }
         }

services/src/main/java/org/keycloak/services/resources/LoginActionsServiceChecks.java

@@ -300,7 +300,7 @@ public class LoginActionsServiceChecks {
     }
 
     public static <T extends JsonWebToken & SingleUseObjectKeyModel> void checkTokenWasNotUsedYet(T token, ActionTokenContext<T> context) throws VerificationException {
-        SingleUseObjectProvider singleUseObjectProvider = context.getSession().getProvider(SingleUseObjectProvider.class);
+        SingleUseObjectProvider singleUseObjectProvider = context.getSession().singleUseObjects();
 
         if (singleUseObjectProvider.get(token.serializeKey()) != null) {
             throw new ExplainedTokenVerificationException(token, Errors.EXPIRED_CODE, Messages.EXPIRED_ACTION);

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/util/saml/HandleArtifactStepBuilder.java

@@ -196,8 +196,8 @@ public class HandleArtifactStepBuilder extends SamlDocumentStepBuilder<ArtifactR
 
         if (beforeStepChecker != null && beforeStepChecker instanceof SessionStateChecker) {
             SessionStateChecker sessionStateChecker = (SessionStateChecker) beforeStepChecker;
-            sessionStateChecker.setUserSessionProvider(session -> session.getProvider(SingleUseObjectProvider.class).get(artifact).get(SamlProtocol.USER_SESSION_ID));
+            sessionStateChecker.setUserSessionProvider(session -> session.singleUseObjects().get(artifact).get(SamlProtocol.USER_SESSION_ID));
-            sessionStateChecker.setClientSessionProvider(session -> session.getProvider(SingleUseObjectProvider.class).get(artifact).get(SamlProtocol.CLIENT_SESSION_ID));
+            sessionStateChecker.setClientSessionProvider(session -> session.singleUseObjects().get(artifact).get(SamlProtocol.CLIENT_SESSION_ID));
         }
 
         HttpPost post =  Soap.createMessage().addToBody(DocumentUtil.getDocument(transformed)).buildHttpPost(authServerSamlUrl);

testsuite/model/src/test/java/org/keycloak/testsuite/model/singleUseObject/SingleUseObjectModelTest.java

@@ -69,7 +69,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
     @Test
     public void testActionTokens() {
         DefaultActionTokenKey key = withRealm(realmId, (session, realm) -> {
-            SingleUseObjectProvider singleUseObjectProvider = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseObjectProvider = session.singleUseObjects();
             int time = Time.currentTime();
             DefaultActionTokenKey actionTokenKey = new DefaultActionTokenKey(userId, UUID.randomUUID().toString(), time + 60, null);
             Map<String, String> notes = new HashMap<>();
@@ -79,7 +79,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         });
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseObjectProvider = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseObjectProvider = session.singleUseObjects();
             Map<String, String> notes = singleUseObjectProvider.get(key.serializeKey());
             Assert.assertNotNull(notes);
             Assert.assertEquals("bar", notes.get("foo"));
@@ -90,7 +90,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         });
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseObjectProvider = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseObjectProvider = session.singleUseObjects();
             Map<String, String> notes = singleUseObjectProvider.get(key.serializeKey());
             Assert.assertNull(notes);
 
@@ -100,7 +100,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         });
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseObjectProvider = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseObjectProvider = session.singleUseObjects();
             Map<String, String> notes = singleUseObjectProvider.get(key.serializeKey());
             Assert.assertNotNull(notes);
             Assert.assertEquals("bar", notes.get("foo"));
@@ -109,7 +109,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         setTimeOffset(70);
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseObjectProvider = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseObjectProvider = session.singleUseObjects();
             Map<String, String> notes = singleUseObjectProvider.get(key.serializeKey());
             notes = singleUseObjectProvider.get(key.serializeKey());
             Assert.assertNull(notes);
@@ -126,14 +126,14 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         notes2.put("baf", "meow");
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseStore = session.singleUseObjects();
             Assert.assertFalse(singleUseStore.replace(key, notes2));
 
             singleUseStore.put(key,  60, notes);
         });
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseStore = session.singleUseObjects();
             Map<String, String> actualNotes = singleUseStore.get(key);
             Assert.assertEquals(notes, actualNotes);
 
@@ -141,7 +141,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         });
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseStore = session.singleUseObjects();
             Map<String, String> actualNotes = singleUseStore.get(key);
             Assert.assertEquals(notes2, actualNotes);
 
@@ -151,12 +151,12 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         });
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseStore = session.singleUseObjects();
             Assert.assertTrue(singleUseStore.putIfAbsent(key, 60));
         });
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseStore = session.singleUseObjects();
             Map<String, String> actualNotes = singleUseStore.get(key);
             assertThat(actualNotes, Matchers.anEmptyMap());
         });
@@ -164,7 +164,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
         setTimeOffset(70);
 
         inComittedTransaction(session -> {
-            SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+            SingleUseObjectProvider singleUseStore = session.singleUseObjects();
             Assert.assertNull(singleUseStore.get(key));
         });
     }
@@ -196,7 +196,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
 
             if (index.incrementAndGet() == 1) {
                 actionTokenKey.set(withRealm(realmId, (session, realm) -> {
-                    SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+                    SingleUseObjectProvider singleUseStore = session.singleUseObjects();
                     singleUseStore.put(key, 60, notes);
 
                     int time = Time.currentTime();
@@ -212,7 +212,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
 
             // check if single-use object/action token is available on all nodes
             inComittedTransaction(session -> {
-                SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+                SingleUseObjectProvider singleUseStore = session.singleUseObjects();
                 while (singleUseStore.get(key) == null || singleUseStore.get(actionTokenKey.get()) == null) {
                     sleep(1000);
                 }
@@ -224,7 +224,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
             // remove objects on one node
             if (index.incrementAndGet() == 5) {
                 inComittedTransaction(session -> {
-                    SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+                    SingleUseObjectProvider singleUseStore = session.singleUseObjects();
                     singleUseStore.remove(key);
                     singleUseStore.remove(actionTokenKey.get());
                 });
@@ -236,7 +236,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest {
 
             // check if single-use object/action token is removed
             inComittedTransaction(session -> {
-                SingleUseObjectProvider singleUseStore = session.getProvider(SingleUseObjectProvider.class);
+                SingleUseObjectProvider singleUseStore = session.singleUseObjects();
 
                 while (singleUseStore.get(key) != null && singleUseStore.get(actionTokenKey.get()) != null) {
                    sleep(1000);