From 75ae7c968aa0e6eaab1444c72de2e6953f3a8ab8 Mon Sep 17 00:00:00 2001 From: Tadaya Tsuyukubo Date: Sun, 28 Apr 2024 08:22:58 -0700 Subject: [PATCH] Add ProxyConnectionFactoryCustomizer See gh-40555 --- .../R2dbcObservationAutoConfiguration.java | 28 ++++++++++-- ...2dbcObservationAutoConfigurationTests.java | 43 ++++++++++++++++++- spring-boot-project/spring-boot/build.gradle | 1 + .../ProxyConnectionFactoryCustomizer.java | 36 ++++++++++++++++ 4 files changed, 102 insertions(+), 6 deletions(-) create mode 100644 spring-boot-project/spring-boot/src/main/java/org/springframework/boot/r2dbc/ProxyConnectionFactoryCustomizer.java diff --git a/spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfiguration.java b/spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfiguration.java index 75f619c1bdf..bcd1cf69b26 100644 --- a/spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfiguration.java +++ b/spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfiguration.java @@ -1,5 +1,5 @@ /* - * Copyright 2012-2023 the original author or authors. + * Copyright 2012-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -33,12 +33,15 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.context.properties.EnableConfigurationProperties; import org.springframework.boot.r2dbc.ConnectionFactoryDecorator; import org.springframework.boot.r2dbc.OptionsCapableConnectionFactory; +import org.springframework.boot.r2dbc.ProxyConnectionFactoryCustomizer; import org.springframework.context.annotation.Bean; +import org.springframework.core.annotation.Order; /** * {@link EnableAutoConfiguration Auto-configuration} for R2DBC observability support. * * @author Moritz Halbritter + * @author Tadaya Tsuyukubo * @since 3.2.0 */ @AutoConfiguration(after = ObservationAutoConfiguration.class) @@ -46,20 +49,37 @@ import org.springframework.context.annotation.Bean; @EnableConfigurationProperties(R2dbcObservationProperties.class) public class R2dbcObservationAutoConfiguration { + /** + * {@code @Order} value of observation customizer. + */ + public static final int R2DBC_PROXY_OBSERVATION_CUSTOMIZER_ORDER = 1000; + @Bean + ConnectionFactoryDecorator connectionFactoryDecorator( + ObjectProvider customizers) { + return (connectionFactory) -> { + ProxyConnectionFactory.Builder builder = ProxyConnectionFactory.builder(connectionFactory); + customizers.orderedStream().forEach((customizer) -> customizer.customize(builder)); + return builder.build(); + }; + } + + @Bean + @Order(R2DBC_PROXY_OBSERVATION_CUSTOMIZER_ORDER) @ConditionalOnBean(ObservationRegistry.class) - ConnectionFactoryDecorator connectionFactoryDecorator(R2dbcObservationProperties properties, + ProxyConnectionFactoryCustomizer proxyConnectionFactoryObservationCustomizer(R2dbcObservationProperties properties, ObservationRegistry observationRegistry, ObjectProvider queryObservationConvention, ObjectProvider queryParametersTagProvider) { - return (connectionFactory) -> { + return (builder) -> { + ConnectionFactory connectionFactory = builder.getConnectionFactory(); HostAndPort hostAndPort = extractHostAndPort(connectionFactory); ObservationProxyExecutionListener listener = new ObservationProxyExecutionListener(observationRegistry, connectionFactory, hostAndPort.host(), hostAndPort.port()); listener.setIncludeParameterValues(properties.isIncludeParameterValues()); queryObservationConvention.ifAvailable(listener::setQueryObservationConvention); queryParametersTagProvider.ifAvailable(listener::setQueryParametersTagProvider); - return ProxyConnectionFactory.builder(connectionFactory).listener(listener).build(); + builder.listener(listener); }; } diff --git a/spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfigurationTests.java b/spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfigurationTests.java index e5ae366b49b..815819030ec 100644 --- a/spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfigurationTests.java +++ b/spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/r2dbc/R2dbcObservationAutoConfigurationTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2012-2023 the original author or authors. + * Copyright 2012-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,6 +16,8 @@ package org.springframework.boot.actuate.autoconfigure.r2dbc; +import java.util.ArrayList; +import java.util.List; import java.util.UUID; import java.util.concurrent.atomic.AtomicReference; @@ -25,6 +27,7 @@ import io.micrometer.observation.ObservationRegistry; import io.r2dbc.spi.ConnectionFactory; import org.awaitility.Awaitility; import org.hamcrest.Matchers; +import org.junit.jupiter.api.Order; import org.junit.jupiter.api.Test; import reactor.core.publisher.Mono; @@ -33,9 +36,12 @@ import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.context.annotation.ImportCandidates; import org.springframework.boot.r2dbc.ConnectionFactoryBuilder; import org.springframework.boot.r2dbc.ConnectionFactoryDecorator; +import org.springframework.boot.r2dbc.ProxyConnectionFactoryCustomizer; import org.springframework.boot.test.context.FilteredClassLoader; import org.springframework.boot.test.context.assertj.AssertableApplicationContext; import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; import static org.assertj.core.api.Assertions.assertThat; @@ -43,6 +49,7 @@ import static org.assertj.core.api.Assertions.assertThat; * Tests for {@link R2dbcObservationAutoConfiguration}. * * @author Moritz Halbritter + * @author Tadaya Tsuyukubo */ class R2dbcObservationAutoConfigurationTests { @@ -78,7 +85,20 @@ class R2dbcObservationAutoConfigurationTests { @Test void shouldNotSupplyBeansIfObservationRegistryIsNotPresent() { this.runnerWithoutObservationRegistry - .run((context) -> assertThat(context).doesNotHaveBean(ConnectionFactoryDecorator.class)); + .run((context) -> assertThat(context).doesNotHaveBean(ProxyConnectionFactoryCustomizer.class)); + } + + @Test + void shouldApplyCustomizers() { + this.runner.withUserConfiguration(ProxyConnectionFactoryCustomizerConfig.class).run((context) -> { + ConnectionFactoryDecorator decorator = context.getBean(ConnectionFactoryDecorator.class); + ConnectionFactory connectionFactory = ConnectionFactoryBuilder + .withUrl("r2dbc:h2:mem:///" + UUID.randomUUID()) + .build(); + decorator.decorate(connectionFactory); + assertThat(context.getBean(ProxyConnectionFactoryCustomizerConfig.class).called).containsExactly("first", + "second"); + }); } @Test @@ -128,4 +148,23 @@ class R2dbcObservationAutoConfigurationTests { } + @Configuration(proxyBeanMethods = false) + private static final class ProxyConnectionFactoryCustomizerConfig { + + private final List called = new ArrayList<>(); + + @Bean + @Order(R2dbcObservationAutoConfiguration.R2DBC_PROXY_OBSERVATION_CUSTOMIZER_ORDER - 1) + ProxyConnectionFactoryCustomizer first() { + return (builder) -> this.called.add("first"); + } + + @Bean + @Order(R2dbcObservationAutoConfiguration.R2DBC_PROXY_OBSERVATION_CUSTOMIZER_ORDER + 1) + ProxyConnectionFactoryCustomizer second() { + return (builder) -> this.called.add("second"); + } + + } + } diff --git a/spring-boot-project/spring-boot/build.gradle b/spring-boot-project/spring-boot/build.gradle index 0e2f235e55b..2de14ff476a 100644 --- a/spring-boot-project/spring-boot/build.gradle +++ b/spring-boot-project/spring-boot/build.gradle @@ -35,6 +35,7 @@ dependencies { optional("io.projectreactor:reactor-tools") optional("io.projectreactor.netty:reactor-netty-http") optional("io.r2dbc:r2dbc-pool") + optional("io.r2dbc:r2dbc-proxy") optional("io.rsocket:rsocket-core") optional("io.rsocket:rsocket-transport-netty") optional("io.undertow:undertow-servlet") diff --git a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/r2dbc/ProxyConnectionFactoryCustomizer.java b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/r2dbc/ProxyConnectionFactoryCustomizer.java new file mode 100644 index 00000000000..b4ce1972bd9 --- /dev/null +++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/r2dbc/ProxyConnectionFactoryCustomizer.java @@ -0,0 +1,36 @@ +/* + * Copyright 2012-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.boot.r2dbc; + +import io.r2dbc.proxy.ProxyConnectionFactory; + +/** + * Callback interface that can be used to customize a + * {@link ProxyConnectionFactory.Builder}. + * + * @author Tadaya Tsuyukubo + * @since 3.3 + */ +public interface ProxyConnectionFactoryCustomizer { + + /** + * Callback to customize a {@link ProxyConnectionFactory.Builder} instance. + * @param builder the builder to customize + */ + void customize(ProxyConnectionFactory.Builder builder); + +}