[Spring] Make ServiceBus JMS connection factory creation customizable

sdk/spring/CHANGELOG.md

@@ -8,6 +8,7 @@ This section includes changes in `spring-cloud-azure-autoconfigure` module.
 
 #### Features Added
 
+- Added `AzureServiceBusJmsConnectionFactoryFactory` to allow applications to customize how `ServiceBusJmsConnectionFactory` instances are created, including support for custom subclasses ([#49676](https://github.com/Azure/azure-sdk-for-java/pull/49676)).
 - Added support for constructing `AadB2cAuthorizationRequestResolver` with a custom `authorizationRequestBaseUri`, aligning Azure AD B2C authorization request resolution with the configurability already available for AAD. ([#49674](https://github.com/Azure/azure-sdk-for-java/pull/49674))
 
 #### Breaking Changes

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsAutoConfiguration.java

@@ -3,10 +3,15 @@
 
 package com.azure.spring.cloud.autoconfigure.implementation.jms;
 
+import com.azure.core.credential.TokenCredential;
+import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
+import com.azure.identity.extensions.implementation.credential.provider.TokenCredentialProvider;
 import com.azure.servicebus.jms.ServiceBusJmsConnectionFactory;
+import com.azure.servicebus.jms.ServiceBusJmsConnectionFactorySettings;
 import com.azure.spring.cloud.autoconfigure.implementation.context.properties.AzureGlobalProperties;
 import com.azure.spring.cloud.autoconfigure.implementation.jms.properties.AzureServiceBusJmsProperties;
 import com.azure.spring.cloud.autoconfigure.implementation.resourcemanager.AzureServiceBusResourceManagerAutoConfiguration;
+import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryFactory;
 import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryCustomizer;
 import com.azure.spring.cloud.core.implementation.util.AzurePasswordlessPropertiesUtils;
 import com.azure.spring.cloud.core.implementation.util.ReflectionUtils;
@@ -19,6 +24,7 @@
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.jms.autoconfigure.JmsAutoConfiguration;
 import org.springframework.boot.jms.autoconfigure.JndiConnectionFactoryAutoConfiguration;
@@ -31,8 +37,11 @@
 import java.net.URI;
 import java.util.EnumMap;
 import java.util.Map;
+import java.util.Properties;
 import java.util.function.BiFunction;
 
+import static com.azure.spring.cloud.autoconfigure.implementation.util.SpringPasswordlessPropertiesUtils.enhancePasswordlessProperties;
+
 /**
  * {@link EnableAutoConfiguration Auto-configuration} for Azure Service Bus JMS support.
  *
@@ -59,6 +68,34 @@ AzureServiceBusJmsProperties serviceBusJmsProperties(AzureGlobalProperties azure
         return mergeAzureProperties(azureGlobalProperties, properties);
     }
 
+    @Bean
+    @ConditionalOnMissingBean
+    AzureServiceBusJmsConnectionFactoryFactory azureServiceBusJmsConnectionFactoryFactory(
+        final AzureServiceBusJmsProperties properties) {
+        if (!properties.isPasswordlessEnabled()) {
+            return () -> new ServiceBusJmsConnectionFactory(
+                properties.getConnectionString(),
+                new ServiceBusJmsConnectionFactorySettings());
+        }
+
+        String hostName =
+            properties.getNamespace()
+                + "."
+                + properties.getProfile().getEnvironment().getServiceBusDomainName();
+        Properties passwordlessProperties = properties.toPasswordlessProperties();
+        enhancePasswordlessProperties(AzureServiceBusJmsProperties.PREFIX, properties, passwordlessProperties);
+        TokenCredentialProvider tokenCredentialProvider = TokenCredentialProvider.createDefault(
+            new TokenCredentialProviderOptions(passwordlessProperties));
+
+        return () -> {
+            TokenCredential tokenCredential = tokenCredentialProvider.get();
+            return new ServiceBusJmsConnectionFactory(
+                tokenCredential,
+                hostName,
+                new ServiceBusJmsConnectionFactorySettings());
+        };
+    }
+
     /**
      * Standard tier does not support the property "com.microsoft:is-client-provider", so remove it.
      */

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryConfiguration.java

@@ -6,6 +6,7 @@
 import com.azure.servicebus.jms.ServiceBusJmsConnectionFactory;
 import com.azure.spring.cloud.autoconfigure.implementation.jms.properties.AzureServiceBusJmsProperties;
 import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryCustomizer;
+import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryFactory;
 import jakarta.jms.ConnectionFactory;
 import org.messaginghub.pooled.jms.JmsPoolConnectionFactory;
 import org.springframework.beans.BeansException;
@@ -50,13 +51,15 @@ class ServiceBusJmsConnectionFactoryConfiguration {
      *
      * @param properties the Azure Service Bus JMS properties
      * @param customizers the list of customizers to apply
+     * @param instanceFactory the factory used to create the ServiceBusJmsConnectionFactory instance
      * @return a configured ServiceBusJmsConnectionFactory instance
      */
     static ServiceBusJmsConnectionFactory createServiceBusJmsConnectionFactory(
         AzureServiceBusJmsProperties properties,
-        java.util.List<AzureServiceBusJmsConnectionFactoryCustomizer> customizers) {
+        java.util.List<AzureServiceBusJmsConnectionFactoryCustomizer> customizers,
+        AzureServiceBusJmsConnectionFactoryFactory instanceFactory) {
         return new ServiceBusJmsConnectionFactoryFactory(properties, customizers)
-            .createConnectionFactory(ServiceBusJmsConnectionFactory.class);
+            .createConnectionFactory(instanceFactory);
     }
 
     /**
@@ -195,10 +198,13 @@ private void registerJmsPoolConnectionFactory(BeanDefinitionRegistry registry) {
 
         private ServiceBusJmsConnectionFactory createServiceBusJmsConnectionFactory() {
             AzureServiceBusJmsProperties serviceBusJmsProperties = beanFactory.getBean(AzureServiceBusJmsProperties.class);
+            AzureServiceBusJmsConnectionFactoryFactory instanceFactory =
+                beanFactory.getBean(AzureServiceBusJmsConnectionFactoryFactory.class);
             ObjectProvider<AzureServiceBusJmsConnectionFactoryCustomizer> factoryCustomizers = beanFactory.getBeanProvider(AzureServiceBusJmsConnectionFactoryCustomizer.class);
             return ServiceBusJmsConnectionFactoryConfiguration.createServiceBusJmsConnectionFactory(
                 serviceBusJmsProperties,
-                factoryCustomizers.orderedStream().collect(Collectors.toList()));
+                factoryCustomizers.orderedStream().collect(Collectors.toList()),
+                instanceFactory);
         }
     }
 }

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryFactory.java

@@ -3,44 +3,34 @@
 
 package com.azure.spring.cloud.autoconfigure.implementation.jms;
 
-import com.azure.core.credential.TokenCredential;
-import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
-import com.azure.identity.extensions.implementation.credential.provider.TokenCredentialProvider;
 import com.azure.servicebus.jms.ServiceBusJmsConnectionFactory;
-import com.azure.servicebus.jms.ServiceBusJmsConnectionFactorySettings;
 import com.azure.spring.cloud.autoconfigure.implementation.jms.properties.AzureServiceBusJmsProperties;
 import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryCustomizer;
+import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryFactory;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
 
-import java.lang.reflect.InvocationTargetException;
 import java.util.Collections;
 import java.util.List;
-import java.util.Properties;
-
-import static com.azure.spring.cloud.autoconfigure.implementation.util.SpringPasswordlessPropertiesUtils.enhancePasswordlessProperties;
 
 class ServiceBusJmsConnectionFactoryFactory {
     private final AzureServiceBusJmsProperties properties;
     private final List<AzureServiceBusJmsConnectionFactoryCustomizer> factoryCustomizers;
-    private final TokenCredentialProvider tokenCredentialProvider;
 
     ServiceBusJmsConnectionFactoryFactory(AzureServiceBusJmsProperties properties,
                                           List<AzureServiceBusJmsConnectionFactoryCustomizer> factoryCustomizers) {
         Assert.notNull(properties, "Properties must not be null");
         this.properties = properties;
         this.factoryCustomizers = (factoryCustomizers != null) ? factoryCustomizers : Collections.emptyList();
-        if (properties.isPasswordlessEnabled()) {
-            Properties passwordlessProperties = properties.toPasswordlessProperties();
-            enhancePasswordlessProperties(AzureServiceBusJmsProperties.PREFIX, properties, passwordlessProperties);
-            this.tokenCredentialProvider = TokenCredentialProvider.createDefault(new TokenCredentialProviderOptions(passwordlessProperties));
-        } else {
-            this.tokenCredentialProvider = null;
-        }
     }
 
-    <T extends ServiceBusJmsConnectionFactory> T createConnectionFactory(Class<T> factoryClass) {
-        T factory = createConnectionFactoryInstance(factoryClass);
+    ServiceBusJmsConnectionFactory createConnectionFactory(
+        AzureServiceBusJmsConnectionFactoryFactory instanceFactory) {
+        Assert.notNull(instanceFactory, "AzureServiceBusJmsConnectionFactoryFactory must not be null");
+        ServiceBusJmsConnectionFactory factory = createConnectionFactoryInstance(instanceFactory);
+        Assert.notNull(
+            factory,
+            "AzureServiceBusJmsConnectionFactoryFactory must create a non-null ServiceBusJmsConnectionFactory");
         setClientId(factory);
         setPrefetchPolicy(factory);
         customize(factory);
@@ -65,26 +55,9 @@ private <T extends ServiceBusJmsConnectionFactory> void setPrefetchPolicy(T fact
             String.valueOf(prefetchProperties.getTopicPrefetch()));
     }
 
-    private <T extends ServiceBusJmsConnectionFactory> T createConnectionFactoryInstance(Class<T> factoryClass) {
-        try {
-            T factory;
-            if (properties.isPasswordlessEnabled()) {
-                String hostName =
-                    properties.getNamespace() + "." + properties.getProfile().getEnvironment().getServiceBusDomainName();
-                TokenCredential tokenCredential = tokenCredentialProvider.get();
-                factory = factoryClass.getConstructor(TokenCredential.class, String.class,
-                                          ServiceBusJmsConnectionFactorySettings.class)
-                                      .newInstance(tokenCredential, hostName,
-                                          new ServiceBusJmsConnectionFactorySettings());
-            } else {
-                factory = factoryClass.getConstructor(String.class, ServiceBusJmsConnectionFactorySettings.class)
-                                      .newInstance(properties.getConnectionString(),
-                                          new ServiceBusJmsConnectionFactorySettings());
-            }
-            return factory;
-        } catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException | IllegalArgumentException | InvocationTargetException ex) {
-            throw new IllegalStateException("Unable to create JmsConnectionFactory", ex);
-        }
+    private ServiceBusJmsConnectionFactory createConnectionFactoryInstance(
+        AzureServiceBusJmsConnectionFactoryFactory instanceFactory) {
+        return instanceFactory.createServiceBusJmsConnectionFactory();
     }
 
     private void customize(ServiceBusJmsConnectionFactory connectionFactory) {

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsContainerConfiguration.java

@@ -6,6 +6,7 @@
 import com.azure.servicebus.jms.ServiceBusJmsConnectionFactory;
 import com.azure.spring.cloud.autoconfigure.implementation.jms.properties.AzureServiceBusJmsProperties;
 import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryCustomizer;
+import com.azure.spring.cloud.autoconfigure.jms.AzureServiceBusJmsConnectionFactoryFactory;
 import jakarta.jms.ConnectionFactory;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -68,9 +69,10 @@ class ServiceBusJmsContainerConfiguration implements DisposableBean {
 
     private final AzureServiceBusJmsProperties azureServiceBusJMSProperties;
     private final ObjectProvider<AzureServiceBusJmsConnectionFactoryCustomizer> factoryCustomizers;
+    private final AzureServiceBusJmsConnectionFactoryFactory instanceFactory;
     private final Environment environment;
     private final JmsProperties jmsProperties;
-    
+
     // Memoized dedicated listener container ConnectionFactory instances to avoid duplicates and enable lifecycle management
     // Use ConnectionFactory type instead of concrete types to avoid NoClassDefFoundError when optional dependencies are missing
     private volatile ConnectionFactory dedicatedCachingConnectionFactory;
@@ -79,10 +81,12 @@ class ServiceBusJmsContainerConfiguration implements DisposableBean {
 
     ServiceBusJmsContainerConfiguration(AzureServiceBusJmsProperties azureServiceBusJMSProperties,
                                         ObjectProvider<AzureServiceBusJmsConnectionFactoryCustomizer> factoryCustomizers,
+                                        AzureServiceBusJmsConnectionFactoryFactory instanceFactory,
                                         Environment environment,
                                         JmsProperties jmsProperties) {
         this.azureServiceBusJMSProperties = azureServiceBusJMSProperties;
         this.factoryCustomizers = factoryCustomizers;
+        this.instanceFactory = instanceFactory;
         this.environment = environment;
         this.jmsProperties = jmsProperties;
     }
@@ -199,7 +203,8 @@ private synchronized ServiceBusJmsConnectionFactory getOrCreateDedicatedServiceB
     private ServiceBusJmsConnectionFactory createServiceBusJmsConnectionFactory() {
         return ServiceBusJmsConnectionFactoryConfiguration.createServiceBusJmsConnectionFactory(
             azureServiceBusJMSProperties,
-            factoryCustomizers.orderedStream().collect(Collectors.toList()));
+            factoryCustomizers.orderedStream().collect(Collectors.toList()),
+            instanceFactory);
     }
 
     private void configureCommonListenerContainerFactory(DefaultJmsListenerContainerFactory jmsListenerContainerFactory) {

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/jms/AzureServiceBusJmsConnectionFactoryFactory.java

@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.spring.cloud.autoconfigure.jms;
+
+import com.azure.servicebus.jms.ServiceBusJmsConnectionFactory;
+
+/**
+ * Defines how a {@link ServiceBusJmsConnectionFactory} instance is created.
+ * <p>
+ * Provide this interface as a Spring bean to customize creation of
+ * {@link ServiceBusJmsConnectionFactory} (or a subclass).
+ * </p>
+ * <p>
+ * The factory can be invoked multiple times in one application context
+ * (for example sender and listener container paths). Implementations should
+ * return a new {@link ServiceBusJmsConnectionFactory} instance per invocation,
+ * or otherwise ensure the returned instance is safe to share.
+ * </p>
+ */
+@FunctionalInterface
+public interface AzureServiceBusJmsConnectionFactoryFactory {
+
+    /**
+     * Creates an instance of {@link ServiceBusJmsConnectionFactory} or a subclass thereof.
+     *
+     * @return an instance of {@link ServiceBusJmsConnectionFactory}
+     */
+    ServiceBusJmsConnectionFactory createServiceBusJmsConnectionFactory();
+}