Merge pull request #225 from sherfert/4.1-bookmark-per-db

Keep one bookmark per database.

Author: sherfert

cypher-shell/src/integration-test/java/org/neo4j/shell/commands/CypherShellVerboseIntegrationTest.java

@@ -190,6 +190,13 @@ public void cypherWithOrder() throws CommandException {
         String serverVersion = shell.getServerVersion();
         assumeThat( version(serverVersion), greaterThanOrEqualTo(version("3.6")));
 
+        // Make sure we are creating a new NEW index
+        try {
+            shell.execute( "DROP INDEX ON :Person(age)" );
+        } catch ( Exception e ) {
+            // ignore if the index didn't exist
+        }
+
         shell.execute( "CREATE INDEX ON :Person(age)" );
         shell.execute( "CALL db.awaitIndexes()" );
 
@@ -262,8 +269,8 @@ public void cypherWithProfileWithMemory() throws CommandException {
 
         //then
         String actual = linePrinter.output();
-        assertThat(actual, containsString("| Plan      | Statement   | Version      | Planner | Runtime       | Time | DbHits | Rows | Memory (Bytes) |")); // First table
-        assertThat(actual, containsString("| Operator        | Details            | Estimated Rows | Rows | DB Hits | Cache H/M | Memory (Bytes) |")); // Second table
+        assertThat(actual.replace( " ", "" ), containsString("|Plan|Statement|Version|Planner|Runtime|Time|DbHits|Rows|Memory(Bytes)|")); // First table
+        assertThat(actual.replace( " ", "" ), containsString("|Operator|Details|EstimatedRows|Rows|DBHits|CacheH/M|Memory(Bytes)|")); // Second table
     }
 
     @Test

cypher-shell/src/main/java/org/neo4j/shell/DatabaseManager.java

@@ -13,6 +13,11 @@ public interface DatabaseManager
 
     String DATABASE_UNAVAILABLE_ERROR_CODE = "Neo.TransientError.General.DatabaseUnavailable";
 
+    /**
+     * Sets the active database name as set by the user.
+     * If the current state is connected, try to reconnect to that database.
+     * If the current state is disconnected, simply update `activeDatabaseAsSetByUser`.
+     */
     void setActiveDatabase(String databaseName) throws CommandException;
 
     String getActiveDatabaseAsSetByUser();

cypher-shell/src/main/java/org/neo4j/shell/state/BoltStateHandler.java

@@ -1,5 +1,6 @@
 package org.neo4j.shell.state;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -51,7 +52,7 @@ public class BoltStateHandler implements TransactionHandler, Connector, Database
     private String activeDatabaseNameAsSetByUser;
     private String actualDatabaseNameAsReportedByServer;
     private final boolean isInteractive;
-    private Bookmark systemBookmark;
+    private final Map<String, Bookmark> bookmarks = new HashMap<>();
     private Transaction tx = null;
 
     public BoltStateHandler(boolean isInteractive) {
@@ -75,15 +76,15 @@ public void setActiveDatabase(String databaseName) throws CommandException
         activeDatabaseNameAsSetByUser = databaseName;
         try {
             if (isConnected()) {
-                reconnect(true);
+                reconnect(databaseName, previousDatabaseName);
             }
         }
         catch (ClientException e) {
             if (isInteractive) {
                 // We want to try to connect to the previous database
                 activeDatabaseNameAsSetByUser = previousDatabaseName;
                 try {
-                    reconnect(true);
+                    reconnect(previousDatabaseName, previousDatabaseName);
                 }
                 catch (Exception e2) {
                     e.addSuppressed(e2);
@@ -161,10 +162,11 @@ public void connect( @Nonnull ConnectionConfig connectionConfig, ThrowingAction<
         }
         final AuthToken authToken = AuthTokens.basic(connectionConfig.username(), connectionConfig.password());
         try {
+            String previousDatabaseName = activeDatabaseNameAsSetByUser;
             try {
-                setActiveDatabase(connectionConfig.database());
+                activeDatabaseNameAsSetByUser = connectionConfig.database();
                 driver = getDriver(connectionConfig, authToken);
-                reconnect(command);
+                reconnect(activeDatabaseNameAsSetByUser, previousDatabaseName, command);
             } catch (org.neo4j.driver.exceptions.ServiceUnavailableException e) {
                 String scheme = connectionConfig.scheme();
                 String fallbackScheme;
@@ -191,7 +193,7 @@ public void connect( @Nonnull ConnectionConfig connectionConfig, ThrowingAction<
                     connectionConfig.encryption(),
                     connectionConfig.database());
                 driver = getDriver(connectionConfig, authToken);
-                reconnect(command);
+                reconnect(activeDatabaseNameAsSetByUser, previousDatabaseName, command);
             }
         } catch (Throwable t) {
             try {
@@ -203,35 +205,24 @@ public void connect( @Nonnull ConnectionConfig connectionConfig, ThrowingAction<
         }
     }
 
-    private void reconnect() throws CommandException {
-        reconnect(true, null);
+    private void reconnect(String databaseToConnectTo, String previousDatabase) throws CommandException {
+        reconnect(databaseToConnectTo, previousDatabase, null);
     }
 
-    private void reconnect(ThrowingAction<CommandException> command) throws CommandException {
-        reconnect(true, command);
-    }
-
-    private void reconnect(boolean keepBokmark) throws CommandException {
-        reconnect(keepBokmark, null);
-    }
-
-    private void reconnect( boolean keepBookmark, ThrowingAction<CommandException> command ) throws CommandException {
+    private void reconnect( String databaseToConnectTo,
+                            String previousDatabase,
+                            ThrowingAction<CommandException> command ) throws CommandException {
         SessionConfig.Builder builder = SessionConfig.builder();
         builder.withDefaultAccessMode( AccessMode.WRITE );
-        if ( !ABSENT_DB_NAME.equals( activeDatabaseNameAsSetByUser ) )
+        if ( !ABSENT_DB_NAME.equals( databaseToConnectTo ) )
         {
-            builder.withDatabase( activeDatabaseNameAsSetByUser );
+            builder.withDatabase( databaseToConnectTo );
         }
-        if ( session != null && keepBookmark )
+        closeSession( previousDatabase );
+        final Bookmark bookmarkForDBToConnectTo = bookmarks.get( databaseToConnectTo );
+        if ( bookmarkForDBToConnectTo != null )
         {
-            // Save the last bookmark and close the session
-            final Bookmark bookmark = session.lastBookmark();
-            session.close();
-            builder.withBookmarks( bookmark );
-        }
-        else if ( systemBookmark != null )
-        {
-            builder.withBookmarks( systemBookmark );
+            builder.withBookmarks( bookmarkForDBToConnectTo );
         }
 
         session = driver.session( builder.build() );
@@ -240,6 +231,22 @@ else if ( systemBookmark != null )
         connect(command);
     }
 
+    /**
+     * Closes the session, if there is any.
+     * Saves a bookmark for the database currently connected to.
+     * @param databaseName the name of the database currently connected to
+     */
+    private void closeSession( String databaseName )
+    {
+        if ( session != null )
+        {
+            // Save the last bookmark and close the session
+            final Bookmark bookmarkForPreviousDB = session.lastBookmark();
+            session.close();
+            bookmarks.put(databaseName, bookmarkForPreviousDB);
+        }
+    }
+
     private void connect( ThrowingAction<CommandException> command) throws CommandException
     {
         ThrowingAction<CommandException> toCall = command == null ? getPing() : () ->
@@ -323,7 +330,7 @@ public Optional<BoltResult> runCypher(@Nonnull String cypher,
             } catch (SessionExpiredException e) {
                 // Server is no longer accepting writes, reconnect and try again.
                 // If it still fails, leave it up to the user
-                reconnect();
+                reconnect(activeDatabaseNameAsSetByUser, activeDatabaseNameAsSetByUser);
                 return getBoltResult(cypher, queryParams);
             }
         }
@@ -348,10 +355,9 @@ public void changePassword( @Nonnull ConnectionConfig connectionConfig )
         try {
             driver = getDriver(connectionConfig, authToken);
 
-            SessionConfig.Builder builder = SessionConfig.builder()
-                    .withDefaultAccessMode(AccessMode.WRITE)
-                    .withDatabase(SYSTEM_DB_NAME);
-            session = driver.session(builder.build());
+            activeDatabaseNameAsSetByUser = SYSTEM_DB_NAME;
+            // Supply empty command, so that we do not run ping.
+            reconnect( SYSTEM_DB_NAME, SYSTEM_DB_NAME, () -> {} );
 
             String command;
             Value parameters;
@@ -370,17 +376,19 @@ public void changePassword( @Nonnull ConnectionConfig connectionConfig )
             connectionConfig.setPassword(connectionConfig.newPassword());
             connectionConfig.setNewPassword(null);
 
-            // Save a system bookmark to make sure we wait for the password change to propagate on reconnection
-            systemBookmark = session.lastBookmark();
-
             silentDisconnect();
         } catch (Throwable t) {
             try {
                 silentDisconnect();
             } catch (Exception e) {
                 t.addSuppressed(e);
             }
-            throw t;
+            if (t instanceof RuntimeException) {
+                throw (RuntimeException) t;
+            }
+            // The only checked exception is CommandException and we know that
+            // we cannot get that since we supply an empty command.
+            throw new RuntimeException(t);
         }
     }
 
@@ -419,9 +427,7 @@ private void resetActualDbName() {
      */
     void silentDisconnect() {
         try {
-            if (session != null) {
-                session.close();
-            }
+            closeSession( activeDatabaseNameAsSetByUser );
             if (driver != null) {
                 driver.close();
             }