feat: add batch processing support in SqsAutoConfiguration

Author: khc41

spring-cloud-aws-autoconfigure/src/main/java/io/awspring/cloud/autoconfigure/sqs/SqsAutoConfiguration.java

@@ -17,11 +17,7 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.awspring.cloud.autoconfigure.AwsAsyncClientCustomizer;
-import io.awspring.cloud.autoconfigure.core.AwsClientBuilderConfigurer;
-import io.awspring.cloud.autoconfigure.core.AwsClientCustomizer;
-import io.awspring.cloud.autoconfigure.core.AwsConnectionDetails;
-import io.awspring.cloud.autoconfigure.core.CredentialsProviderAutoConfiguration;
-import io.awspring.cloud.autoconfigure.core.RegionProviderAutoConfiguration;
+import io.awspring.cloud.autoconfigure.core.*;
 import io.awspring.cloud.sqs.config.SqsBootstrapConfiguration;
 import io.awspring.cloud.sqs.config.SqsListenerConfigurer;
 import io.awspring.cloud.sqs.config.SqsMessageListenerContainerFactory;
@@ -30,6 +26,7 @@
 import io.awspring.cloud.sqs.listener.errorhandler.ErrorHandler;
 import io.awspring.cloud.sqs.listener.interceptor.AsyncMessageInterceptor;
 import io.awspring.cloud.sqs.listener.interceptor.MessageInterceptor;
+import io.awspring.cloud.sqs.operations.BatchingSqsClientAdapter;
 import io.awspring.cloud.sqs.operations.SqsTemplate;
 import io.awspring.cloud.sqs.operations.SqsTemplateBuilder;
 import io.awspring.cloud.sqs.support.converter.MessagingMessageConverter;
@@ -48,8 +45,11 @@
 import org.springframework.boot.context.properties.PropertyMapper;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Import;
+import org.springframework.context.annotation.Primary;
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
 import software.amazon.awssdk.services.sqs.SqsAsyncClientBuilder;
