From cc1aac9eaf06f32b493bc1ed3dd26aae0346109c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Dom=C3=ADnguez?= Date: Thu, 18 Dec 2025 23:35:08 -0300 Subject: [PATCH] added optional db param neo4j --- .env.example | 1 + blarify/cli/commands/create.py | 2 ++ blarify/repositories/graph_db_manager/neo4j_manager.py | 10 ++++++---- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.env.example b/.env.example index 54c4dab3..3a3d1ff2 100644 --- a/.env.example +++ b/.env.example @@ -5,6 +5,7 @@ NEO4J_URI=bolt://localhost:7687 NEO4J_USERNAME=neo4j NEO4J_PASSWORD=your_password_here +NEO4J_DATABASE= # Optional: specify database name (defaults to Neo4j's default database) # Project Configuration ROOT_PATH=/Users/pepemanu/Desktop/Trabajo/Blar/Dev/blarify diff --git a/blarify/cli/commands/create.py b/blarify/cli/commands/create.py index c442d22a..7e715fe0 100644 --- a/blarify/cli/commands/create.py +++ b/blarify/cli/commands/create.py @@ -266,6 +266,7 @@ def add_arguments(parser: ArgumentParser) -> None: ) parser.add_argument("--neo4j-username", help="Neo4j username") parser.add_argument("--neo4j-password", help="Neo4j password") + parser.add_argument("--neo4j-database", help="Neo4j database name (optional, defaults to Neo4j's default database)") # Graph building options parser.add_argument( @@ -385,6 +386,7 @@ def execute(args: Namespace) -> int: uri=args.neo4j_uri, user=args.neo4j_username, password=args.neo4j_password, + database=args.neo4j_database, repo_id=repo_id, entity_id=args.entity_id, ) diff --git a/blarify/repositories/graph_db_manager/neo4j_manager.py b/blarify/repositories/graph_db_manager/neo4j_manager.py index 4842f5cb..2701c22f 100644 --- a/blarify/repositories/graph_db_manager/neo4j_manager.py +++ b/blarify/repositories/graph_db_manager/neo4j_manager.py @@ -34,11 +34,13 @@ def __init__( uri: Optional[str] = None, user: Optional[str] = None, password: Optional[str] = None, + database: Optional[str] = None, max_connections: int = 50, ): uri = uri or os.getenv("NEO4J_URI") user = user or os.getenv("NEO4J_USERNAME") password = password or os.getenv("NEO4J_PASSWORD") + self.database = database or os.getenv("NEO4J_DATABASE") if not uri or not user or not password: raise ValueError("Missing required Neo4j connection parameters") @@ -83,7 +85,7 @@ def create_nodes(self, nodeList: List[Any]): if self.repo_id is None: raise ValueError("repo_id is required for creating nodes. Cannot create nodes with entity-wide scope.") - with self.driver.session() as session: + with self.driver.session(database=self.database) as session: session.execute_write( self._create_nodes_txn, nodeList, @@ -102,7 +104,7 @@ def create_edges(self, edgesList: List[Any]): total_edges = len(edgesList) logger.info(f"Creating {total_edges} edges in batches of {batch_size}") - with self.driver.session() as session: + with self.driver.session(database=self.database) as session: for i in range(0, total_edges, batch_size): batch = edgesList[i:i + batch_size] logger.info(f"Processing edges batch {i // batch_size + 1}/{(total_edges + batch_size - 1) // batch_size} ({i}/{total_edges})") @@ -205,7 +207,7 @@ def detatch_delete_nodes_with_path(self, path: str): if self.repo_id is None: raise ValueError("repo_id is required for deleting nodes. Cannot delete nodes with entity-wide scope.") - with self.driver.session() as session: + with self.driver.session(database=self.database) as session: result = session.run( """ MATCH (n {path: $path, repoId: $repo_id, entityId: $entity_id}) @@ -243,7 +245,7 @@ def query( parameters["entity_id"] = self.entity_id try: - with self.driver.session() as session: + with self.driver.session(database=self.database) as session: result = [] if transaction: result = session.execute_write(