SQL commenter for R2DBC

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'

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<Void> 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));
+  }
+}

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<Object> {
+
+  private final CoreSubscriber<Object> delegate;
+
+  public ConnectionDecorator(CoreSubscriber<Object> 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();
+  }
+}

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<Object> publisher = (Publisher<Object>) object;
+
+    return Mono.from(publisher)
+        .transform(
+            Operators.liftPublisher(
+                (publisher1, coreSubscriber) -> new ConnectionDecorator(coreSubscriber)));
+  }
+}

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);
+    }
+  }
+}

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();
+  }
+}

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<Object> 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();
+  }
+}

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<? extends Connection> 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));
+  }
+}

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);
+  }
+}