Keycloak UserStorageProviderFactory with configurable JPA Peristence Context

custom-ds.xml

<?xml version="1.0" encoding="UTF-8"?>
<datasources xmlns="http://www.jboss.org/ironjacamar/schema"
             xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
             xsi:schemaLocation="
             http://www.jboss.org/ironjacamar/schema
             http://docs.jboss.org/ironjacamar/schema/datasources_1_0.xsd">
  <xa-datasource use-java-context="true" enabled="true" jndi-name="java:jboss/datasources/CustomDS" pool-name="CustomDS">
    <xa-datasource-property name="URL">
      jdbc:h2:${jboss.server.data.dir}/custom;AUTO_SERVER=TRUE;MULTI_THREADED=TRUE
    </xa-datasource-property>
    <driver>h2</driver>
    <security>
      <user-name>sa</user-name>
      <password>sa</password>
    </security>
  </xa-datasource>
</datasources>

CustomUserStorageProviderFactory.java

package com.github.bertramn.keycloak;

import org.keycloak.component.ComponentModel;
import org.keycloak.component.ComponentValidationException;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.RealmModel;
import org.keycloak.provider.ProviderConfigProperty;
import org.keycloak.provider.ProviderConfigurationBuilder;
import org.keycloak.storage.UserStorageProviderFactory;

import javax.naming.InitialContext;
import javax.persistence.EntityManagerFactory;
import javax.persistence.Persistence;
import java.util.Collections;
import java.util.List;
import java.util.Properties;

public class CustomUserStorageProviderFactory implements UserStorageProviderFactory<CustomUserStorageProvider> {

  /**
   * the name of the persistence unit
   */
  private static final String PERSISTENCE_UNIT_NAME = "custom-default";

  /**
   * The provider id
   */
  private static final String PROVIDER_NAME = "custom";

  /**
   * the provider configuration metadata definition
   */
  private static final List<ProviderConfigProperty> configMetadata;

  static {
    configMetadata = ProviderConfigurationBuilder.create()
      .property().name("datasource")
      .type(ProviderConfigProperty.STRING_TYPE)
      .label("Data Source Name")
      .defaultValue("java:jboss/datasources/CustomDS")
      .helpText("The jndi name of the XA datasource used by this federation provider.")
      .add()
      .build();
  }

  /**
   * The JPA entity manager factory used by the provider.
   */
  private EntityManagerFactory managedEntityManagerFactory;

  /**
   * The SPI configuration parameter used to enable debug on entity manager and other internals.
   */
  private Boolean debug = false;

  @Override
  public String getId() {
    return PROVIDER_NAME;
  }

  @Override
  public List<ProviderConfigProperty> getConfigProperties() {
    return configMetadata;
  }

  @Override
  public void validateConfiguration(KeycloakSession session, RealmModel realm, ComponentModel config) throws ComponentValidationException {

    String datasource = config.getConfig().getFirst("datasource");

    if (datasource == null) {
      throw new ComponentValidationException("datasource name is not configured");
    }

    // on configuration update we need to close the entity manager factory as we might
    // have to connect to a different data source
    closeEntityManagerFactory();

  }

  @Override
  public CustomUserStorageProvider create(KeycloakSession session, ComponentModel model) {

    String datasource = model.getConfig().getFirst("datasource");

    debug = Boolean.parseBoolean(model.getConfig().getFirst("debug"));

    try {
      InitialContext ctx = new InitialContext();

      CustomUserStorageProvider provider = (CustomUserStorageProvider) ctx.lookup("java:global/custom-federation/custom-federation-ejb/" + CustomUserStorageProvider.class.getSimpleName());

      provider.setModel(model);
      provider.setSession(session);
      provider.setEntityManagerFactory(getEntityManagerFactory(datasource));
      return provider;

    } catch (Exception e) {
      throw new RuntimeException(e);
    }

  }

  /**
   * Create a JTA enabled and enrolled {@link EntityManagerFactory} that can be injected into the {@link CustomUserStorageProvider}.
   *
   * @param datasourceName the JNDI name of the XA datasource
   *
   * @return a configured entity manager factory
   */
  private EntityManagerFactory getEntityManagerFactory(String datasourceName) {

    if (managedEntityManagerFactory == null) {

      Properties p = new Properties();
      p.put("hibernate.connection.datasource", datasourceName);
      p.put("hibernate.transaction.jta.platform", "org.hibernate.engine.transaction.jta.platform.internal.JBossAppServerJtaPlatform");
      p.put("current_session_context_class", "jta");
      p.put("hibernate.ddl-auto", "create");
      p.put("hibernate.show_sql", debug.toString());
      p.put("hibernate.format_sql", debug.toString());

      // Adding "hibernate.classLoaders" property is critical for this to work with keycloak!!!
      p.put("hibernate.classLoaders", Collections.singletonList(getClass().getClassLoader()));

      managedEntityManagerFactory = Persistence.createEntityManagerFactory(PERSISTENCE_UNIT_NAME, p);

    }

    return managedEntityManagerFactory;

  }

  @Override
  public void close() {
    closeEntityManagerFactory();
  }

  /**
   * Closes the application managed entity manager factory and frees handle for GC.
   */
  private void closeEntityManagerFactory() {
    if (managedEntityManagerFactory != null) {
      try {
        managedEntityManagerFactory.close();
      } catch (Exception e) {
        e.printStackTrace();
      } finally {
        managedEntityManagerFactory = null;
      }

    }
  }

}

persistence.xml

<?xml version="1.0" encoding="UTF-8"?>
<persistence version="2.0"
             xmlns="http://java.sun.com/xml/ns/persistence"
             xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
             xsi:schemaLocation="http://java.sun.com/xml/ns/persistence http://java.sun.com/xml/ns/persistence/persistence_2_0.xsd">
  <persistence-unit name="custom-default" transaction-type="JTA">
    <class>com.github.bertramn.keycloak.DatabaseUser</class>
  </persistence-unit>
</persistence>