added neo4j db name

bdavis · bdavis · commit 2e4d3a5802ae · 2026-02-24T10:10:22.000-08:00
diff --git a/src/codegraphcontext/cli/config_manager.py b/src/codegraphcontext/cli/config_manager.py
@@ -16,7 +16,7 @@
 CONFIG_FILE = CONFIG_DIR / ".env"
 
 # Database credential keys (stored in same .env file but not managed as config)
-DATABASE_CREDENTIAL_KEYS = {"NEO4J_URI", "NEO4J_USERNAME", "NEO4J_PASSWORD"}
+DATABASE_CREDENTIAL_KEYS = {"NEO4J_URI", "NEO4J_USERNAME", "NEO4J_PASSWORD", "NEO4J_DATABASE"}
 
 # Default configuration values
 DEFAULT_CONFIG = {
diff --git a/src/codegraphcontext/cli/main.py b/src/codegraphcontext/cli/main.py
@@ -288,7 +288,11 @@ def _load_credentials():
             os.environ.get("NEO4J_PASSWORD")
         ])
         if has_neo4j_creds:
-            console.print("[cyan]Using database: Neo4j[/cyan]")
+            neo4j_db = os.environ.get("NEO4J_DATABASE")
+            if neo4j_db:
+                console.print(f"[cyan]Using database: Neo4j (database: {neo4j_db})[/cyan]")
+            else:
+                console.print("[cyan]Using database: Neo4j[/cyan]")
         else:
             console.print("[yellow]⚠ DEFAULT_DATABASE=neo4j but credentials not found. Falling back to FalkorDB.[/yellow]")
     else:
@@ -700,7 +704,7 @@ def doctor():
             
             if uri and username and password:
                 console.print(f"   [cyan]Testing Neo4j connection to {uri}...[/cyan]")
-                is_connected, error_msg = DatabaseManager.test_connection(uri, username, password)
+                is_connected, error_msg = DatabaseManager.test_connection(uri, username, password, database=os.environ.get("NEO4J_DATABASE"))
                 if is_connected:
                     console.print(f"   [green]✓[/green] Neo4j connection successful")
                 else:
diff --git a/src/codegraphcontext/cli/setup_wizard.py b/src/codegraphcontext/cli/setup_wizard.py
@@ -377,7 +377,7 @@ def configure_mcp_client():
                     if line and not line.startswith("#") and "=" in line:
                         key, value = line.split("=", 1)
                         key = key.strip()
-                        if key in ["NEO4J_URI", "NEO4J_USERNAME", "NEO4J_PASSWORD"]:
+                        if key in ["NEO4J_URI", "NEO4J_USERNAME", "NEO4J_PASSWORD", "NEO4J_DATABASE"]:
                             env_vars[key] = value.strip()
         except Exception:
             pass
diff --git a/src/codegraphcontext/core/database.py b/src/codegraphcontext/core/database.py
@@ -10,6 +10,24 @@
 
 from codegraphcontext.utils.debug_log import debug_log, info_logger, error_logger, warning_logger
 
+class Neo4jDriverWrapper:
+    """
+    A simple wrapper around the Neo4j Driver to inject a database name into session() calls.
+    """
+    def __init__(self, driver: Driver, database: str = None):
+        self._driver = driver
+        self._database = database
+
+    def session(self, **kwargs):
+        """Proxy method to get a session from the underlying driver."""
+        if self._database and 'database' not in kwargs:
+            kwargs["database"] = self._database
+        return self._driver.session(**kwargs)
+    
+    def close(self):
+        """Proxy method to close the underlying driver."""
+        self._driver.close()
+
 class DatabaseManager:
     """
     Manages the Neo4j database driver as a singleton to ensure only one
@@ -42,6 +60,7 @@ def __init__(self):
         self.neo4j_uri = os.getenv('NEO4J_URI')
         self.neo4j_username = os.getenv('NEO4J_USERNAME', 'neo4j')
         self.neo4j_password = os.getenv('NEO4J_PASSWORD')
+        self.neo4j_database = os.getenv('NEO4J_DATABASE') # Optional, if not set, will use default database configured in Neo4j
         self._initialized = True
 
     def get_driver(self) -> Driver:
@@ -53,7 +72,7 @@ def get_driver(self) -> Driver:
             ValueError: If Neo4j credentials are not set in environment variables.
 
         Returns:
-            The active Neo4j Driver instance.
+            The a wrapper for Neo4j Driver instance.
         """
         if self._driver is None:
             with self._lock:
@@ -100,7 +119,7 @@ def get_driver(self) -> Driver:
                             self._driver.close()
                         self._driver = None
                         raise
-        return self._driver
+        return Neo4jDriverWrapper(self._driver, database=self.neo4j_database)
 
     def close_driver(self):
         """Closes the Neo4j driver connection if it exists."""
@@ -116,7 +135,10 @@ def is_connected(self) -> bool:
         if self._driver is None:
             return False
         try:
-            with self._driver.session() as session:
+            session_kwargs = {}
+            if self.neo4j_database:
+                session_kwargs['database'] = self.neo4j_database
+            with self._driver.session(**session_kwargs) as session:
                 session.run("RETURN 1").consume()
             return True
         except Exception:
@@ -163,7 +185,7 @@ def validate_config(uri: str, username: str, password: str) -> Tuple[bool, Optio
         return True, None
 
     @staticmethod
-    def test_connection(uri: str, username: str, password: str) -> Tuple[bool, Optional[str]]:
+    def test_connection(uri: str, username: str, password: str, database: str=None) -> Tuple[bool, Optional[str]]:
         """
         Tests the Neo4j database connection.
         
@@ -206,7 +228,10 @@ def test_connection(uri: str, username: str, password: str) -> Tuple[bool, Optio
             # Now test Neo4j authentication
             driver = GraphDatabase.driver(uri, auth=(username, password))
             
-            with driver.session() as session:
+            session_kwargs = {}
+            if database:
+                session_kwargs['database'] = database # Pass database to session if provided
+            with driver.session(**session_kwargs) as session:
                 result = session.run("RETURN 'Connection successful' as status")
                 result.single()
             
