LicenceServerUserStorageProvider.java

import org.jboss.logging.Logger;
import org.keycloak.component.ComponentModel;
import org.keycloak.models.*;
import org.keycloak.storage.StorageId;
import org.keycloak.storage.UserStorageProvider;
import org.keycloak.storage.user.UserLookupProvider;
import org.keycloak.storage.user.UserQueryProvider;

import javax.ejb.Local;
import javax.ejb.Remove;
import javax.ejb.Stateful;
import javax.persistence.EntityManager;
import javax.persistence.PersistenceContext;
import javax.persistence.TypedQuery;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Stateful
@Local(LicenceServerUserStorageProvider.class)
public class LicenceServerUserStorageProvider implements
        UserStorageProvider,
        UserLookupProvider,
        UserQueryProvider {

    protected KeycloakSession keycloakSession;
    protected ComponentModel model;

    private static final Logger LOGGER = Logger.getLogger(LicenceServerUserStorageProvider.class);
    @PersistenceContext
    protected EntityManager em;

    protected Map<String, UserModel> loadedUsers = new HashMap<>();

    public LicenceServerUserStorageProvider(KeycloakSession session, ComponentModel model) {
        this.keycloakSession = session;
        this.model = model;
    }

    public LicenceServerUserStorageProvider() {
    }

    public void setModel(ComponentModel model) {
        this.model = model;
    }

    public void setKeycloakSession(KeycloakSession keycloakSession) {
        this.keycloakSession = keycloakSession;
    }

    @Override
    @Remove
    public void close() {

    }

    @Override
    public UserModel getUserById(String s, RealmModel realmModel) {
        String persistenceId = StorageId.externalId(s);
        LicenceDevice device = em.find(LicenceDevice.class, Long.valueOf(persistenceId));
        if (device == null) {
            return null;
        }
        return new UserAdapter(keycloakSession, realmModel, model, device);
    }

    @Override
    public UserModel getUserByUsername(String s, RealmModel realmModel) {
        TypedQuery<LicenceDevice> query = em.createNamedQuery("findByUid", LicenceDevice.class);
        query.setParameter("uid", s);
        List<LicenceDevice> result = query.getResultList();
        if (result.isEmpty()) {
            return null;
        }
        return new UserAdapter(keycloakSession, realmModel, model, result.get(0));
    }

    @Override
    public UserModel getUserByEmail(String s, RealmModel realmModel) {
        return null;
    }

    @Override
    public int getUsersCount(RealmModel realmModel) {
        Object count = em.createNamedQuery("getDeviceCount").getSingleResult();
        return ((Number) count).intValue();
    }

    @Override
    public List<UserModel> getUsers(RealmModel realmModel) {
        return getUsers(realmModel, -1, -1);
    }

    @Override
    public List<UserModel> getUsers(RealmModel realmModel, int firstResult, int maxResults) {
        TypedQuery<LicenceDevice> query = em.createNamedQuery("getAllDevices", LicenceDevice.class);

        return getUserModels(realmModel, firstResult, maxResults, query);
    }

    @Override
    public List<UserModel> searchForUser(String s, RealmModel realmModel) {
        return searchForUser(s, realmModel, -1, -1);
    }

    @Override
    public List<UserModel> searchForUser(String s, RealmModel realmModel, int firstResult, int maxResults) {
        TypedQuery<LicenceDevice> query = em.createNamedQuery("searchDevices", LicenceDevice.class);
        query.setParameter("search", "%"+s+"%");

        return getUserModels(realmModel, firstResult, maxResults, query);
    }

    private List<UserModel> getUserModels(RealmModel realmModel, int firstResult, int maxResults, TypedQuery<LicenceDevice> query) {
        if (firstResult != -1) {
            query.setFirstResult(firstResult);
        }

        if (maxResults != -1) {
            query.setMaxResults(maxResults);
        }

        List<LicenceDevice> devices = query.getResultList();
        LOGGER.info(devices.toString());
        return devices.stream()
                .map(licenceDevice -> new UserAdapter(keycloakSession, realmModel, model, licenceDevice))
                .collect(Collectors.toList());
    }

    @Override
    public List<UserModel> searchForUser(Map<String, String> map, RealmModel realmModel) {
        return Collections.emptyList();
    }

    @Override
    public List<UserModel> searchForUser(Map<String, String> map, RealmModel realmModel, int i, int i1) {
        return Collections.emptyList();
    }

    @Override
    public List<UserModel> getGroupMembers(RealmModel realmModel, GroupModel groupModel, int i, int i1) {
        return Collections.emptyList();
    }

    @Override
    public List<UserModel> getGroupMembers(RealmModel realmModel, GroupModel groupModel) {
        return Collections.emptyList();
    }

    @Override
    public List<UserModel> searchForUserByUserAttribute(String s, String s1, RealmModel realmModel) {
        return Collections.emptyList();
    }
}