+import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration;
+import software.amazon.awssdk.services.sqs.batchmanager.SqsAsyncBatchManager;
 import software.amazon.awssdk.services.sqs.model.Message;
 
 /**
@@ -82,9 +82,35 @@ public SqsAsyncClient sqsAsyncClient(AwsClientBuilderConfigurer awsClientBuilder
 			ObjectProvider<AwsConnectionDetails> connectionDetails,
 			ObjectProvider<SqsAsyncClientCustomizer> sqsAsyncClientCustomizers,
 			ObjectProvider<AwsAsyncClientCustomizer> awsAsyncClientCustomizers) {
-		return awsClientBuilderConfigurer.configureAsyncClient(SqsAsyncClient.builder(), this.sqsProperties,
-				connectionDetails.getIfAvailable(), configurer.getIfAvailable(),
-				sqsAsyncClientCustomizers.orderedStream(), awsAsyncClientCustomizers.orderedStream()).build();
+        return awsClientBuilderConfigurer.configureAsyncClient(SqsAsyncClient.builder(), this.sqsProperties,
+                connectionDetails.getIfAvailable(), configurer.getIfAvailable(),
+                sqsAsyncClientCustomizers.orderedStream(), awsAsyncClientCustomizers.orderedStream()).build();
+    }
+
+	@ConditionalOnProperty(name = "spring.cloud.aws.sqs.batch.enabled", havingValue = "true")
+	@Bean
+	@Primary
+	public SqsAsyncClient batchSqsAsyncClient(SqsAsyncClient sqsAsyncClient) {
+		SqsAsyncBatchManager batchManager = createBatchManager(sqsAsyncClient);
+		return new BatchingSqsClientAdapter(batchManager);
+	}
+
+	private SqsAsyncBatchManager createBatchManager(SqsAsyncClient sqsAsyncClient) {
+		return SqsAsyncBatchManager.builder()
+			.client(sqsAsyncClient)
+			.overrideConfiguration(this::configurationProperties)
+			.build();
+	}
+
+	private void configurationProperties(BatchOverrideConfiguration.Builder options) {
+		PropertyMapper mapper = PropertyMapper.get().alwaysApplyingWhenNonNull();
+		mapper.from(this.sqsProperties.getBatch().getMaxNumberOfMessages()).to(options::maxBatchSize);
+		mapper.from(this.sqsProperties.getBatch().getSendBatchFrequency()).to(options::sendRequestFrequency);
+		mapper.from(this.sqsProperties.getBatch().getWaitTimeSeconds()).to(options::receiveMessageMinWaitDuration);
+		mapper.from(this.sqsProperties.getBatch().getVisibilityTimeout()).to(options::receiveMessageVisibilityTimeout);
+		mapper.from(this.sqsProperties.getBatch().getSystemAttributeNames())
+			.to(options::receiveMessageSystemAttributeNames);
+		mapper.from(this.sqsProperties.getBatch().getAttributeNames()).to(options::receiveMessageAttributeNames);
 	}
 
 	@ConditionalOnMissingBean
@@ -93,15 +119,15 @@ public SqsTemplate sqsTemplate(SqsAsyncClient sqsAsyncClient, ObjectProvider<Obj
 			ObjectProvider<ObservationRegistry> observationRegistryProvider,
 			ObjectProvider<SqsTemplateObservation.Convention> observationConventionProvider,
 			MessagingMessageConverter<Message> messageConverter) {
-		SqsTemplateBuilder builder = SqsTemplate.builder().sqsAsyncClient(sqsAsyncClient)
-				.messageConverter(messageConverter);
+        SqsTemplateBuilder builder = SqsTemplate.builder().sqsAsyncClient(sqsAsyncClient)
+                .messageConverter(messageConverter);
 		objectMapperProvider.ifAvailable(om -> setMapperToConverter(messageConverter, om));
-		if (this.sqsProperties.isObservationEnabled()) {
-			observationRegistryProvider
-					.ifAvailable(registry -> builder.configure(options -> options.observationRegistry(registry)));
-			observationConventionProvider
-					.ifAvailable(convention -> builder.configure(options -> options.observationConvention(convention)));
-		}
+        if (this.sqsProperties.isObservationEnabled()) {
+            observationRegistryProvider
+                    .ifAvailable(registry -> builder.configure(options -> options.observationRegistry(registry)));
+            observationConventionProvider
+                    .ifAvailable(convention -> builder.configure(options -> options.observationConvention(convention)));
+        }
 		if (sqsProperties.getQueueNotFoundStrategy() != null) {
 			builder.configure((options) -> options.queueNotFoundStrategy(sqsProperties.getQueueNotFoundStrategy()));
 		}
@@ -127,12 +153,12 @@ public SqsMessageListenerContainerFactory<Object> defaultSqsListenerContainerFac
 		interceptors.forEach(factory::addMessageInterceptor);
 		asyncInterceptors.forEach(factory::addMessageInterceptor);
 		objectMapperProvider.ifAvailable(om -> setMapperToConverter(messagingMessageConverter, om));
-		if (this.sqsProperties.isObservationEnabled()) {
-			observationRegistry
-					.ifAvailable(registry -> factory.configure(options -> options.observationRegistry(registry)));
-			observationConventionProvider
-					.ifAvailable(convention -> factory.configure(options -> options.observationConvention(convention)));
-		}
+        if (this.sqsProperties.isObservationEnabled()) {
+            observationRegistry
+                    .ifAvailable(registry -> factory.configure(options -> options.observationRegistry(registry)));
+            observationConventionProvider
+                    .ifAvailable(convention -> factory.configure(options -> options.observationConvention(convention)));
+        }
 		factory.configure(options -> options.messageConverter(messagingMessageConverter));
 		return factory;
 	}