diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py index b8bd4c617..9e9942818 100644 --- a/cognee/api/v1/datasets/routers/get_datasets_router.py +++ b/cognee/api/v1/datasets/routers/get_datasets_router.py @@ -52,6 +52,31 @@ def get_datasets_router(): await delete_dataset(dataset) + @router.delete("/{dataset_id}/data/{data_id}", response_model=None, responses={404: {"model": ErrorResponseDTO}}) + async def delete_data(dataset_id: str, data_id: str, user: User = Depends(get_authenticated_user)): + from cognee.modules.data.methods import get_data, delete_data + from cognee.modules.data.methods import get_dataset + + # Check if user has permission to access dataset and data by trying to get the dataset + dataset = await get_dataset(user.id, dataset_id) + + #TODO: Handle situation differently if user doesn't have permission to access data? + if dataset is None: + raise HTTPException( + status_code=404, + detail=f"Dataset ({dataset_id}) not found." + ) + + data = await get_data(user.id, data_id) + + if data is None: + raise HTTPException( + status_code=404, + detail=f"Dataset ({data_id}) not found." + ) + + await delete_data(data) + @router.get("/{dataset_id}/graph", response_model=str) async def get_dataset_graph(dataset_id: str, user: User = Depends(get_authenticated_user)): from cognee.shared.utils import render_graph diff --git a/cognee/modules/data/methods/__init__.py b/cognee/modules/data/methods/__init__.py index a904060f5..34f943359 100644 --- a/cognee/modules/data/methods/__init__.py +++ b/cognee/modules/data/methods/__init__.py @@ -6,6 +6,8 @@ from .get_dataset import get_dataset from .get_datasets import get_datasets from .get_datasets_by_name import get_datasets_by_name from .get_dataset_data import get_dataset_data +from .get_data import get_data # Delete from .delete_dataset import delete_dataset +from .delete_data import delete_data \ No newline at end of file diff --git a/cognee/modules/data/methods/delete_data.py b/cognee/modules/data/methods/delete_data.py new file mode 100644 index 000000000..4fe437d92 --- /dev/null +++ b/cognee/modules/data/methods/delete_data.py @@ -0,0 +1,8 @@ +from cognee.modules.data.models import Data +from cognee.infrastructure.databases.relational import get_relational_engine + + +async def delete_data(data: Data): + db_engine = get_relational_engine() + + return await db_engine.delete_data_by_id(data.__tablename__, data.id) diff --git a/cognee/modules/data/methods/delete_dataset.py b/cognee/modules/data/methods/delete_dataset.py index b0fe96c42..c2205144d 100644 --- a/cognee/modules/data/methods/delete_dataset.py +++ b/cognee/modules/data/methods/delete_dataset.py @@ -4,4 +4,4 @@ from cognee.infrastructure.databases.relational import get_relational_engine async def delete_dataset(dataset: Dataset): db_engine = get_relational_engine() - return await db_engine.delete_table(dataset.id) + return await db_engine.delete_data_by_id(dataset.__tablename__, dataset.id) diff --git a/cognee/modules/data/methods/get_data.py b/cognee/modules/data/methods/get_data.py new file mode 100644 index 000000000..dac3eb202 --- /dev/null +++ b/cognee/modules/data/methods/get_data.py @@ -0,0 +1,12 @@ +from uuid import UUID +from typing import Optional +from cognee.infrastructure.databases.relational import get_relational_engine +from ..models import Data + +async def get_data(user_id: UUID, data_id: UUID) -> Optional[Data]: + db_engine = get_relational_engine() + + async with db_engine.get_async_session() as session: + data = await session.get(Data, data_id) + + return data \ No newline at end of file