diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py index 31e3fa67d..1ba96a232 100644 --- a/cognee/api/v1/datasets/routers/get_datasets_router.py +++ b/cognee/api/v1/datasets/routers/get_datasets_router.py @@ -76,7 +76,7 @@ def get_datasets_router() -> APIRouter: message=f"Dataset ({dataset_id}) not found." ) - data = await get_data(data_id) + data = await get_data(user.id, data_id) if data is None: raise EntityNotFoundError( @@ -141,6 +141,7 @@ def get_datasets_router() -> APIRouter: @router.get("/{dataset_id}/data/{data_id}/raw", response_class=FileResponse) async def get_raw_data(dataset_id: str, data_id: str, user: User = Depends(get_authenticated_user)): + from cognee.modules.data.methods import get_data from cognee.modules.data.methods import get_dataset, get_dataset_data dataset = await get_dataset(user.id, dataset_id) @@ -164,7 +165,10 @@ def get_datasets_router() -> APIRouter: if len(matching_data) == 0: raise EntityNotFoundError(message= f"Data ({data_id}) not found in dataset ({dataset_id}).") - data = matching_data[0] + data = await get_data(user.id, data_id) + + if data is None: + raise EntityNotFoundError(message=f"Data ({data_id}) not found in dataset ({dataset_id}).") return data.raw_data_location diff --git a/cognee/modules/data/methods/get_data.py b/cognee/modules/data/methods/get_data.py index b07401463..9161e3df2 100644 --- a/cognee/modules/data/methods/get_data.py +++ b/cognee/modules/data/methods/get_data.py @@ -3,10 +3,11 @@ from typing import Optional from cognee.infrastructure.databases.relational import get_relational_engine from ..models import Data -async def get_data(data_id: UUID) -> Optional[Data]: +async def get_data(user_id: UUID, data_id: UUID) -> Optional[Data]: """Retrieve data by ID. Args: + user_id (UUID): user ID data_id (UUID): ID of the data to retrieve Returns: @@ -17,4 +18,7 @@ async def get_data(data_id: UUID) -> Optional[Data]: async with db_engine.get_async_session() as session: data = await session.get(Data, data_id) + if data and data.owner_id != user_id: + return None + return data \ No newline at end of file