use a custom user-model-transaction

flintstones-userprovider/src/main/java/dasniko/keycloak/user/flintstones/FlintstonesUserStorageProvider.java

@@ -3,7 +3,6 @@
 import dasniko.keycloak.user.flintstones.repo.Credential;
 import dasniko.keycloak.user.flintstones.repo.FlintstoneUser;
 import dasniko.keycloak.user.flintstones.repo.FlintstonesApiClient;
-import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.keycloak.component.ComponentModel;
 import org.keycloak.credential.CredentialInput;
@@ -25,18 +24,12 @@
 import org.keycloak.storage.user.UserQueryProvider;
 import org.keycloak.storage.user.UserRegistrationProvider;
 
-import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
-/**
- * @author Niko Köbler, http://www.n-k.de, @dasniko
- */
 @Slf4j
-@RequiredArgsConstructor
 public class FlintstonesUserStorageProvider implements UserStorageProvider,
 	UserLookupProvider, UserQueryProvider, CredentialInputUpdater, CredentialInputValidator,
 	UserRegistrationProvider {
@@ -45,10 +38,15 @@ public class FlintstonesUserStorageProvider implements UserStorageProvider,
 	private final ComponentModel model;
 	private final FlintstonesApiClient apiClient;
 
-	// map of loaded users in this transaction
-	private final Map<String, UserModel> loadedUsers = new HashMap<>();
-	// list of newly created users in this transaction
-	private final List<FlintstoneUserAdapter> newUsers = new ArrayList<>();
+	// user handling in this transaction
+	UserModelTransaction tx = new UserModelTransaction(this::updateUser);
+
+	public FlintstonesUserStorageProvider(KeycloakSession session, ComponentModel model, FlintstonesApiClient apiClient) {
+		this.session = session;
+		this.model = model;
+		this.apiClient = apiClient;
+		session.getTransactionManager().enlistAfterCompletion(tx);
+	}
 
 	@Override
 	public boolean supportsCredentialType(String credentialType) {
@@ -120,13 +118,13 @@ public UserModel getUserByEmail(RealmModel realm, String email) {
 	}
 
 	private UserModel findUser(RealmModel realm, String identifier, Function<String, FlintstoneUser> fnFindUser) {
-		UserModel adapter = loadedUsers.get(identifier);
+		UserModel adapter = tx.findUser(identifier);
 		if (adapter == null) {
 			FlintstoneUser user = fnFindUser.apply(identifier);
 			log.debug("Received user data for identifier <{}> from repository: {}", identifier, user);
 			if (user != null) {
 				adapter = new FlintstoneUserAdapter(session, realm, model, user);
-				loadedUsers.put(identifier, adapter);
+				tx.addUser(identifier, adapter);
 			}
 		} else {
 			log.debug("Found user data for {} in loadedUsers.", identifier);
@@ -168,8 +166,7 @@ public UserModel addUser(RealmModel realm, String username) {
 			apiClient.createUser(flintstoneUser);
 			flintstoneUser = apiClient.getUserByUsername(username);
 			FlintstoneUserAdapter newUser = new FlintstoneUserAdapter(session, realm, model, flintstoneUser);
-			newUsers.add(newUser);
-			loadedUsers.put(username, newUser);
+			tx.addUser(username, newUser);
 			return newUser;
 		}
 		return null;
@@ -183,16 +180,6 @@ public boolean removeUser(RealmModel realm, UserModel user) {
 
 	@Override
 	public void close() {
-		newUsers.forEach(newUser -> {
-			apiClient.updateUser(newUser.getUser());
-			loadedUsers.remove(newUser.getUsername());
-		});
-		loadedUsers.values().forEach(user -> {
-			FlintstoneUserAdapter userAdapter = (FlintstoneUserAdapter) user;
-			if (userAdapter.isDirty()) {
-				apiClient.updateUser(userAdapter.getUser());
-			}
-		});
 	}
 
 	private boolean syncUsers() {
@@ -203,4 +190,11 @@ private boolean usePasswordPolicy() {
 		return model.get(FlintstonesUserStorageProviderFactory.USE_PASSWORD_POLICY, false);
 	}
 
+	private void updateUser(UserModel user) {
+		FlintstoneUserAdapter userAdapter = (FlintstoneUserAdapter) user;
+		if (userAdapter.isDirty()) {
+			apiClient.updateUser(userAdapter.getUser());
+		}
+	}
+
 }

flintstones-userprovider/src/main/java/dasniko/keycloak/user/flintstones/UserModelTransaction.java

@@ -0,0 +1,37 @@
+package dasniko.keycloak.user.flintstones;
+
+import org.keycloak.models.AbstractKeycloakTransaction;
+import org.keycloak.models.UserModel;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Consumer;
+
+public class UserModelTransaction extends AbstractKeycloakTransaction {
+
+	private final Consumer<UserModel> userConsumer;
+	private final Map<String, UserModel> loadedUsers = new HashMap<>();
+
+	public UserModelTransaction(Consumer<UserModel> userConsumer) {
+		this.userConsumer = userConsumer;
+	}
+
+	public void addUser(String identifier, UserModel userModel) {
+		loadedUsers.put(identifier, userModel);
+	}
+
+	public UserModel findUser(String identifier) {
+		return loadedUsers.get(identifier);
+	}
+
+	@Override
+	protected void commitImpl() {
+		loadedUsers.values().forEach(userConsumer);
+	}
+
+	@Override
+	protected void rollbackImpl() {
+		// maybe do some more checks on loadedUsers...
+		loadedUsers.clear();
+	}
+}