diff --git a/extension/persistence/eclipselink/src/main/java/io/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java b/extension/persistence/eclipselink/src/main/java/io/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java index a49b6071fb..c57509ee7b 100644 --- a/extension/persistence/eclipselink/src/main/java/io/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java +++ b/extension/persistence/eclipselink/src/main/java/io/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java @@ -54,17 +54,21 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; import java.util.stream.Collectors; import javax.xml.parsers.DocumentBuilder; import javax.xml.parsers.DocumentBuilderFactory; +import javax.xml.parsers.ParserConfigurationException; import javax.xml.xpath.XPath; import javax.xml.xpath.XPathConstants; +import javax.xml.xpath.XPathExpressionException; import javax.xml.xpath.XPathFactory; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; +import org.jetbrains.annotations.TestOnly; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.w3c.dom.Document; @@ -80,6 +84,10 @@ public class PolarisEclipseLinkMetaStoreSessionImpl implements PolarisMetaStoreS private static final Logger LOG = LoggerFactory.getLogger(PolarisEclipseLinkMetaStoreSessionImpl.class); + // Cache to hold the EntityManagerFactory for each realm. Each realm needs a separate + // EntityManagerFactory since it connects to different databases + private static final ConcurrentHashMap realmFactories = + new ConcurrentHashMap<>(); private final EntityManagerFactory emf; private final ThreadLocal localSession = new ThreadLocal<>(); private final PolarisEclipseLinkStore store; @@ -100,20 +108,30 @@ public PolarisEclipseLinkMetaStoreSessionImpl( @NotNull RealmContext realmContext, @Nullable String confFile, @Nullable String persistenceUnitName) { - - emf = createEntityManagerFactory(realmContext, confFile, persistenceUnitName); LOG.debug("Create EclipseLink Meta Store Session for {}", realmContext.getRealmIdentifier()); + emf = createEntityManagerFactory(realmContext, confFile, persistenceUnitName); // init store this.store = store; this.storageIntegrationProvider = storageIntegrationProvider; } - /** Load the persistence unit properties from a given configuration file */ + /** + * Create EntityManagerFactory. + * + *

The EntityManagerFactory creation is expensive, so we are caching and reusing it for each + * realm. + */ private EntityManagerFactory createEntityManagerFactory( @NotNull RealmContext realmContext, @Nullable String confFile, @Nullable String persistenceUnitName) { + String realm = realmContext.getRealmIdentifier(); + EntityManagerFactory factory = realmFactories.getOrDefault(realm, null); + if (factory != null) { + return factory; + } + ClassLoader prevClassLoader = Thread.currentThread().getContextClassLoader(); try { persistenceUnitName = persistenceUnitName == null ? "polaris" : persistenceUnitName; @@ -131,31 +149,44 @@ private EntityManagerFactory createEntityManagerFactory( if (prefixUrl == null) { prefixUrl = new File(jarPrefixPath).toURI().toURL(); } + + LOG.info( + "Created a new ClassLoader with the jar {} in classpath to load the config file", + prefixUrl); + URLClassLoader currentClassLoader = new URLClassLoader(new URL[] {prefixUrl}, this.getClass().getClassLoader()); + + LOG.debug("Update ClassLoader in current thread temporarily"); Thread.currentThread().setContextClassLoader(currentClassLoader); } Map properties = loadProperties(confFile, persistenceUnitName); // Replace database name in JDBC URL with realm if (properties.containsKey(JDBC_URL)) { - properties.put( - JDBC_URL, - properties.get(JDBC_URL).replace("{realm}", realmContext.getRealmIdentifier())); + properties.put(JDBC_URL, properties.get(JDBC_URL).replace("{realm}", realm)); } properties.put(ECLIPSELINK_PERSISTENCE_XML, confFile); - return Persistence.createEntityManagerFactory(persistenceUnitName, properties); - } catch (Exception e) { + factory = Persistence.createEntityManagerFactory(persistenceUnitName, properties); + realmFactories.putIfAbsent(realm, factory); + + return factory; + } catch (IOException e) { throw new RuntimeException(e); } finally { Thread.currentThread().setContextClassLoader(prevClassLoader); } } + @TestOnly + static void clearEntityManagerFactories() { + realmFactories.clear(); + } + /** Load the persistence unit properties from a given configuration file */ private Map loadProperties( - @NotNull String confFile, @NotNull String persistenceUnitName) throws Exception { + @NotNull String confFile, @NotNull String persistenceUnitName) throws IOException { try { InputStream input = Thread.currentThread().getContextClassLoader().getResourceAsStream(confFile); @@ -176,13 +207,16 @@ private Map loadProperties( } return properties; - } catch (SAXException | IOException e) { + } catch (XPathExpressionException + | ParserConfigurationException + | SAXException + | IOException e) { String str = String.format( "Cannot find or parse the configuration file %s for persistence-unit %s", confFile, persistenceUnitName); - LOG.error(str); - throw new Exception(str); + LOG.error(str, e); + throw new IOException(str); } } diff --git a/extension/persistence/eclipselink/src/test/java/com/snowflake/polaris/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreTest.java b/extension/persistence/eclipselink/src/test/java/io/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreTest.java similarity index 93% rename from extension/persistence/eclipselink/src/test/java/com/snowflake/polaris/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreTest.java rename to extension/persistence/eclipselink/src/test/java/io/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreTest.java index 7d707fc6a1..32b4169d60 100644 --- a/extension/persistence/eclipselink/src/test/java/com/snowflake/polaris/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreTest.java +++ b/extension/persistence/eclipselink/src/test/java/io/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreTest.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.snowflake.polaris.persistence.impl.eclipselink; +package io.polaris.extension.persistence.impl.eclipselink; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -26,8 +26,6 @@ import io.polaris.core.persistence.PolarisMetaStoreManagerImpl; import io.polaris.core.persistence.PolarisMetaStoreManagerTest; import io.polaris.core.persistence.PolarisTestMetaStoreManager; -import io.polaris.extension.persistence.impl.eclipselink.PolarisEclipseLinkMetaStoreSessionImpl; -import io.polaris.extension.persistence.impl.eclipselink.PolarisEclipseLinkStore; import java.time.ZoneId; import java.util.stream.Stream; import org.junit.jupiter.api.extension.ExtensionContext; @@ -63,6 +61,9 @@ protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() { @ParameterizedTest() @ArgumentsSource(CreateStoreSessionArgs.class) void testCreateStoreSession(String confFile, boolean success) { + // Clear cache to prevent reuse EntityManagerFactory + PolarisEclipseLinkMetaStoreSessionImpl.clearEntityManagerFactories(); + PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl(); PolarisEclipseLinkStore store = new PolarisEclipseLinkStore(diagServices); try {