refactor: change neo4j_aura to neo4j_aura_dev

This commit is contained in:
Igor Ilic 2025-12-11 13:05:23 +01:00
parent f1c5b9a55f
commit 0a1ed79340
2 changed files with 17 additions and 7 deletions

View file

@ -1,5 +1,5 @@
from cognee.infrastructure.databases.graph.neo4j_driver.Neo4jAuraDatasetDatabaseHandler import (
Neo4jAuraDatasetDatabaseHandler,
from cognee.infrastructure.databases.graph.neo4j_driver.Neo4jAuraDevDatasetDatabaseHandler import (
Neo4jAuraDevDatasetDatabaseHandler,
)
from cognee.infrastructure.databases.vector.lancedb.LanceDBDatasetDatabaseHandler import (
LanceDBDatasetDatabaseHandler,
@ -9,8 +9,8 @@ from cognee.infrastructure.databases.graph.kuzu.KuzuDatasetDatabaseHandler impor
)
supported_dataset_database_handlers = {
"neo4j_aura": {
"handler_instance": Neo4jAuraDatasetDatabaseHandler,
"neo4j_aura_dev": {
"handler_instance": Neo4jAuraDevDatasetDatabaseHandler,
"handler_provider": "neo4j",
},
"lancedb": {"handler_instance": LanceDBDatasetDatabaseHandler, "handler_provider": "lancedb"},

View file

@ -12,9 +12,18 @@ from cognee.modules.users.models import User, DatasetDatabase
from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
class Neo4jAuraDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
"""
Handler for interacting with Neo4j Aura Dataset databases.
Handler for a quick development PoC integration of Cognee multi-user and permission mode with Neo4j Aura databases.
This handler creates a new Neo4j Aura instance for each Cognee dataset created.
Improvements needed to be production ready:
- Secret management for client credentials, currently secrets are encrypted and stored in the Cognee relational database,
a secret manager or a similar system should be used instead.
Quality of life improvements:
- Allow configuration of different Neo4j Aura plans and regions.
- Requests should be made async, currently a blocking requests library is used.
"""
@classmethod
@ -34,7 +43,7 @@ class Neo4jAuraDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
if graph_config.graph_database_provider != "neo4j":
raise ValueError(
"Neo4jAuraDatasetDatabaseHandler can only be used with Neo4j graph database provider."
"Neo4jAuraDevDatasetDatabaseHandler can only be used with Neo4j graph database provider."
)
graph_db_name = f"{dataset_id}"
@ -134,6 +143,7 @@ class Neo4jAuraDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
) -> DatasetDatabase:
"""
Resolve and decrypt connection info for the Neo4j dataset database.
In this case, decrypt the password stored in the database.
Args:
dataset_database: DatasetDatabase instance containing encrypted connection info.