Fix nullability in ConnectionInfo

Mark `getConnectionId()` and `getOriginalConnection()` methods with
`@Nullable` and updated their javadoc.

Fixes #151

Author: ttddyy

src/main/java/io/r2dbc/proxy/callback/ConnectionFactoryCallbackHandler.java

@@ -61,8 +61,11 @@ public Object invoke(Object proxy, Method method, @Nullable Object[] args) throw
             executionInfo.setMethodArgs(args);
             executionInfo.setTarget(target);
 
-            // The actual connection is not yet created, but in order to populate connection level context-view,
-            // create an empty connection info here and set it to the execution info
+            // The actual connection has not been created yet, but to populate the connection-level context view (mainly for ValueStore),
+            // an empty ConnectionInfo is created here and set into the execution info.
+            // NOTE: Until it is associated with an actual connection, retrieving the connection (`getOriginalConnection`) and connection ID (`getConnectionId`)
+            // will return null. For example, during `beforeMethod` of `ConnectionFactory#create()`.
+            // This is a design decision to always provide a connection-level ValueStore.
             DefaultConnectionInfo connectionInfo = new DefaultConnectionInfo();
             executionInfo.setConnectionInfo(connectionInfo);
 

src/main/java/io/r2dbc/proxy/callback/DefaultConnectionInfo.java

@@ -21,6 +21,7 @@
 import io.r2dbc.proxy.core.ValueStore;
 import io.r2dbc.proxy.util.Assert;
 import io.r2dbc.spi.Connection;
+import reactor.util.annotation.Nullable;
 
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
@@ -91,11 +92,13 @@ public void setClosed(boolean closed) {
     }
 
     @Override
+    @Nullable
     public Connection getOriginalConnection() {
         return this.originalConnection;
     }
 
     @Override
+    @Nullable
     public String getConnectionId() {
         return this.connectionId;
     }

src/main/java/io/r2dbc/proxy/callback/StatementCallbackHandler.java

@@ -73,12 +73,6 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
             return handleCommonMethod(methodName, this.statement, args, this.connectionInfo.getOriginalConnection());
         }
 
-//        if ("unwrap".equals(methodName)) {
-//            return this.statement;
-//        } else if ("unwrapConnection".equals(methodName)) {
-//            return this.connectionInfo.getOriginalConnection();
-//        }
-
         if ("bind".equals(methodName) || "bindNull".equals(methodName)) {
 
             BoundValue boundValue;

src/main/java/io/r2dbc/proxy/core/ConnectionInfo.java

@@ -17,28 +17,38 @@
 package io.r2dbc.proxy.core;
 
 
+import io.r2dbc.proxy.listener.ProxyExecutionListener;
 import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactory;
+import reactor.util.annotation.Nullable;
 
 /**
  * Hold {@link Connection} related information.
  *
  * @author Tadaya Tsuyukubo
  */
 public interface ConnectionInfo {
-
     /**
      * Retrieve original {@link Connection}.
      *
-     * @return connection
+     * @return connection; {@code null} is returned when {@link ConnectionInfo} is evaluated
+     * before it is associated with an actual connection, for example, during
+     * {@link ProxyExecutionListener#beforeMethod(MethodExecutionInfo)} for
+     * {@link ConnectionFactory#create()}.
      */
+    @Nullable
     Connection getOriginalConnection();
 
     /**
      * Get ID for the connection.
      *
-     * @return connection ID
+     * @return connection ID; {@code null} is returned when {@link ConnectionInfo} is evaluated
+     * before it is associated with an actual connection, for example, during
+     * {@link ProxyExecutionListener#beforeMethod(MethodExecutionInfo)} for
+     * {@link ConnectionFactory#create()}.
      * @see io.r2dbc.proxy.callback.ConnectionIdManager
      */
+    @Nullable
     String getConnectionId();
 
     /**

src/test/java/io/r2dbc/proxy/callback/ConnectionFactoryCallbackHandlerTest.java

@@ -19,6 +19,7 @@
 import io.r2dbc.proxy.core.ConnectionInfo;
 import io.r2dbc.proxy.core.MethodExecutionInfo;
 import io.r2dbc.proxy.listener.LastExecutionAwareListener;
+import io.r2dbc.proxy.listener.ProxyExecutionListener;
 import io.r2dbc.proxy.listener.ProxyMethodExecutionListener;
 import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.ConnectionFactory;
@@ -114,6 +115,63 @@ void createConnection() throws Throwable {
         assertThat(afterMethod.getResult()).isSameAs(originalConnection);
     }
 
+    @Test
+    void connectionInfoInBeforeAndAfterMethod() throws Throwable {
+        ConnectionFactory connectionFactory = mock(ConnectionFactory.class);
+        Connection originalConnection = mock(Connection.class);
+        ConnectionIdManager idManager = mock(ConnectionIdManager.class);
+
+        // mock call to original ConnectionFactory#create()
+        doReturn(Mono.just(originalConnection)).when(connectionFactory).create();
+
+        String connectionId = "100";
+        when(idManager.getId(originalConnection)).thenReturn(connectionId);
+
+        AtomicReference<Connection> connectionBeforeMethod = new AtomicReference<>();
+        AtomicReference<Connection> connectionAfterMethod = new AtomicReference<>();
+        AtomicReference<String> connectionIdBeforeMethod = new AtomicReference<>();
+        AtomicReference<String> connectionIdAfterMethod = new AtomicReference<>();
+        ProxyExecutionListener listener = new ProxyExecutionListener() {
+
+            @Override
+            public void beforeMethod(MethodExecutionInfo executionInfo) {
+                ConnectionInfo connectionInfo = executionInfo.getConnectionInfo();
+                assertThat(connectionInfo).isNotNull();
+                connectionBeforeMethod.set(connectionInfo.getOriginalConnection());
+                connectionIdBeforeMethod.set(connectionInfo.getConnectionId());
+            }
+
+            @Override
+            public void afterMethod(MethodExecutionInfo executionInfo) {
+                ConnectionInfo connectionInfo = executionInfo.getConnectionInfo();
+                assertThat(connectionInfo).isNotNull();
+                connectionAfterMethod.set(connectionInfo.getOriginalConnection());
+                connectionIdAfterMethod.set(connectionInfo.getConnectionId());
+            }
+        };
+
+        ProxyConfig proxyConfig = ProxyConfig.builder()
+            .connectionIdManager(idManager)
+            .listener(listener)
+            .build();
+
+        ConnectionFactoryCallbackHandler handler = new ConnectionFactoryCallbackHandler(connectionFactory, proxyConfig);
+
+        Object result = handler.invoke(connectionFactory, CREATE_METHOD, null);
+
+        assertThat(result).isInstanceOf(Publisher.class);
+
+        StepVerifier.create((Publisher<?>) result)
+            .expectSubscription()
+            .expectNextCount(1)
+            .verifyComplete();
+
+        assertThat(connectionBeforeMethod).hasValue(null);
+        assertThat(connectionIdBeforeMethod).hasValue(null);
+        assertThat(connectionAfterMethod).hasValue(originalConnection);
+        assertThat(connectionIdAfterMethod).hasValue("100");
+    }
+
     @Test
     void getMetadata() throws Throwable {