Async invocation of routing procedures

lutovich · lutovich · commit d0f9341261f6 · 2017-10-03T12:11:53.000+02:00
This commit makes `RoutingProcedureRunner` able to invoke routing
procedures asynchronously using given connection future. Server
version check is performed to determine which procedure to call.
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureRunner.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureRunner.java
@@ -20,12 +20,17 @@
 package org.neo4j.driver.internal.cluster;
 
 import java.util.List;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.CompletionStage;
 
 import org.neo4j.driver.ResultResourcesHandler;
 import org.neo4j.driver.internal.NetworkSession;
+import org.neo4j.driver.internal.async.AsyncConnection;
+import org.neo4j.driver.internal.async.QueryRunner;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.v1.Record;
 import org.neo4j.driver.v1.Statement;
+import org.neo4j.driver.v1.StatementResultCursor;
 import org.neo4j.driver.v1.exceptions.ClientException;
 
 import static org.neo4j.driver.internal.util.ServerVersion.v3_2_0;
@@ -47,16 +52,7 @@ public RoutingProcedureRunner( RoutingContext context )
 
     public RoutingProcedureResponse run( Connection connection )
     {
-        Statement procedure;
-        if( version( connection.server().version() ).greaterThanOrEqual( v3_2_0 ) )
-        {
-            procedure = new Statement( "CALL " + GET_ROUTING_TABLE,
-                    parameters( GET_ROUTING_TABLE_PARAM, context.asMap() ) );
-        }
-        else
-        {
-            procedure = new Statement( "CALL " + GET_SERVERS );
-        }
+        Statement procedure = procedureStatement( connection.server().version() );
 
         try
         {
@@ -68,8 +64,58 @@ public RoutingProcedureResponse run( Connection connection )
         }
     }
 
+    public CompletionStage<RoutingProcedureResponse> run( CompletionStage<AsyncConnection> connectionStage )
+    {
+        return connectionStage.thenCompose( connection ->
+        {
+            Statement procedure = procedureStatement( connection.serverInfo().version() );
+            return runProcedure( connection, procedure ).handle( ( records, error ) ->
+            {
+                if ( error != null )
+                {
+                    return handleError( procedure, error );
+                }
+                else
+                {
+                    return new RoutingProcedureResponse( procedure, records );
+                }
+            } );
+        } );
+    }
+
     List<Record> runProcedure( Connection connection, Statement procedure )
     {
         return NetworkSession.run( connection, procedure, ResultResourcesHandler.NO_OP ).list();
     }
+
+    CompletionStage<List<Record>> runProcedure( AsyncConnection connection, Statement procedure )
+    {
+        return QueryRunner.runAsync( connection, procedure )
+                .thenCompose( StatementResultCursor::listAsync );
+    }
+
+    private Statement procedureStatement( String serverVersionString )
+    {
+        if ( version( serverVersionString ).greaterThanOrEqual( v3_2_0 ) )
+        {
+            return new Statement( "CALL " + GET_ROUTING_TABLE,
+                    parameters( GET_ROUTING_TABLE_PARAM, context.asMap() ) );
+        }
+        else
+        {
+            return new Statement( "CALL " + GET_SERVERS );
+        }
+    }
+
+    private RoutingProcedureResponse handleError( Statement procedure, Throwable error )
+    {
+        if ( error instanceof ClientException )
+        {
+            return new RoutingProcedureResponse( procedure, error );
+        }
+        else
+        {
+            throw new CompletionException( error );
+        }
+    }
 }
diff --git a/driver/src/test/java/org/neo4j/driver/internal/cluster/RoutingProcedureRunnerTest.java b/driver/src/test/java/org/neo4j/driver/internal/cluster/RoutingProcedureRunnerTest.java
@@ -22,19 +22,31 @@
 
 import java.net.URI;
 import java.util.List;
+import java.util.concurrent.CompletionStage;
 
+import org.neo4j.driver.internal.async.AsyncConnection;
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.summary.InternalServerInfo;
 import org.neo4j.driver.v1.Record;
 import org.neo4j.driver.v1.Statement;
 import org.neo4j.driver.v1.Value;
+import org.neo4j.driver.v1.exceptions.ClientException;
 
+import static java.util.Arrays.asList;
 import static java.util.Collections.EMPTY_MAP;
-import static org.hamcrest.MatcherAssert.assertThat;
+import static java.util.Collections.singletonList;
+import static java.util.concurrent.CompletableFuture.completedFuture;
 import static org.hamcrest.core.IsEqual.equalTo;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
+import static org.neo4j.driver.internal.async.Futures.failedFuture;
+import static org.neo4j.driver.internal.async.Futures.getBlocking;
 import static org.neo4j.driver.internal.cluster.RoutingProcedureRunner.GET_ROUTING_TABLE;
 import static org.neo4j.driver.internal.cluster.RoutingProcedureRunner.GET_ROUTING_TABLE_PARAM;
 import static org.neo4j.driver.internal.cluster.RoutingProcedureRunner.GET_SERVERS;
@@ -58,6 +70,20 @@ public void shouldCallGetRoutingTableWithEmptyMap() throws Throwable
                 new Statement( "CALL " + GET_ROUTING_TABLE, parameters( GET_ROUTING_TABLE_PARAM, EMPTY_MAP ) ) ) );
     }
 
