diff --git a/cognee/modules/graph/methods/__init__.py b/cognee/modules/graph/methods/__init__.py index e0752c6b0..920c66263 100644 --- a/cognee/modules/graph/methods/__init__.py +++ b/cognee/modules/graph/methods/__init__.py @@ -1 +1,7 @@ from .get_formatted_graph_data import get_formatted_graph_data + +from .delete_data_related_nodes import delete_data_related_nodes +from .delete_data_related_edges import delete_data_related_edges + +from .delete_dataset_related_nodes import delete_dataset_related_nodes +from .delete_dataset_related_edges import delete_dataset_related_edges diff --git a/cognee/modules/graph/methods/delete_data_related_edges.py b/cognee/modules/graph/methods/delete_data_related_edges.py new file mode 100644 index 000000000..374bde5f9 --- /dev/null +++ b/cognee/modules/graph/methods/delete_data_related_edges.py @@ -0,0 +1,14 @@ +from uuid import UUID +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from cognee.infrastructure.databases.relational import with_async_session +from cognee.modules.graph.models import Edge + + +@with_async_session +async def delete_data_related_edges(data_id: UUID, session: AsyncSession): + edges = (await session.scalars(select(Edge).where(Edge.data_id == data_id))).all() + + await session.execute(delete(Edge).where(Edge.id.in_([edge.id for edge in edges]))) + diff --git a/cognee/modules/graph/methods/delete_data_related_nodes.py b/cognee/modules/graph/methods/delete_data_related_nodes.py new file mode 100644 index 000000000..c6bb7c95d --- /dev/null +++ b/cognee/modules/graph/methods/delete_data_related_nodes.py @@ -0,0 +1,14 @@ +from uuid import UUID +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from cognee.infrastructure.databases.relational import with_async_session +from cognee.modules.graph.models import Node + + +@with_async_session +async def delete_data_related_nodes(data_id: UUID, session: AsyncSession): + nodes = (await session.scalars(select(Node).where(Node.data_id == data_id))).all() + + await session.execute(delete(Node).where(Node.id.in_([node.id for node in nodes]))) + diff --git a/cognee/modules/graph/methods/delete_dataset_related_edges.py b/cognee/modules/graph/methods/delete_dataset_related_edges.py new file mode 100644 index 000000000..8ac0c63f1 --- /dev/null +++ b/cognee/modules/graph/methods/delete_dataset_related_edges.py @@ -0,0 +1,14 @@ +from uuid import UUID +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from cognee.infrastructure.databases.relational import with_async_session +from cognee.modules.graph.models import Edge + + +@with_async_session +async def delete_dataset_related_edges(dataset_id: UUID, session: AsyncSession): + edges = (await session.scalars(select(Edge).where(Edge.dataset_id == dataset_id))).all() + + await session.execute(delete(Edge).where(Edge.id.in_([edge.id for edge in edges]))) + diff --git a/cognee/modules/graph/methods/delete_dataset_related_nodes.py b/cognee/modules/graph/methods/delete_dataset_related_nodes.py new file mode 100644 index 000000000..7e251189f --- /dev/null +++ b/cognee/modules/graph/methods/delete_dataset_related_nodes.py @@ -0,0 +1,14 @@ +from uuid import UUID +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from cognee.infrastructure.databases.relational import with_async_session +from cognee.modules.graph.models import Node + + +@with_async_session +async def delete_dataset_related_nodes(dataset_id: UUID, session: AsyncSession): + nodes = (await session.scalars(select(Node).where(Node.dataset_id == dataset_id))).all() + + await session.execute(delete(Node).where(Node.id.in_([node.id for node in nodes]))) + diff --git a/cognee/tests/unit/modules/graph/methods/test_delete_data_related_edges.py b/cognee/tests/unit/modules/graph/methods/test_delete_data_related_edges.py new file mode 100644 index 000000000..dca5c1558 --- /dev/null +++ b/cognee/tests/unit/modules/graph/methods/test_delete_data_related_edges.py @@ -0,0 +1,44 @@ +import pytest +from uuid import uuid4 +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from cognee.modules.graph.methods import delete_data_related_edges + + +class DummyScalarResult: + def __init__(self, items): + self._items = items + + def all(self): + return self._items + + +class FakeEdge: + def __init__(self, edge_id): + self.id = edge_id + + +@pytest.mark.asyncio +async def test_delete_data_related_edges_deletes_found_rows(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([FakeEdge(1), FakeEdge(2)])) + session.execute = AsyncMock() + + await delete_data_related_edges(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delete_data_related_edges_handles_empty_list(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([])) + session.execute = AsyncMock() + + await delete_data_related_edges(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() + diff --git a/cognee/tests/unit/modules/graph/methods/test_delete_data_related_nodes.py b/cognee/tests/unit/modules/graph/methods/test_delete_data_related_nodes.py new file mode 100644 index 000000000..44ef20476 --- /dev/null +++ b/cognee/tests/unit/modules/graph/methods/test_delete_data_related_nodes.py @@ -0,0 +1,44 @@ +import pytest +from uuid import uuid4 +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from cognee.modules.graph.methods import delete_data_related_nodes + + +class DummyScalarResult: + def __init__(self, items): + self._items = items + + def all(self): + return self._items + + +class FakeNode: + def __init__(self, node_id): + self.id = node_id + + +@pytest.mark.asyncio +async def test_delete_data_related_nodes_deletes_found_rows(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([FakeNode(1), FakeNode(2)])) + session.execute = AsyncMock() + + await delete_data_related_nodes(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delete_data_related_nodes_handles_empty_list(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([])) + session.execute = AsyncMock() + + await delete_data_related_nodes(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() + diff --git a/cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_edges.py b/cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_edges.py new file mode 100644 index 000000000..210a98bd5 --- /dev/null +++ b/cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_edges.py @@ -0,0 +1,44 @@ +import pytest +from uuid import uuid4 +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from cognee.modules.graph.methods import delete_dataset_related_edges + + +class DummyScalarResult: + def __init__(self, items): + self._items = items + + def all(self): + return self._items + + +class FakeEdge: + def __init__(self, edge_id): + self.id = edge_id + + +@pytest.mark.asyncio +async def test_delete_dataset_related_edges_deletes_found_rows(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([FakeEdge(1), FakeEdge(2)])) + session.execute = AsyncMock() + + await delete_dataset_related_edges(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delete_dataset_related_edges_handles_empty_list(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([])) + session.execute = AsyncMock() + + await delete_dataset_related_edges(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() + diff --git a/cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_nodes.py b/cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_nodes.py new file mode 100644 index 000000000..f52ad58ce --- /dev/null +++ b/cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_nodes.py @@ -0,0 +1,44 @@ +import pytest +from uuid import uuid4 +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from cognee.modules.graph.methods import delete_dataset_related_nodes + + +class DummyScalarResult: + def __init__(self, items): + self._items = items + + def all(self): + return self._items + + +class FakeNode: + def __init__(self, node_id): + self.id = node_id + + +@pytest.mark.asyncio +async def test_delete_dataset_related_nodes_deletes_found_rows(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([FakeNode(1), FakeNode(2)])) + session.execute = AsyncMock() + + await delete_dataset_related_nodes(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delete_dataset_related_nodes_handles_empty_list(): + session = SimpleNamespace() + session.scalars = AsyncMock(return_value=DummyScalarResult([])) + session.execute = AsyncMock() + + await delete_dataset_related_nodes(uuid4(), session=session) + + session.scalars.assert_awaited_once() + session.execute.assert_awaited_once() +