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

fix: import datasets directly in delete cli command

Browse source
This commit is contained in:
Boris Arzentar 2025-10-14 15:08:30 +02:00
parent 344b057d44
commit ed272776e2
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
3 changed files with 22 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
cognee/cli/commands/delete_command.py
View file

@ -35,7 +35,7 @@ Be careful with deletion operations as they are irreversible.
def execute(self, args: argparse.Namespace) -> None: def execute(self, args: argparse.Namespace) -> None:
try: try:
# Import cognee here to avoid circular imports # Import cognee here to avoid circular imports
import cognee from cognee.api.v1.datasets.datasets import datasets as cognee_datasets
# Validate arguments # Validate arguments
if not any( if not any(
@ -107,7 +107,7 @@ Be careful with deletion operations as they are irreversible.
raise CliCommandException( raise CliCommandException(
"No user ID provided for '--all' deletion. Please specify using --user-id param." "No user ID provided for '--all' deletion. Please specify using --user-id param."
) )
await cognee.datasets.delete_all(user_id=args.user_id) await cognee_datasets.delete_all(user_id=args.user_id)
elif hasattr(args, "dataset_name") or hasattr(args, "dataset_id"): elif hasattr(args, "dataset_name") or hasattr(args, "dataset_id"):
dataset_id = getattr(args, "dataset_id", None) dataset_id = getattr(args, "dataset_id", None)
@ -129,11 +129,11 @@ Be careful with deletion operations as they are irreversible.
"No user ID provided for deletion. Please specify using --user-id param." "No user ID provided for deletion. Please specify using --user-id param."
) )
await cognee.datasets.delete_dataset( await cognee_datasets.delete_dataset(
dataset_id=dataset_id, user_id=args.user_id dataset_id=dataset_id, user_id=args.user_id
) )
elif hasattr(args, "dataset_id") and hasattr(args, "data_id"): elif hasattr(args, "dataset_id") and hasattr(args, "data_id"):
await cognee.datasets.delete_data(args.dataset_id, args.data_id) await cognee_datasets.delete_data(args.dataset_id, args.data_id)
except Exception as e: except Exception as e:
raise CliCommandInnerException(f"Failed to delete: {str(e)}") from e raise CliCommandInnerException(f"Failed to delete: {str(e)}") from e

21
cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py
View file

@ -286,10 +286,11 @@ class TestDeleteCommand:
assert "force" in actions assert "force" in actions
@patch("cognee.modules.data.methods.get_deletion_counts.get_user") @patch("cognee.modules.data.methods.get_deletion_counts.get_user")
@patch("cognee.api.v1.datasets.datasets.datasets.delete_dataset")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")
@patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
def test_execute_delete_dataset_with_confirmation( def test_execute_delete_dataset_with_confirmation(
self, mock_asyncio_run, mock_confirm, get_user_mock self, mock_asyncio_run, mock_confirm, delete_dataset_mock, get_user_mock
): ):
"""Test execute delete dataset with user confirmation""" """Test execute delete dataset with user confirmation"""
expected_user_id = uuid4() expected_user_id = uuid4()
@ -297,20 +298,16 @@ class TestDeleteCommand:
get_user_mock.return_value = User(id=expected_user_id) get_user_mock.return_value = User(id=expected_user_id)
mock_cognee = MagicMock() command = DeleteCommand()
mock_cognee.datasets.delete_dataset = AsyncMock() args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
with patch.dict(sys.modules, {"cognee": mock_cognee}): mock_confirm.return_value = True
command = DeleteCommand()
args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
mock_confirm.return_value = True command.execute(args)
command.execute(args) delete_dataset_mock.assert_awaited_once_with(
dataset_id=expected_dataset_id, user_id=expected_user_id
mock_cognee.datasets.delete_dataset.assert_awaited_once_with( )
dataset_id=expected_dataset_id, user_id=expected_user_id
)
@patch("cognee.cli.commands.delete_command.get_deletion_counts") @patch("cognee.cli.commands.delete_command.get_deletion_counts")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")

21
cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py
View file

@ -381,26 +381,23 @@ class TestDeleteCommandEdgeCases:
"""Test edge cases for DeleteCommand""" """Test edge cases for DeleteCommand"""
@patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
@patch("cognee.api.v1.datasets.datasets.datasets.delete_all")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")
def test_delete_all_with_user_id(self, fmt_confirm_mock, async_run_mock): def test_delete_all_with_user_id(self, fmt_confirm_mock, delete_all_mock, async_run_mock):
"""Test delete command with both --all and --user-id""" """Test delete command with both --all and --user-id"""
fmt_confirm_mock.return_value = True fmt_confirm_mock.return_value = True
expected_user_id = uuid4() expected_user_id = uuid4()
mock_cognee = MagicMock() command = DeleteCommand()
mock_cognee.datasets.delete_all = AsyncMock() args = argparse.Namespace(
user_id=expected_user_id,
all=True,
)
with patch.dict(sys.modules, {"cognee": mock_cognee}): command.execute(args)
command = DeleteCommand()
args = argparse.Namespace(
user_id=expected_user_id,
all=True,
)
command.execute(args) delete_all_mock.assert_called_once_with(user_id=expected_user_id)
mock_cognee.datasets.delete_all.assert_called_once_with(user_id=expected_user_id)
@patch("cognee.cli.commands.delete_command.get_deletion_counts") @patch("cognee.cli.commands.delete_command.get_deletion_counts")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")