gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: Add user verification for accessing data

Browse source 
Verify user has access to data before returning it

Feature COG-656
This commit is contained in:
Igor Ilic 2024-12-13 13:54:45 +01:00
parent 1180839469
commit 43187e4d63
2 changed files with 11 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
cognee/api/v1/datasets/routers/get_datasets_router.py
View file

@ -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

6
cognee/modules/data/methods/get_data.py
View file

@ -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