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

Code before:
    graph_db_name = f"{dataset_id}.pkl"

    if graph_config.graph_database_provider in HYBRID_DBS:
        vector_db_name = graph_db_name
    else:
        if vector_config.vector_db_provider == "lancedb":
            vector_db_name = f"{dataset_id}.lance.db"
        else:
            vector_db_name = f"{dataset_id}.db"

    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 not in GRAPH_DBS_WITH_MULTI_USER_SUPPORT:
            raise EnvironmentError(
                f"Multi-user is currently not supported for the graph database provider: {graph_config.graph_database_provider}. "
                f"Supported providers are: {', '.join(GRAPH_DBS_WITH_MULTI_USER_SUPPORT)}. Either use one of these"
                f"providers, or disable BACKEND_ACCESS_CONTROL"
            )
        if vector_config.vector_db_provider not in VECTOR_DBS_WITH_MULTI_USER_SUPPORT:
            raise EnvironmentError(
                f"Multi-user is currently not supported for the vector database provider: {vector_config.vector_db_provider}. "
                f"Supported providers are: {', '.join(VECTOR_DBS_WITH_MULTI_USER_SUPPORT)}. Either use one of these"
                f"providers, or disable BACKEND_ACCESS_CONTROL"
            )

        # 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_config.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,
        )

Review comment (refactor):
I think we can remove this, new changes with default multi-user mode should cover these errors, but let's verfiy

Code after:
    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:
            session.add(record)