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

feat: Enable database deletion when dataset is deleted

Browse source
This commit is contained in:
Igor Ilic 2025-12-12 17:43:31 +01:00
parent cb16dc0a7e
commit 47c02781e5
2 changed files with 28 additions and 4 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

@ -208,21 +208,21 @@ def get_datasets_router() -> APIRouter:
},
)
from cognee.modules.data.methods import get_dataset, delete_dataset
from cognee.modules.data.methods import delete_dataset
dataset = await get_dataset(user.id, dataset_id)
dataset = await get_authorized_existing_datasets([dataset_id], "delete", user)
if dataset is None:
raise DatasetNotFoundError(message=f"Dataset ({str(dataset_id)}) not found.")
await delete_dataset(dataset)
await delete_dataset(dataset[0])
@router.delete(
"/{dataset_id}/data/{data_id}",
response_model=None,
responses={404: {"model": ErrorResponseDTO}},
)
async def delete_data(
async def delete_data( # TODO: Is this endpoint needed? It seems redundant with data deletion elsewhere
dataset_id: UUID, data_id: UUID, user: User = Depends(get_authenticated_user)
):
"""

24
cognee/modules/data/methods/delete_dataset.py
View file

@ -1,8 +1,32 @@
from cognee.modules.users.models import DatasetDatabase
from sqlalchemy import select
from cognee.modules.data.models import Dataset
from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import (
get_vector_dataset_database_handler,
)
from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import (
get_graph_dataset_database_handler,
)
from cognee.infrastructure.databases.relational import get_relational_engine
async def delete_dataset(dataset: Dataset):
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
stmt = select(DatasetDatabase).where(
DatasetDatabase.dataset_id == dataset.id,
)
dataset_database: DatasetDatabase = await session.scalar(stmt)
if dataset_database:
graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database)
vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database)
await graph_dataset_database_handler["handler_instance"].delete_dataset(
dataset_database
)
await vector_dataset_database_handler["handler_instance"].delete_dataset(
dataset_database
)
return await db_engine.delete_entity_by_id(dataset.__tablename__, dataset.id)