diff --git a/java/sqlcommenter-java/build.gradle b/java/sqlcommenter-java/build.gradle index c965bcf4..818ceb67 100644 --- a/java/sqlcommenter-java/build.gradle +++ b/java/sqlcommenter-java/build.gradle @@ -61,6 +61,21 @@ dependencies { compile (group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.1.5.RELEASE') { exclude group: 'org.apache.logging.log4j', module: 'log4j-to-slf4j' } + + compile (group: 'org.springframework.boot', name: 'spring-boot-starter-webflux', version: '2.7.7') { + exclude group: 'org.apache.logging.log4j', module: 'log4j-to-slf4j' + } + + compile (group: 'org.springframework.boot', name: 'spring-boot-starter-aop', version: '2.7.7') { + exclude group: 'org.apache.logging.log4j', module: 'log4j-to-slf4j' + } + + testImplementation 'io.projectreactor:reactor-test:3.5.2' + + compile 'io.r2dbc:r2dbc-proxy:1.1.0.RELEASE' + compile 'io.r2dbc:r2dbc-spi:1.0.0.RELEASE' + testImplementation 'io.r2dbc:r2dbc-h2:1.0.0.RELEASE' + compile 'org.apache.logging.log4j:log4j-to-slf4j:2.16.0' compile 'org.hibernate:hibernate-core:5.4.3.Final' diff --git a/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/filter/SpringSQLCommenterWebFilter.java b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/filter/SpringSQLCommenterWebFilter.java new file mode 100644 index 00000000..725daece --- /dev/null +++ b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/filter/SpringSQLCommenterWebFilter.java @@ -0,0 +1,34 @@ +package com.google.cloud.sqlcommenter.filter; + +import com.google.cloud.sqlcommenter.threadlocalstorage.State; +import org.springframework.web.method.HandlerMethod; +import org.springframework.web.reactive.result.method.annotation.RequestMappingHandlerMapping; +import org.springframework.web.server.ServerWebExchange; +import org.springframework.web.server.WebFilter; +import org.springframework.web.server.WebFilterChain; +import reactor.core.publisher.Mono; + +public class SpringSQLCommenterWebFilter implements WebFilter { + + private final RequestMappingHandlerMapping handlerMapping; + + public SpringSQLCommenterWebFilter(RequestMappingHandlerMapping handlerMapping) { + this.handlerMapping = handlerMapping; + } + + @Override + public Mono filter(ServerWebExchange exchange, WebFilterChain chain) { + HandlerMethod handlerMethod = + (HandlerMethod) this.handlerMapping.getHandler(exchange).toFuture().getNow(null); + + State state = + State.newBuilder() + .withActionName(handlerMethod.getMethod().getName()) + .withFramework("spring") + .withControllerName( + handlerMethod.getBeanType().getSimpleName().replace("Controller", "")) + .build(); + + return chain.filter(exchange).contextWrite(ctx -> ctx.put("state", state)); + } +} diff --git a/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionDecorator.java b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionDecorator.java new file mode 100644 index 00000000..c27c36e4 --- /dev/null +++ b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionDecorator.java @@ -0,0 +1,43 @@ +package com.google.cloud.sqlcommenter.r2dbc; + +import io.r2dbc.spi.Connection; +import java.lang.reflect.Proxy; +import org.reactivestreams.Subscription; +import reactor.core.CoreSubscriber; + +public class ConnectionDecorator implements CoreSubscriber { + + private final CoreSubscriber delegate; + + public ConnectionDecorator(CoreSubscriber delegate) { + this.delegate = delegate; + } + + @Override + public void onSubscribe(Subscription s) { + this.delegate.onSubscribe(s); + } + + @Override + public void onNext(Object o) { + assert o instanceof Connection; + Connection connection = (Connection) o; + + Object proxied = + Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + new ConnectionInvocationHandler(connection, delegate.currentContext())); + this.delegate.onNext(proxied); + } + + @Override + public void onError(Throwable t) { + this.delegate.onError(t); + } + + @Override + public void onComplete() { + this.delegate.onComplete(); + } +} diff --git a/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionFactoryAspect.java b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionFactoryAspect.java new file mode 100644 index 00000000..7ea69692 --- /dev/null +++ b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionFactoryAspect.java @@ -0,0 +1,25 @@ +package com.google.cloud.sqlcommenter.r2dbc; + +import org.aspectj.lang.ProceedingJoinPoint; +import org.aspectj.lang.annotation.Around; +import org.aspectj.lang.annotation.Aspect; +import org.reactivestreams.Publisher; +import reactor.core.publisher.Mono; +import reactor.core.publisher.Operators; + +@Aspect +public class ConnectionFactoryAspect { + + @Around("execution(* io.r2dbc.spi.ConnectionFactory.create(..)) ") + public Object beforeSampleCreation(ProceedingJoinPoint joinPoint) throws Throwable { + Object object = joinPoint.proceed(); + + @SuppressWarnings("unchecked") + Publisher publisher = (Publisher) object; + + return Mono.from(publisher) + .transform( + Operators.liftPublisher( + (publisher1, coreSubscriber) -> new ConnectionDecorator(coreSubscriber))); + } +} diff --git a/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionInvocationHandler.java b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionInvocationHandler.java new file mode 100644 index 00000000..fb47d20a --- /dev/null +++ b/java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionInvocationHandler.java @@ -0,0 +1,35 @@ +package com.google.cloud.sqlcommenter.r2dbc; + +import com.google.cloud.sqlcommenter.threadlocalstorage.State; +import io.r2dbc.spi.Connection; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import reactor.util.context.ContextView; + +public class ConnectionInvocationHandler implements InvocationHandler { + + private final Connection connection; + private final ContextView contextView; + + public ConnectionInvocationHandler(Connection connection, ContextView contextView) { + this.connection = connection; + this.contextView = contextView; + } + + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + + if ("createStatement".equals(methodName)) { + String query = (String) args[0]; + if (contextView != null && contextView.hasKey("state")) { + State state = contextView.get("state"); + query = state.formatAndAppendToSQL(query); + } + + return method.invoke(connection, query); + } else { + return method.invoke(connection, args); + } + } +} diff --git a/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/filter/SpringSQLCommenterWebFilterTest.java b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/filter/SpringSQLCommenterWebFilterTest.java new file mode 100644 index 00000000..9fa78036 --- /dev/null +++ b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/filter/SpringSQLCommenterWebFilterTest.java @@ -0,0 +1,39 @@ +package com.google.cloud.sqlcommenter.filter; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.springframework.mock.http.server.reactive.MockServerHttpRequest; +import org.springframework.mock.web.server.MockServerWebExchange; +import org.springframework.web.reactive.result.method.RequestMappingInfo; +import org.springframework.web.reactive.result.method.annotation.RequestMappingHandlerMapping; +import org.springframework.web.server.WebFilterChain; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +@RunWith(JUnit4.class) +public class SpringSQLCommenterWebFilterTest { + + @Test + public void testPreHandlePlacesStateInContextView() throws NoSuchMethodException { + RequestMappingHandlerMapping requestMappingHandlerMapping = new RequestMappingHandlerMapping(); + RequestMappingInfo info = RequestMappingInfo.paths("/test").build(); + requestMappingHandlerMapping.registerMapping( + info, + this, + SpringSQLCommenterWebFilterTest.class.getMethod("testPreHandlePlacesStateInContextView")); + + SpringSQLCommenterWebFilter springSQLCommenterWebFilter = + new SpringSQLCommenterWebFilter(requestMappingHandlerMapping); + + WebFilterChain filterChain = filterExchange -> Mono.empty(); + + MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.get("/test")); + + StepVerifier.create(springSQLCommenterWebFilter.filter(exchange, filterChain)) + .expectAccessibleContext() + .hasKey("state") + .then() + .verifyComplete(); + } +} diff --git a/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionDecoratorTest.java b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionDecoratorTest.java new file mode 100644 index 00000000..4254cb63 --- /dev/null +++ b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionDecoratorTest.java @@ -0,0 +1,61 @@ +package com.google.cloud.sqlcommenter.r2dbc; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import io.r2dbc.spi.Connection; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.reactivestreams.Subscription; +import reactor.core.CoreSubscriber; + +@RunWith(JUnit4.class) +public class ConnectionDecoratorTest { + + private CoreSubscriber delegate; + private ConnectionDecorator connectionDecorator; + + @Before + public void setUp() throws Exception { + delegate = mock(CoreSubscriber.class); + connectionDecorator = new ConnectionDecorator(delegate); + } + + @Test + public void testOnSubscribe() { + Subscription mocked = mock(Subscription.class); + connectionDecorator.onSubscribe(mocked); + verify(delegate, times(1)).onSubscribe(mocked); + } + + @Test + public void testOnNext() { + Connection connection = mock(Connection.class); + connectionDecorator.onNext(connection); + verify(delegate, times(1)).currentContext(); + verify(delegate, times(1)).onNext(any()); + } + + @Test(expected = AssertionError.class) + public void testOnNextNonConnection() { + Object object = new Object(); + connectionDecorator.onNext(object); + } + + @Test + public void testOnError() { + Throwable mocked = mock(Throwable.class); + connectionDecorator.onError(mocked); + verify(delegate, times(1)).onError(mocked); + } + + @Test + public void testOnComplete() { + connectionDecorator.onComplete(); + verify(delegate, times(1)).onComplete(); + } +} diff --git a/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionFactoryAspectTest.java b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionFactoryAspectTest.java new file mode 100644 index 00000000..1b694d5e --- /dev/null +++ b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionFactoryAspectTest.java @@ -0,0 +1,50 @@ +package com.google.cloud.sqlcommenter.r2dbc; + +import static org.junit.Assert.assertEquals; + +import com.google.cloud.sqlcommenter.threadlocalstorage.State; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionFactory; +import io.r2dbc.spi.Statement; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.reactivestreams.Publisher; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.test.annotation.DirtiesContext; +import org.springframework.test.context.ContextConfiguration; +import org.springframework.test.context.junit4.SpringJUnit4ClassRunner; +import org.springframework.test.util.ReflectionTestUtils; +import reactor.core.publisher.Mono; +import reactor.util.context.Context; + +@RunWith(SpringJUnit4ClassRunner.class) +@ContextConfiguration(classes = R2DBCConfiguration.class) +@DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_EACH_TEST_METHOD) +public class ConnectionFactoryAspectTest { + + private final String stmt1 = "SELECT * from FOO"; + + @Autowired private ConnectionFactory connectionFactory; + + @Test + public void name() { + State state = + State.newBuilder() + .withControllerName("Order") + .withFramework("spring") + .withActionName("add") + .build(); + Context updatedContext = Context.of("state", state); + + Publisher conn = connectionFactory.create(); + + Statement statement = + Mono.from(conn) + .contextWrite(updatedContext) + .map(connection -> connection.createStatement(stmt1)) + .block(); + + String value = (String) ReflectionTestUtils.getField(statement, null, "sql"); + assertEquals(value, state.formatAndAppendToSQL(stmt1)); + } +} diff --git a/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionInvocationHandlerTest.java b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionInvocationHandlerTest.java new file mode 100644 index 00000000..2d740405 --- /dev/null +++ b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionInvocationHandlerTest.java @@ -0,0 +1,77 @@ +package com.google.cloud.sqlcommenter.r2dbc; + +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.cloud.sqlcommenter.threadlocalstorage.State; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.Statement; +import java.lang.reflect.Proxy; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import reactor.util.context.Context; + +@RunWith(JUnit4.class) +public class ConnectionInvocationHandlerTest { + + private final String stmt1 = "SELECT * from FOO"; + + private Connection connection = mock(Connection.class); + private Connection proxiedConnection; + private Context context = Context.empty(); + private ConnectionInvocationHandler connectionInvocationHandler; + + @Before + public void setUp() throws Exception { + connection = mock(Connection.class); + when(connection.createStatement(anyString())).thenAnswer(i -> mock(Statement.class)); + + connectionInvocationHandler = new ConnectionInvocationHandler(connection, context); + proxiedConnection = proxyConnection(connectionInvocationHandler); + } + + @Test + public void testInvokeCreateStatementEmptyState() { + proxiedConnection.createStatement(stmt1); + verify(connection).createStatement(eq(stmt1)); + } + + @Test + public void testInvokeCreateStatement() { + State state = + State.newBuilder() + .withControllerName("Order") + .withFramework("spring") + .withActionName("add") + .build(); + Context updatedContext = context.put("state", state); + + connectionInvocationHandler = new ConnectionInvocationHandler(connection, updatedContext); + proxiedConnection = proxyConnection(connectionInvocationHandler); + + proxiedConnection.createStatement(stmt1); + String appendedState = state.formatAndAppendToSQL(stmt1); + verify(connection).createStatement(eq(appendedState)); + } + + @Test + public void testInvokeDelegate() { + proxiedConnection.beginTransaction(); + proxiedConnection.createStatement(stmt1); + verify(connection, times(1)).beginTransaction(); + } + + private Connection proxyConnection(ConnectionInvocationHandler connectionInvocationHandler) { + return (Connection) + Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + connectionInvocationHandler); + } +} diff --git a/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/R2DBCConfiguration.java b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/R2DBCConfiguration.java new file mode 100644 index 00000000..83026003 --- /dev/null +++ b/java/sqlcommenter-java/src/test/java/com/google/cloud/sqlcommenter/r2dbc/R2DBCConfiguration.java @@ -0,0 +1,25 @@ +package com.google.cloud.sqlcommenter.r2dbc; + +import io.r2dbc.h2.H2ConnectionFactory; +import io.r2dbc.spi.ConnectionFactory; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.context.annotation.EnableAspectJAutoProxy; + +@Configuration +@EnableAspectJAutoProxy +public class R2DBCConfiguration { + + @Bean + public ConnectionFactory connectionFactory() { + return new H2ConnectionFactory( + io.r2dbc.h2.H2ConnectionConfiguration.builder() + .url("mem:testdb;MODE=PostgreSQL;DB_CLOSE_DELAY=-1;") + .build()); + } + + @Bean + public ConnectionFactoryAspect connectionFactoryAspect() { + return new ConnectionFactoryAspect(); + } +}