HHH-15422 Pick up ConnectionProvider from BeanContainer if not explicit set

hibernate-core/src/main/java/org/hibernate/engine/jdbc/connections/internal/ConnectionProviderInitiator.java

@@ -4,6 +4,7 @@
  */
 package org.hibernate.engine.jdbc.connections.internal;
 
+import java.lang.reflect.InvocationTargetException;
 import java.sql.Connection;
 import java.util.Collection;
 import java.util.HashSet;
@@ -19,6 +20,11 @@ import org.hibernate.engine.jdbc.connections.spi.ConnectionProvider;
 import org.hibernate.internal.CoreLogging;
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.internal.util.StringHelper;
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.internal.FallbackBeanInstanceProducer;
+import org.hibernate.resource.beans.internal.Helper;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
+import org.hibernate.resource.beans.spi.ManagedBeanRegistry;
 import org.hibernate.service.spi.ServiceRegistryImplementor;
 
 import static java.sql.Connection.TRANSACTION_NONE;
@@ -102,6 +108,7 @@ public class ConnectionProviderInitiator implements StandardServiceInitiator<Con
 			return null;
 		}
 
+		final BeanContainer beanContainer = Helper.allowExtensionsInCdi( registry ) ? registry.requireService( ManagedBeanRegistry.class ).getBeanContainer() : null;
 		final StrategySelector strategySelector = registry.requireService( StrategySelector.class );
 		final Object explicitSetting = configurationValues.get( CONNECTION_PROVIDER );
 		if ( explicitSetting != null ) {
@@ -111,25 +118,25 @@ public class ConnectionProviderInitiator implements StandardServiceInitiator<Con
 			}
 			else if ( explicitSetting instanceof Class<?> providerClass ) {
 				LOG.instantiatingExplicitConnectionProvider( providerClass.getName() );
-				return instantiateExplicitConnectionProvider( providerClass );
+				return instantiateExplicitConnectionProvider( providerClass, beanContainer );
 			}
 			else {
 				final String providerName = nullIfEmpty( explicitSetting.toString() );
 				if ( providerName != null ) {
-					return instantiateNamedConnectionProvider(providerName, strategySelector);
+					return instantiateNamedConnectionProvider(providerName, strategySelector, beanContainer);
 				}
 			}
 		}
 
-		return instantiateConnectionProvider( configurationValues, strategySelector );
+		return instantiateConnectionProvider( configurationValues, strategySelector, beanContainer );
 	}
 
