Repository: topoteretes/cognee
File: cognee/context_global_variables.py (Python)

Code before:
#       for different async tasks, threads and processes
vector_db_config = ContextVar("vector_db_config", default=None)
graph_db_config = ContextVar("graph_db_config", default=None)
session_user = ContextVar("session_user", default=None)

vector_dbs_with_multi_user_support = ["lancedb"]
graph_dbs_with_multi_user_support = ["kuzu"]


async def set_session_user_context_variable(user):
    session_user.set(user)


def check_multi_user_support():
    graph_db_config = get_graph_context_config()
    vector_db_config = get_vectordb_context_config()
    if (
        graph_db_config["graph_database_provider"] in graph_dbs_with_multi_user_support
        and vector_db_config["vector_db_provider"] in vector_dbs_with_multi_user_support
    ):
        return True
    else:
        return False


def check_backend_access_control_mode():
    backend_access_control = os.environ.get("ENABLE_BACKEND_ACCESS_CONTROL", None)
    if backend_access_control is None:
        # If backend access control is not defined in environment variables,
        # enable it by default if graph and vector DBs can support it, otherwise disable it
        multi_user_support = check_multi_user_support()
        if multi_user_support:
            return True
        else:
            return False
    elif backend_access_control.lower() == "true":
        # If enabled, ensure that the current graph and vector DBs can support it
        multi_user_support = check_multi_user_support()
        if not multi_user_support:
            raise EnvironmentError(
                "ENABLE_BACKEND_ACCESS_CONTROL is set to true but the current graph and/or vector databases do not support multi-user access control. Please use supported databases or disable backend access control."
            )
        else:
            return True
    else:
        # If explicitly disabled, return false
        return False


async def set_database_global_context_variables(dataset: Union[str, UUID], user_id: UUID):
    """

Review comment (suggestion):
Tip: Could you name it a bit more deterministic: like `backend_access_control_enabled()`. It will help us better understand the codebase

Code after:
#       for different async tasks, threads and processes
vector_db_config = ContextVar("vector_db_config", default=None)
graph_db_config = ContextVar("graph_db_config", default=None)
session_user = ContextVar("session_user", default=None)

vector_dbs_with_multi_user_support = ["lancedb"]
graph_dbs_with_multi_user_support = ["kuzu"]


async def set_session_user_context_variable(user):
    session_user.set(user)


def multi_user_support_possible():
    graph_db_config = get_graph_context_config()
    vector_db_config = get_vectordb_context_config()
    return (
        graph_db_config["graph_database_provider"] in graph_dbs_with_multi_user_support
        and vector_db_config["vector_db_provider"] in vector_dbs_with_multi_user_support
    )


def backend_access_control_enabled():
    backend_access_control = os.environ.get("ENABLE_BACKEND_ACCESS_CONTROL", None)
    if backend_access_control is None:
        # If backend access control is not defined in environment variables,
        # enable it by default if graph and vector DBs can support it, otherwise disable it
        return multi_user_support_possible()
    elif backend_access_control.lower() == "true":
        # If enabled, ensure that the current graph and vector DBs can support it
        multi_user_support = multi_user_support_possible()
        if not multi_user_support:
            raise EnvironmentError(
                "ENABLE_BACKEND_ACCESS_CONTROL is set to true but the current graph and/or vector databases do not support multi-user access control. Please use supported databases or disable backend access control."
            )
        return True
    return False


async def set_database_global_context_variables(dataset: Union[str, UUID], user_id: UUID):
    """
    If backend access control is enabled this function will ensure all datasets have their own databases,
    access to which will be enforced by given permissions.
    Database name will be determined by dataset_id and LanceDB and KuzuDB use will be enforced.

    Note: This is only currently supported by the following databases:
          Relational: SQLite, Postgres
          Vector: LanceDB
          Graph: KuzuDB

    Args: