export_users.py

import csv

from django.contrib.auth import get_user_model
from django.core.exceptions import ObjectDoesNotExist
from django.core.management.base import BaseCommand

from ... import settings as app_settings

User = get_user_model()


class Command(BaseCommand):
    help = "Exports user data to a CSV file"

    def add_arguments(self, parser):
        parser.add_argument(
            "--exclude-fields",
            dest="exclude_fields",
            default="",
            help="Comma-separated list of fields to exclude from export",
        )
        parser.add_argument(
            "--filename",
            dest="filename",
            default="openwisp_exported_users.csv",
            help=(
                "Filename for the exported CSV, defaults to"
                ' "openwisp_exported_users.csv"'
            ),
        )

    def handle(self, *args, **options):
        fields = app_settings.EXPORT_USERS_COMMAND_CONFIG.get("fields", []).copy()
        # Get the fields to be excluded from the command-line argument
        exclude_fields = options.get("exclude_fields").split(",")
        # Remove excluded fields from the export fields
        fields = [field for field in fields if field not in exclude_fields]
        # Fetch all user data in a single query using select_related for related models
        queryset = User.objects.select_related(
            *app_settings.EXPORT_USERS_COMMAND_CONFIG.get("select_related", []),
        ).order_by("date_joined")

        # Prepare a CSV writer
        filename = options.get("filename")
        csv_file = open(filename, "w", newline="")
        csv_writer = csv.writer(csv_file)

        # Write header row
        csv_writer.writerow(fields)

        # Write data rows
        for user in queryset.iterator():
            data_row = []
            for field in fields:
                # Extract the value from related models
                if "." in field:
                    related_model, related_field = field.split(".")
                    try:
                        related_value = getattr(
                            getattr(user, related_model), related_field
                        )
                    except ObjectDoesNotExist:
                        data_row.append("")
                    else:
                        data_row.append(related_value)
                elif field == "organizations":
                    organizations = []
                    for org_id, user_perm in user.organizations_dict.items():
                        organizations.append(f'({org_id},{user_perm["is_admin"]})')
                    data_row.append("\n".join(organizations))
                else:
                    data_row.append(getattr(user, field))
            csv_writer.writerow(data_row)

        # Close the CSV file
        csv_file.close()
        self.stdout.write(
            self.style.SUCCESS(f"User data exported successfully to {filename}!")
        )