From 6c9ac2453ae12bef1d27fc2fa78161d9be660667 Mon Sep 17 00:00:00 2001
From: Alex Thomas <alexthomas93@users.noreply.github.com>
Date: Tue, 7 Jan 2025 10:13:53 +0000
Subject: [PATCH] Disables Neo4jGraph driver warnings (#29)

* Disables Neo4jGraph driver warnings

* Fixed linting error

* More linting fixes
---
 .../langchain_neo4j/graphs/neo4j_graph.py     | 22 +++++++--
 .../unit_tests/graphs/test_neo4j_graph.py     | 49 ++++++++++++++++++-
 2 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/libs/neo4j/langchain_neo4j/graphs/neo4j_graph.py b/libs/neo4j/langchain_neo4j/graphs/neo4j_graph.py
index 237c5d5..b1ced4b 100644
--- a/libs/neo4j/langchain_neo4j/graphs/neo4j_graph.py
+++ b/libs/neo4j/langchain_neo4j/graphs/neo4j_graph.py
@@ -310,6 +310,7 @@ class Neo4jGraph(GraphStore):
     enhanced_schema (bool): A flag whether to scan the database for
             example values and use them in the graph schema. Default is False.
     driver_config (Dict): Configuration passed to Neo4j Driver.
+            Defaults to {"notifications_min_severity", "OFF"} if not set.
 
     *Security note*: Make sure that the database connection uses credentials
         that are narrowly-scoped to only include necessary permissions.
@@ -365,9 +366,10 @@ def __init__(
             {"database": database}, "database", "NEO4J_DATABASE", "neo4j"
         )
 
-        self._driver = neo4j.GraphDatabase.driver(
-            url, auth=auth, **(driver_config or {})
-        )
+        if driver_config is None:
+            driver_config = {}
+        driver_config.setdefault("notifications_min_severity", "OFF")
+        self._driver = neo4j.GraphDatabase.driver(url, auth=auth, **driver_config)
         self._database = database
         self.timeout = timeout
         self.sanitize = sanitize
@@ -377,6 +379,20 @@ def __init__(
         # Verify connection
         try:
             self._driver.verify_connectivity()
+        except neo4j.exceptions.ConfigurationError as e:
+            # If notification filtering is not supported
+            if "Notification filtering is not supported" in str(e):
+                # Retry without notifications_min_severity
+                driver_config.pop("notifications_min_severity", None)
+                self._driver = neo4j.GraphDatabase.driver(
+                    url, auth=auth, **driver_config
+                )
+                self._driver.verify_connectivity()
+            else:
+                raise ValueError(
+                    "Could not connect to Neo4j database. "
+                    "Please ensure that the driver config is correct"
+                )
         except neo4j.exceptions.ServiceUnavailable:
             raise ValueError(
                 "Could not connect to Neo4j database. "
diff --git a/libs/neo4j/tests/unit_tests/graphs/test_neo4j_graph.py b/libs/neo4j/tests/unit_tests/graphs/test_neo4j_graph.py
index 265af31..60b79b1 100644
--- a/libs/neo4j/tests/unit_tests/graphs/test_neo4j_graph.py
+++ b/libs/neo4j/tests/unit_tests/graphs/test_neo4j_graph.py
@@ -3,7 +3,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from neo4j.exceptions import ClientError, Neo4jError
+from neo4j.exceptions import ClientError, ConfigurationError, Neo4jError
 
 from langchain_neo4j.graphs.neo4j_graph import (
     LIST_LIMIT,
@@ -201,7 +201,52 @@ def test_neo4j_graph_init_with_empty_credentials() -> None:
         Neo4jGraph(
             url="bolt://localhost:7687", username="", password="", refresh_schema=False
         )
-        mock_driver.assert_called_with("bolt://localhost:7687", auth=None)
+        mock_driver.assert_called_with(
+            "bolt://localhost:7687", auth=None, notifications_min_severity="OFF"
+        )
+
+
+def test_neo4j_graph_init_notification_filtering_err() -> None:
+    """Test the __init__ method when notification filtering is disabled."""
+    with patch("neo4j.GraphDatabase.driver", autospec=True) as mock_driver:
+        mock_driver_instance = MagicMock()
+        mock_driver.return_value = mock_driver_instance
+        err = ConfigurationError("Notification filtering is not supported")
+        mock_driver_instance.verify_connectivity.side_effect = [err, None]
+        Neo4jGraph(
+            url="bolt://localhost:7687",
+            username="username",
+            password="password",
+            refresh_schema=False,
+        )
+        mock_driver.assert_any_call(
+            "bolt://localhost:7687",
+            auth=("username", "password"),
+            notifications_min_severity="OFF",
+        )
+        # The first call verify_connectivity should fail causing the driver to be
+        # recreated without the notifications_min_severity parameter
+        mock_driver.assert_any_call(
+            "bolt://localhost:7687",
+            auth=("username", "password"),
+        )
+
+
+def test_neo4j_graph_init_driver_config_err() -> None:
+    """Test the __init__ method with an incorrect driver config."""
+    with patch("neo4j.GraphDatabase.driver", autospec=True) as mock_driver:
+        mock_driver_instance = MagicMock()
+        mock_driver.return_value = mock_driver_instance
+        err = ConfigurationError()
+        mock_driver_instance.verify_connectivity.side_effect = err
+        with pytest.raises(ValueError) as exc_info:
+            Neo4jGraph(
+                url="bolt://localhost:7687",
+                username="username",
+                password="password",
+                refresh_schema=False,
+            )
+        assert "Please ensure that the driver config is correct" in str(exc_info.value)
 
 
 def test_init_apoc_procedure_not_found(