From 6bef2231281bb6c1b701c09ab66490bec57ba1ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Oskar=20Damkj=C3=A6r?=
 <OskarDamkjaer@users.noreply.github.com>
Date: Mon, 16 Oct 2023 13:19:54 +0200
Subject: [PATCH] Ensure server automatically reconnects when connection is
 lost (#111)

* properly reconnects when connection is lost

* avoid using node js type directly
---
 .../src/neo4j-connection.ts                   |  2 +-
 .../src/neo4j-schema-poller.ts                | 34 +++++++++++++++++++
 packages/server/src/server.ts                 | 28 ++++++---------
 3 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/packages/neo4j-schema-poller/src/neo4j-connection.ts b/packages/neo4j-schema-poller/src/neo4j-connection.ts
index 871ef3a1..fac40edb 100644
--- a/packages/neo4j-schema-poller/src/neo4j-connection.ts
+++ b/packages/neo4j-schema-poller/src/neo4j-connection.ts
@@ -66,7 +66,7 @@ export class Neo4jConnection {
 
   async healthcheck() {
     try {
-      await this.driver.verifyConnectivity();
+      await this.driver.verifyConnectivity({ database: this.currentDb });
       return true;
     } catch (error) {
       return false;
diff --git a/packages/neo4j-schema-poller/src/neo4j-schema-poller.ts b/packages/neo4j-schema-poller/src/neo4j-schema-poller.ts
index 89b8f47a..98e986f5 100644
--- a/packages/neo4j-schema-poller/src/neo4j-schema-poller.ts
+++ b/packages/neo4j-schema-poller/src/neo4j-schema-poller.ts
@@ -6,6 +6,7 @@ import { listDatabases } from './queries/databases.js';
 export class Neo4jSchemaPoller {
   public connection?: Neo4jConnection;
   public metadata?: MetadataPoller;
+  private reconnectionTimeout?: ReturnType<typeof setTimeout>;
 
   async connect(
     url: string,
@@ -47,10 +48,43 @@ export class Neo4jSchemaPoller {
     this.metadata.startBackgroundPolling();
   }
 
+  async persistentConnect(
+    url: string,
+    credentials: { username: string; password: string },
+    config: { driverConfig?: Config; appName: string },
+  ) {
+    const shouldHaveConnection = this.connection !== undefined;
+    const connectionAlive = await this.connection?.healthcheck();
+
+    if (!connectionAlive) {
+      if (shouldHaveConnection) {
+        console.error('Connection to Neo4j dropped');
+        this.disconnect();
+      }
+
+      try {
+        await this.connect(url, credentials, config);
+        // eslint-disable-next-line no-console
+        console.log('Established connection to Neo4j');
+      } catch (error) {
+        console.error(
+          `Unable to connect to Neo4j: ${String(
+            error,
+          )}. Retrying in 30 seconds.`,
+        );
+      }
+    }
+
+    this.reconnectionTimeout = setTimeout(() => {
+      void this.persistentConnect(url, credentials, config);
+    }, 30000);
+  }
+
   disconnect() {
     this.connection?.dispose();
     this.metadata?.stopBackgroundPolling();
     this.connection = undefined;
     this.metadata = undefined;
+    clearTimeout(this.reconnectionTimeout);
   }
 }
diff --git a/packages/server/src/server.ts b/packages/server/src/server.ts
index c7f903f2..10942587 100644
--- a/packages/server/src/server.ts
+++ b/packages/server/src/server.ts
@@ -95,30 +95,24 @@ connection.onCompletion(doAutoCompletion(documents, neo4jSdk));
 
 connection.onDidChangeConfiguration(
   (params: { settings: { cypherLSP: CypherLSPSettings } }) => {
-    const neo4jConfig = params.settings.cypherLSP.neo4j;
-
     neo4jSdk.disconnect();
+
+    const neo4jConfig = params.settings.cypherLSP.neo4j;
     if (
+      neo4jSdk.connection === undefined &&
       neo4jConfig.connect &&
       neo4jConfig.password &&
       neo4jConfig.URL &&
       neo4jConfig.user
     ) {
-      neo4jSdk
-        .connect(
-          neo4jConfig.URL,
-          {
-            username: neo4jConfig.user,
-            password: neo4jConfig.password,
-          },
-          { appName: 'cypher-language-server' },
-        )
-        .then(() => {
-          neo4jSdk.metadata.startBackgroundPolling();
-        })
-        .catch((error) => {
-          console.error(`Unable to connect to Neo4j: ${String(error)}`);
-        });
+      void neo4jSdk.persistentConnect(
+        neo4jConfig.URL,
+        {
+          username: neo4jConfig.user,
+          password: neo4jConfig.password,
+        },
+        { appName: 'cypher-language-server' },
+      );
     }
   },
 );