-	private ConnectionProvider instantiateNamedConnectionProvider(String providerName, StrategySelector strategySelector) {
+	private ConnectionProvider instantiateNamedConnectionProvider(String providerName, StrategySelector strategySelector, BeanContainer beanContainer) {
 		LOG.instantiatingExplicitConnectionProvider( providerName );
 		final Class<?> providerClass =
 				strategySelector.selectStrategyImplementor( ConnectionProvider.class, providerName );
 		try {
-			return instantiateExplicitConnectionProvider( providerClass );
+			return instantiateExplicitConnectionProvider( providerClass, beanContainer );
 		}
 		catch (Exception e) {
 			throw new HibernateException(
@@ -140,7 +147,7 @@ public class ConnectionProviderInitiator implements StandardServiceInitiator<Con
 	}
 
 	private ConnectionProvider instantiateConnectionProvider(
-			Map<String, Object> configurationValues, StrategySelector strategySelector) {
+			Map<String, Object> configurationValues, StrategySelector strategySelector, BeanContainer beanContainer) {
 		if ( configurationValues.containsKey( DATASOURCE ) ) {
 			return new DatasourceConnectionProviderImpl();
 		}
@@ -149,9 +156,9 @@ public class ConnectionProviderInitiator implements StandardServiceInitiator<Con
 				getSingleRegisteredProvider( strategySelector );
 		if ( singleRegisteredProvider != null ) {
 			try {
-				return singleRegisteredProvider.newInstance();
+				return singleRegisteredProvider.getConstructor().newInstance();
 			}
-			catch (IllegalAccessException | InstantiationException e) {
+			catch (IllegalAccessException | NoSuchMethodException | InvocationTargetException | InstantiationException e) {
 				throw new HibernateException( "Could not instantiate singular-registered ConnectionProvider", e );
 			}
 		}
@@ -177,11 +184,47 @@ public class ConnectionProviderInitiator implements StandardServiceInitiator<Con
 			return new DriverManagerConnectionProviderImpl();
 		}
 		else {
-			LOG.noAppropriateConnectionProvider();
+			if (beanContainer != null) {
-			return new UserSuppliedConnectionProviderImpl();
+				return beanContainer.getBean(
+						ConnectionProvider.class,
+						new BeanContainer.LifecycleOptions() {
+							@Override
+							public boolean canUseCachedReferences() {
+								return true;
+							}
+
+							@Override
+							public boolean useJpaCompliantCreation() {
+								return true;
+							}
+						},
+						new BeanInstanceProducer() {
+
+							@Override
+							public <B> B produceBeanInstance(Class<B> beanType) {
+								return (B) noAppropriateConnectionProvider();
+							}
+
+							@Override
+							public <B> B produceBeanInstance(String name, Class<B> beanType) {
+								return (B) noAppropriateConnectionProvider();
+							}
+
+						}
+				).getBeanInstance();
+			}
+			else {
+				return noAppropriateConnectionProvider();
+			}
+
 		}
 	}
 
+	private ConnectionProvider noAppropriateConnectionProvider() {
+		LOG.noAppropriateConnectionProvider();
+		return new UserSuppliedConnectionProviderImpl();
+	}
+
 	private Class<? extends ConnectionProvider> getSingleRegisteredProvider(StrategySelector strategySelector) {
 		final Collection<Class<? extends ConnectionProvider>> implementors =
 				strategySelector.getRegisteredStrategyImplementors( ConnectionProvider.class );
@@ -190,9 +233,28 @@ public class ConnectionProviderInitiator implements StandardServiceInitiator<Con
 				: null;
 	}
 
-	private ConnectionProvider instantiateExplicitConnectionProvider(Class<?> providerClass) {
+	private ConnectionProvider instantiateExplicitConnectionProvider(Class<?> providerClass, BeanContainer beanContainer) {
 		try {
-			return (ConnectionProvider) providerClass.newInstance();
+			if ( beanContainer != null ) {
+				return (ConnectionProvider) beanContainer.getBean(
+						providerClass,
+						new BeanContainer.LifecycleOptions() {
+							@Override
+							public boolean canUseCachedReferences() {
+								return true;
+							}
+
+							@Override
+							public boolean useJpaCompliantCreation() {
+								return true;
+							}
+						},
+						FallbackBeanInstanceProducer.INSTANCE
+				).getBeanInstance();
+			}
+			else {
+				return (ConnectionProvider) providerClass.getConstructor().newInstance();
+			}
 		}
 		catch (Exception e) {
 			throw new HibernateException( "Could not instantiate connection provider [" + providerClass.getName() + "]", e );
@@ -201,7 +263,7 @@ public class ConnectionProviderInitiator implements StandardServiceInitiator<Con
 
 	private static ConnectionProvider instantiateProvider(StrategySelector selector, String strategy) {
 		try {
-			return selector.selectStrategyImplementor( ConnectionProvider.class, strategy ).newInstance();
+			return selector.selectStrategyImplementor( ConnectionProvider.class, strategy ).getConstructor().newInstance();
 		}
 		catch ( Exception e ) {
 			LOG.providerClassNotFound(strategy);

hibernate-core/src/test/java/org/hibernate/orm/test/connections/ConnectionProviderFromBeanContainerTest.java

@@ -0,0 +1,104 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.connections;
+
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.dialect.H2Dialect;
+import org.hibernate.engine.jdbc.connections.spi.ConnectionProvider;
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.container.spi.ContainedBean;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
+import org.hibernate.service.ServiceRegistry;
+import org.hibernate.testing.junit4.BaseUnitTestCase;
+import org.hibernate.testing.orm.junit.RequiresDialect;
+import org.hibernate.testing.util.ServiceRegistryUtil;
+import org.junit.Test;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.hibernate.cfg.JdbcSettings.CONNECTION_PROVIDER;
+import static org.junit.Assert.assertSame;
+
+/**
+ * @author Yanming Zhou
+ */
+@RequiresDialect(H2Dialect.class)
+public class ConnectionProviderFromBeanContainerTest extends BaseUnitTestCase {
+
+	private final ConnectionProvider dummyConnectionProvider = new DummyConnectionProvider();
+
+	private Map<String, Object> createSettings() {
+		Map<String, Object> settings = new HashMap<>();
+		settings.put( AvailableSettings.ALLOW_EXTENSIONS_IN_CDI, "true" );
+		settings.put( AvailableSettings.BEAN_CONTAINER, new BeanContainer() {
+			@Override
+			@SuppressWarnings("unchecked")
+			public <B> ContainedBean<B> getBean(
+					Class<B> beanType,
+					LifecycleOptions lifecycleOptions,
+					BeanInstanceProducer fallbackProducer) {
+				return () -> (B) ( beanType == DummyConnectionProvider.class ?
+						dummyConnectionProvider : fallbackProducer.produceBeanInstance( beanType ) );
+			}
+
+			@Override
+			public <B> ContainedBean<B> getBean(
+					String name,
+					Class<B> beanType,
+					LifecycleOptions lifecycleOptions,
+					BeanInstanceProducer fallbackProducer) {
+				return () -> (B) fallbackProducer.produceBeanInstance( beanType );
+			}
+
+			@Override
+			public void stop() {
+
+			}
+		} );
+		return settings;
+	}
+
+	@Test
+	public void testProviderFromBeanContainerInUse() {
+		Map<String, Object> settings = createSettings();
+		settings.putIfAbsent( CONNECTION_PROVIDER, DummyConnectionProvider.class.getName() );
+		try ( ServiceRegistry serviceRegistry = ServiceRegistryUtil.serviceRegistryBuilder()
+				.applySettings( settings ).build() ) {
+			ConnectionProvider providerInUse = serviceRegistry.getService( ConnectionProvider.class );
+			assertSame( dummyConnectionProvider, providerInUse );
+		}
+	}
+
+	public static class DummyConnectionProvider implements ConnectionProvider {
+
+		@Override
+		public boolean isUnwrappableAs(Class<?> unwrapType) {
+			return false;
+		}
+
+		@Override
+		public <T> T unwrap(Class<T> unwrapType) {
+			return null;
+		}
+
+		@Override
+		public Connection getConnection() throws SQLException {
+			return null;
+		}
+
+		@Override
+		public void closeConnection(Connection connection) throws SQLException {
+
+		}
+
+		@Override
+		public boolean supportsAggressiveRelease() {
+			return false;
+		}
+	};
+}