40 lines
1.5 KiB
Python
40 lines
1.5 KiB
Python
from uuid import UUID
|
|
from typing import Union
|
|
|
|
from cognee.modules.users.methods import get_default_user
|
|
from cognee.modules.users.models import User
|
|
from cognee.modules.data.exceptions import DatasetNotFoundError
|
|
from cognee.modules.data.methods import (
|
|
get_authorized_existing_datasets,
|
|
load_or_create_datasets,
|
|
check_dataset_name,
|
|
)
|
|
|
|
|
|
async def authorized_user_datasets(user: User, datasets: Union[str, list[str], list[UUID]]):
|
|
# If no user is provided use default user
|
|
if user is None:
|
|
user = await get_default_user()
|
|
|
|
# Convert datasets to list
|
|
if isinstance(datasets, str) or isinstance(datasets, UUID):
|
|
datasets = [datasets]
|
|
|
|
# Get datasets user wants write permissions for (verify user has permissions if datasets are provided as well)
|
|
# NOTE: If a user wants to write to a dataset he does not own it must be provided through UUID
|
|
existing_datasets = await get_authorized_existing_datasets(datasets, "write", user)
|
|
|
|
if not datasets:
|
|
# Get datasets from database if none sent.
|
|
authorized_datasets = existing_datasets
|
|
else:
|
|
# If dataset matches an existing Dataset (by name or id), reuse it. Otherwise, create a new Dataset.
|
|
authorized_datasets = await load_or_create_datasets(datasets, existing_datasets, user)
|
|
|
|
if not authorized_datasets:
|
|
raise DatasetNotFoundError("There are no datasets to work with.")
|
|
|
|
for dataset in authorized_datasets:
|
|
check_dataset_name(dataset.name)
|
|
|
|
return user, authorized_datasets
|