Added filter for sql commenter

Author: gkatzioura

java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/filter/SpringSQLCommenterWebFilter.java

@@ -0,0 +1,36 @@
+package com.google.cloud.sqlcommenter.filter;
+
+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 com.google.cloud.sqlcommenter.threadlocalstorage.State;
+
+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,44 @@
+package com.google.cloud.sqlcommenter.r2dbc;
+
+import java.lang.reflect.Proxy;
+
+import org.reactivestreams.Subscription;
+
+import io.r2dbc.spi.Connection;
+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,27 @@
+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 io.r2dbc.spi.Connection;
+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 {
+        Publisher<? extends Connection> publisher = (Publisher<? extends Connection>) joinPoint.proceed();
+
+
+        return Mono.from(publisher)
+                   .cast(Object.class)
+                   .transform(Operators.liftPublisher((publisher1, coreSubscriber) -> new ConnectionDecorator(coreSubscriber)));
+    }
+
+}

java/sqlcommenter-java/src/main/java/com/google/cloud/sqlcommenter/r2dbc/ConnectionInvocationHandler.java

@@ -0,0 +1,38 @@
+package com.google.cloud.sqlcommenter.r2dbc;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+
+import com.google.cloud.sqlcommenter.threadlocalstorage.State;
+
+import io.r2dbc.spi.Connection;
+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) {
+                State state = contextView.get("state");
+                query = state.formatAndAppendToSQL(query);
+            }
+
+            return method.invoke(connection, query);
+        } else {
+            return method.invoke(connection, args);
+        }
+    }
+
+}