[Spring] Support custom ServiceBus JMS factory creation

rujche · rujche · commit b68fb6effadc · 2026-06-30T14:26:39.000+08:00
diff --git a/sdk/spring/CHANGELOG.md b/sdk/spring/CHANGELOG.md
@@ -8,6 +8,8 @@ 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.
+
 #### Breaking Changes
 
 #### Bugs Fixed
diff --git a/sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsAutoConfiguration.java b/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,24 @@ AzureServiceBusJmsProperties serviceBusJmsProperties(AzureGlobalProperties azure
         return mergeAzureProperties(azureGlobalProperties, properties);
     }
 
+    @Bean
+    @ConditionalOnMissingBean
+    AzureServiceBusJmsConnectionFactoryFactory azureServiceBusJmsConnectionFactoryFactory(final AzureServiceBusJmsProperties properties) {
+        return () -> {
+            if (properties.isPasswordlessEnabled()) {
+                String hostName =
+                    properties.getNamespace() + "." + properties.getProfile().getEnvironment().getServiceBusDomainName();
+                Properties passwordlessProperties = properties.toPasswordlessProperties();
+                enhancePasswordlessProperties(AzureServiceBusJmsProperties.PREFIX, properties, passwordlessProperties);
+                TokenCredentialProvider tokenCredentialProvider = TokenCredentialProvider.createDefault(new TokenCredentialProviderOptions(passwordlessProperties));
+                TokenCredential tokenCredential = tokenCredentialProvider.get();
+                return new ServiceBusJmsConnectionFactory(tokenCredential, hostName, new ServiceBusJmsConnectionFactorySettings());
+            } else {
+                return new ServiceBusJmsConnectionFactory(properties.getConnectionString(), new ServiceBusJmsConnectionFactorySettings());
+            }
+        };
+    }
+
     /**
      * Standard tier does not support the property "com.microsoft:is-client-provider", so remove it.
      */
diff --git a/sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryConfiguration.java b/sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryConfiguration.java
@@ -3,9 +3,14 @@
 
 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 jakarta.jms.ConnectionFactory;
 import org.messaginghub.pooled.jms.JmsPoolConnectionFactory;
 import org.springframework.beans.BeansException;
@@ -26,8 +31,10 @@
 import org.springframework.jms.connection.CachingConnectionFactory;
 import org.springframework.util.ClassUtils;
 
+import java.util.Properties;
 import java.util.stream.Collectors;
 
+import static com.azure.spring.cloud.autoconfigure.implementation.util.SpringPasswordlessPropertiesUtils.enhancePasswordlessProperties;
 import static org.springframework.beans.factory.support.BeanDefinitionBuilder.genericBeanDefinition;
 
 /**
@@ -56,7 +63,19 @@ static ServiceBusJmsConnectionFactory createServiceBusJmsConnectionFactory(
         AzureServiceBusJmsProperties properties,
         java.util.List<AzureServiceBusJmsConnectionFactoryCustomizer> customizers) {
         return new ServiceBusJmsConnectionFactoryFactory(properties, customizers)
-            .createConnectionFactory(ServiceBusJmsConnectionFactory.class);
+            .createConnectionFactory(() -> {
+                if (properties.isPasswordlessEnabled()) {
+                    String hostName =
+                        properties.getNamespace() + "." + properties.getProfile().getEnvironment().getServiceBusDomainName();
+                    Properties passwordlessProperties = properties.toPasswordlessProperties();
+                    enhancePasswordlessProperties(AzureServiceBusJmsProperties.PREFIX, properties, passwordlessProperties);
+                    TokenCredentialProvider tokenCredentialProvider = TokenCredentialProvider.createDefault(new TokenCredentialProviderOptions(passwordlessProperties));
+                    TokenCredential tokenCredential = tokenCredentialProvider.get();
+                    return new ServiceBusJmsConnectionFactory(tokenCredential, hostName, new ServiceBusJmsConnectionFactorySettings());
+                } else {
+                    return new ServiceBusJmsConnectionFactory(properties.getConnectionString(), new ServiceBusJmsConnectionFactorySettings());
+                }
+            });
     }
 
     /**
@@ -195,10 +214,11 @@ 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()));
+            return new ServiceBusJmsConnectionFactoryFactory(serviceBusJmsProperties,
+                factoryCustomizers.orderedStream().collect(Collectors.toList()))
+                .createConnectionFactory(instanceFactory);
         }
     }
 }
diff --git a/sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryFactory.java b/sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryFactory.java
@@ -3,44 +3,29 @@
 
 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) {
+        ServiceBusJmsConnectionFactory factory = createConnectionFactoryInstance(instanceFactory);
         setClientId(factory);
         setPrefetchPolicy(factory);
         customize(factory);
@@ -65,26 +50,8 @@ 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) {
diff --git a/sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/jms/AzureServiceBusJmsConnectionFactoryFactory.java b/sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/jms/AzureServiceBusJmsConnectionFactoryFactory.java
@@ -0,0 +1,20 @@
+// 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;
+
+/**
+ * The interface used to define how the {@link ServiceBusJmsConnectionFactory} instance is created.
+ */
+@FunctionalInterface
+public interface AzureServiceBusJmsConnectionFactoryFactory {
+
+    /**
+     * Creates an instance of {@link ServiceBusJmsConnectionFactory} or a subclass thereof.
+     *
+     * @return an instance of {@link ServiceBusJmsConnectionFactory}
+     */
+    ServiceBusJmsConnectionFactory createServiceBusJmsConnectionFactory();
+}
diff --git a/sdk/spring/spring-cloud-azure-autoconfigure/src/test/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryConfigurationTests.java b/sdk/spring/spring-cloud-azure-autoconfigure/src/test/java/com/azure/spring/cloud/autoconfigure/implementation/jms/ServiceBusJmsConnectionFactoryConfigurationTests.java
@@ -3,8 +3,11 @@
 
 package com.azure.spring.cloud.autoconfigure.implementation.jms;
 
+import com.azure.core.credential.TokenCredential;
 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.jms.AzureServiceBusJmsConnectionFactoryFactory;
 import jakarta.jms.Connection;
 import jakarta.jms.ConnectionFactory;
 import jakarta.jms.Destination;
@@ -19,6 +22,7 @@
 import org.springframework.boot.jms.autoconfigure.JmsAutoConfiguration;
 import org.springframework.boot.test.context.FilteredClassLoader;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.PropertySource;
 import org.springframework.jms.connection.CachingConnectionFactory;
@@ -180,6 +184,72 @@ void fallbackToServiceBusConnectionFactoryWhenNoCachingOrPoolClassesPresent(Stri
             });
     }
 
+    @Test
+    void useCustomServiceBusJmsConnectionFactoryClassForServiceBusFactory() {
+        this.contextRunner
+            .withUserConfiguration(CustomConnectionFactoryClassConfiguration.class)
+            .withPropertyValues(
+                "spring.jms.servicebus.pricing-tier=premium",
+                "spring.jms.servicebus.pool.enabled=false",
+                "spring.jms.cache.enabled=false"
+            )
+            .run(context -> {
+                assertThat(context).hasSingleBean(ServiceBusJmsConnectionFactory.class);
+                assertThat(context.getBean(ServiceBusJmsConnectionFactory.class))
+                    .isInstanceOf(CustomServiceBusJmsConnectionFactory.class);
+            });
+    }
+
+    @Test
+    void useCustomServiceBusJmsConnectionFactoryClassForPasswordlessServiceBusFactory() {
+        this.contextRunner
+            .withUserConfiguration(CustomConnectionFactoryClassConfiguration.class)
+            .withPropertyValues(
+                "spring.jms.servicebus.pricing-tier=premium",
+                "spring.jms.servicebus.passwordless-enabled=true",
+                "spring.jms.servicebus.namespace=test-namespace",
+                "spring.jms.servicebus.pool.enabled=false",
+                "spring.jms.cache.enabled=false"
+            )
+            .run(context -> {
+                assertThat(context).hasSingleBean(ServiceBusJmsConnectionFactory.class);
+                assertThat(context.getBean(ServiceBusJmsConnectionFactory.class))
+                    .isInstanceOf(CustomServiceBusJmsConnectionFactory.class);
+            });
+    }
+
+    @Test
+    void useCustomServiceBusJmsConnectionFactoryClassForCachingFactory() {
+        this.contextRunner
+            .withUserConfiguration(CustomConnectionFactoryClassConfiguration.class)
+            .withPropertyValues(
+                "spring.jms.servicebus.pricing-tier=premium",
+                "spring.jms.cache.enabled=true"
+            )
+            .run(context -> {
+                assertThat(context).hasSingleBean(CachingConnectionFactory.class);
+                CachingConnectionFactory cachingConnectionFactory = context.getBean(CachingConnectionFactory.class);
+                assertThat(cachingConnectionFactory.getTargetConnectionFactory())
+                    .isInstanceOf(CustomServiceBusJmsConnectionFactory.class);
+            });
+    }
+
+    @Test
+    void useCustomServiceBusJmsConnectionFactoryClassForPoolingFactory() {
+        this.contextRunner
+            .withUserConfiguration(CustomConnectionFactoryClassConfiguration.class)
+            .withPropertyValues(
+                "spring.jms.servicebus.pricing-tier=premium",
+                "spring.jms.servicebus.pool.enabled=true"
+            )
+            .run(context -> {
+                assertThat(context).hasSingleBean(JmsPoolConnectionFactory.class);
+                JmsPoolConnectionFactory poolConnectionFactory = context.getBean(JmsPoolConnectionFactory.class);
+                assertThat(poolConnectionFactory.getConnectionFactory())
+                    .isInstanceOf(CustomServiceBusJmsConnectionFactory.class);
+            });
+    }
+
     @Test
     void cachingConnectionFactoryReusesSameProducerForSameDestination() throws Exception {
         // Create mock objects for JMS components
@@ -276,4 +346,25 @@ private Session createServiceBusJmsSession(Session innerSession) throws Exceptio
     static class AdditionalPropertySourceConfiguration {
 
     }
+
+    @Configuration
+    static class CustomConnectionFactoryClassConfiguration {
+        @Bean
+        AzureServiceBusJmsConnectionFactoryFactory connectionFactoryFactory() {
+            return () -> new CustomServiceBusJmsConnectionFactory(
+                String.format(CONNECTION_STRING_FORMAT, "test-namespace"),
+                new ServiceBusJmsConnectionFactorySettings());
+        }
+    }
+
+    static class CustomServiceBusJmsConnectionFactory extends ServiceBusJmsConnectionFactory {
+        public CustomServiceBusJmsConnectionFactory(String connectionString, ServiceBusJmsConnectionFactorySettings settings) {
+            super(connectionString, settings);
+        }
+
+        public CustomServiceBusJmsConnectionFactory(TokenCredential tokenCredential, String host,
+            ServiceBusJmsConnectionFactorySettings settings) {
+            super(tokenCredential, host, settings);
+        }
+    }
 }
