Repository: topoteretes/cognee
File: cognee/modules/data/methods/get_deletion_counts.py (Python)

Code before:
        if all_data:
            users_query = select(User)
            users = (await session.execute(users_query)).scalars().all()
            users_ids = [u.id for u in users]
            user_count = len(users_ids)
            dataset_count = (
                await session.execute(select(func.count()).select_from(Dataset))
            ).scalar_one()
            data_entry_count = 0
            for user_id in users_ids:
                datasets_query = select(Dataset).where(Dataset.owner_id == user_id)
                user_datasets = (await session.execute(datasets_query)).scalars().all()
                if dataset_count > 0:
                    dataset_ids = [d.id for d in user_datasets]
                    # Count all data entries across all of the user's datasets
                    data_count_query = (
                        select(func.count())
                        .select_from(DatasetData)
                        .where(DatasetData.dataset_id.in_(dataset_ids))
                    )
                    data_entry_count += (await session.execute(data_count_query)).scalar_one()
            return {
                "datasets": dataset_count,
                "data_entries": data_entry_count,
                "users": user_count,
            }

        # Placeholder for user_id logic
        if user_id:
            counts = {}
            # TODO: How to query a user by its id?
            user = await get_default_user()
            if user:
                counts["users"] = 1
                # Find all datasets owned by this user
                datasets_query = select(Dataset).where(Dataset.owner_id == user.id)
                user_datasets = (await session.execute(datasets_query)).scalars().all()
                dataset_count = len(user_datasets)
                counts["datasets"] = dataset_count
                if dataset_count > 0:
                    dataset_ids = [d.id for d in user_datasets]
                    # Count all data entries across all of the user's datasets
                    data_count_query = (
                        select(func.count())
                        .select_from(DatasetData)
                        .where(DatasetData.dataset_id.in_(dataset_ids))
                    )
                    data_entry_count = (await session.execute(data_count_query)).scalar_one()
                    counts["data_entries"] = data_entry_count
                else:
                    counts["data_entries"] = 0

Review comment (refactor):
Let's simplify this.

Since `all_data` means everything will be deleted, you can just get the total counts of users/datasets/data without doing things like querying for datasets for a given user.

Code after:
                raise CliCommandException(
                    f"No Dataset exists with the name {dataset_name}", error_code=1
                )

            # Count data entries linked to this dataset
            count_query = (
                select(func.count())
                .select_from(DatasetData)
                .where(DatasetData.dataset_id == dataset.id)
            )
            data_entry_count = (await session.execute(count_query)).scalar_one()
            counts.users = 1
            counts.datasets = 1
            counts.entries = data_entry_count
            return counts

        elif all_data:
            # Simplified logic: Get total counts directly from the tables.
            counts.datasets = (
                await session.execute(select(func.count()).select_from(Dataset))
            ).scalar_one()
            counts.entries = (
                await session.execute(select(func.count()).select_from(Data))
            ).scalar_one()
            counts.users = (
                await session.execute(select(func.count()).select_from(User))
            ).scalar_one()
            return counts

        # Placeholder for user_id logic
        elif user_id:
            user = None
            try:
                user_uuid = UUID(user_id)
                user = await get_user(user_uuid)
            except (ValueError, EntityNotFoundError):
                raise CliCommandException(f"No User exists with ID {user_id}", error_code=1)
            counts.users = 1
            # Find all datasets owned by this user
            datasets_query = select(Dataset).where(Dataset.owner_id == user.id)
            user_datasets = (await session.execute(datasets_query)).scalars().all()
            dataset_count = len(user_datasets)
            counts.datasets = dataset_count
            if dataset_count > 0:
                dataset_ids = [d.id for d in user_datasets]
                # Count all data entries across all of the user's datasets
                data_count_query = (
                    select(func.count())
                    .select_from(DatasetData)
                    .where(DatasetData.dataset_id.in_(dataset_ids))
                )