+    @Test
+    public void shouldCallGetRoutingTableWithEmptyMapAsync()
+    {
+        RoutingProcedureRunner runner = new TestRoutingProcedureRunner( RoutingContext.EMPTY,
+                completedFuture( asList( mock( Record.class ), mock( Record.class ) ) ) );
+
+        RoutingProcedureResponse response = getBlocking( runner.run( connectionStage( "Neo4j/3.2.1" ) ) );
+
+        assertTrue( response.isSuccess() );
+        assertEquals( 2, response.records().size() );
+        assertEquals( new Statement( "CALL " + GET_ROUTING_TABLE, parameters( GET_ROUTING_TABLE_PARAM, EMPTY_MAP ) ),
+                response.procedure() );
+    }
+
     @Test
     public void shouldCallGetRoutingTableWithParam() throws Throwable
     {
@@ -77,6 +103,23 @@ public void shouldCallGetRoutingTableWithParam() throws Throwable
                 new Statement( "CALL " + GET_ROUTING_TABLE, expectedParams ) ) );
     }
 
+    @Test
+    public void shouldCallGetRoutingTableWithParamAsync()
+    {
+        URI uri = URI.create( "bolt+routing://localhost/?key1=value1&key2=value2" );
+        RoutingContext context = new RoutingContext( uri );
+
+        RoutingProcedureRunner runner = new TestRoutingProcedureRunner( context,
+                completedFuture( singletonList( mock( Record.class ) ) ) );
+
+        RoutingProcedureResponse response = getBlocking( runner.run( connectionStage( "Neo4j/3.2.1" ) ) );
+
+        assertTrue( response.isSuccess() );
+        assertEquals( 1, response.records().size() );
+        Value expectedParams = parameters( GET_ROUTING_TABLE_PARAM, context.asMap() );
+        assertEquals( new Statement( "CALL " + GET_ROUTING_TABLE, expectedParams ), response.procedure() );
+    }
+
     @Test
     public void shouldCallGetServers() throws Throwable
     {
@@ -95,11 +138,72 @@ public void shouldCallGetServers() throws Throwable
                 new Statement( "CALL " + GET_SERVERS ) ) );
     }
 
+    @Test
+    public void shouldCallGetServersAsync()
+    {
+        URI uri = URI.create( "bolt+routing://localhost/?key1=value1&key2=value2" );
+        RoutingContext context = new RoutingContext( uri );
+
+        RoutingProcedureRunner runner = new TestRoutingProcedureRunner( context,
+                completedFuture( asList( mock( Record.class ), mock( Record.class ) ) ) );
+
+        RoutingProcedureResponse response = getBlocking( runner.run( connectionStage( "Neo4j/3.1.8" ) ) );
+
+        assertTrue( response.isSuccess() );
+        assertEquals( 2, response.records().size() );
+        assertEquals( new Statement( "CALL " + GET_SERVERS ), response.procedure() );
+    }
+
+    @Test
+    public void shouldReturnFailedResponseOnClientException()
+    {
+        ClientException error = new ClientException( "Hi" );
+        RoutingProcedureRunner runner = new TestRoutingProcedureRunner( RoutingContext.EMPTY, failedFuture( error ) );
+
+        RoutingProcedureResponse response = getBlocking( runner.run( connectionStage( "Neo4j/3.2.2" ) ) );
+
+        assertFalse( response.isSuccess() );
+        assertEquals( error, response.error() );
+    }
+
+    @Test
+    public void shouldReturnFailedStageOnError()
+    {
+        Exception error = new Exception( "Hi" );
+        RoutingProcedureRunner runner = new TestRoutingProcedureRunner( RoutingContext.EMPTY, failedFuture( error ) );
+
+        try
+        {
+            getBlocking( runner.run( connectionStage( "Neo4j/3.2.2" ) ) );
+            fail( "Exception expected" );
+        }
+        catch ( Exception e )
+        {
+            assertEquals( error, e );
+        }
+    }
+
+    private static CompletionStage<AsyncConnection> connectionStage( String serverVersion )
+    {
+        AsyncConnection connection = mock( AsyncConnection.class );
+        InternalServerInfo serverInfo = new InternalServerInfo( new BoltServerAddress( "123:45" ), serverVersion );
+        when( connection.serverInfo() ).thenReturn( serverInfo );
+        return completedFuture( connection );
+    }
+
     private static class TestRoutingProcedureRunner extends RoutingProcedureRunner
     {
+        final CompletionStage<List<Record>> runProcedureResult;
+
         TestRoutingProcedureRunner( RoutingContext context )
+        {
+            this( context, null );
+        }
+
+        TestRoutingProcedureRunner( RoutingContext context, CompletionStage<List<Record>> runProcedureResult )
         {
             super( context );
+            this.runProcedureResult = runProcedureResult;
         }
 
         @Override
@@ -108,6 +212,12 @@ List<Record> runProcedure( Connection connection, Statement procedure )
             // I do not want any network traffic
             return null;
         }
+
+        @Override
+        CompletionStage<List<Record>> runProcedure( AsyncConnection connection, Statement procedure )
+        {
+            return runProcedureResult;
+        }
     }
 
 }
