Repository: topoteretes/cognee
File: cognee/infrastructure/databases/utils/get_or_create_dataset_database.py (Python)

Code before:
    db_engine = get_relational_engine()

    dataset_id = await get_unique_dataset_id(dataset, user)

    vector_db_name = f"{dataset_id}.db"
    graph_db_name = f"{dataset_id}.pkl"

    vector_config = get_vectordb_config()
    graph_config = get_graph_config()

    async with db_engine.get_async_session() as session:
        # Create dataset if it doesn't exist
        if isinstance(dataset, str):
            dataset = await create_dataset(dataset, user, session)

        # Try to fetch an existing row first
        stmt = select(DatasetDatabase).where(
            DatasetDatabase.owner_id == user.id,
            DatasetDatabase.dataset_id == dataset_id,
        )
        existing: DatasetDatabase = await session.scalar(stmt)
        if existing:
            return existing

        # Check if we support multi-user for this provider. If not, use default
        if graph_config.graph_database_provider in graph_dbs_with_multi_user_support:
            graph_provider = graph_config.graph_database_provider
            graph_url = graph_config.graph_database_url
            graph_key = graph_config.graph_database_key
        else:
            graph_provider = default_graph_db_provider
            graph_url = default_graph_db_url
            graph_key = default_graph_db_key

        if vector_config.vector_db_provider in vector_dbs_with_multi_user_support:
            vector_provider = vector_config.vector_db_provider
            vector_url = vector_config.vector_db_url
            vector_key = vector_config.vector_db_key
        else:
            vector_provider = default_vector_db_provider
            vector_url = default_vector_db_url
            vector_key = default_vector_db_key

        # If there are no existing rows build a new row
        record = DatasetDatabase(
            owner_id=user.id,
            dataset_id=dataset_id,
            vector_database_name=vector_db_name,
            graph_database_name=graph_db_name,
            vector_database_provider=vector_provider,
            graph_database_provider=graph_provider,

Review comment (suggestion):
Cognee will not work if the DBs don't support mulit-user mode and are trying to use Cognee with multi-user mode so better raise an exception 

Code after:
        graph_db_name = f"{dataset_id}.pkl"
    else:
        graph_db_name = f"{dataset_id}"

    if vector_config.vector_db_provider == "lancedb":
        vector_db_name = f"{dataset_id}.lance.db"
    else:
        vector_db_name = f"{dataset_id}"

    base_config = get_base_config()
    databases_directory_path = os.path.join(
        base_config.system_root_directory, "databases", str(user.id)
    )

    # Determine vector database URL
    if vector_config.vector_db_provider == "lancedb":
        vector_db_url = os.path.join(databases_directory_path, vector_config.vector_db_name)
    else:
        vector_db_url = vector_config.vector_database_url

    # Determine graph database URL

    async with db_engine.get_async_session() as session:
        # Create dataset if it doesn't exist
        if isinstance(dataset, str):
            dataset = await create_dataset(dataset, user, session)

        # Try to fetch an existing row first
        stmt = select(DatasetDatabase).where(
            DatasetDatabase.owner_id == user.id,
            DatasetDatabase.dataset_id == dataset_id,
        )
        existing: DatasetDatabase = await session.scalar(stmt)
        if existing:
            return existing

        # If there are no existing rows build a new row
        record = DatasetDatabase(
            owner_id=user.id,
            dataset_id=dataset_id,
            vector_database_name=vector_db_name,
            graph_database_name=graph_db_name,
            vector_database_provider=vector_config.vector_db_provider,
            graph_database_provider=graph_config.graph_database_provider,
            vector_database_url=vector_db_url,
            graph_database_url=graph_config.graph_database_url,
            vector_database_key=vector_config.vector_db_key,
            graph_database_key=graph_config.graph_database_key,
        )

        try: