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

added minimal delete for relational

Browse source
This commit is contained in:
vasilije 2026-01-06 11:53:57 +01:00
parent 34c6652939
commit 97c9cf2d8a
9 changed files with 238 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cognee/modules/graph/methods/__init__.py
View file

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

14
cognee/modules/graph/methods/delete_data_related_edges.py Normal file
View file

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

14
cognee/modules/graph/methods/delete_data_related_nodes.py Normal file
View file

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

14
cognee/modules/graph/methods/delete_dataset_related_edges.py Normal file
View file

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

14
cognee/modules/graph/methods/delete_dataset_related_nodes.py Normal file
View file

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

44
cognee/tests/unit/modules/graph/methods/test_delete_data_related_edges.py Normal file
View file

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

44
cognee/tests/unit/modules/graph/methods/test_delete_data_related_nodes.py Normal file
View file

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

44
cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_edges.py Normal file
View file

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

44
cognee/tests/unit/modules/graph/methods/test_delete_dataset_related_nodes.py Normal file
View file

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