From d0927265946dc249e03629240c1713091d3ebf1c Mon Sep 17 00:00:00 2001
From: sky-code <skycodes@gmail.com>
Date: Mon, 30 Sep 2024 01:02:38 -0400
Subject: [PATCH] add BackgroundJobExportMixin

---
 example/winners/admin.py                 |  6 ++-
 import_export_celery/admin.py            | 51 ++++++++++++++++++++++++
 import_export_celery/models/exportjob.py | 25 ++++++++----
 3 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/example/winners/admin.py b/example/winners/admin.py
index 9e7a117..2f20ab5 100644
--- a/example/winners/admin.py
+++ b/example/winners/admin.py
@@ -2,13 +2,15 @@
 from django.contrib import admin
 
 from import_export.admin import ImportExportMixin
+
+from import_export_celery.admin import BackgroundJobExportMixin
 from import_export_celery.admin_actions import create_export_job_action
 
 from . import models
 
-
 @admin.register(models.Winner)
-class WinnerAdmin(ImportExportMixin, admin.ModelAdmin):
+class WinnerAdmin(BackgroundJobExportMixin, ImportExportMixin, admin.ModelAdmin):
     list_display = ("name",)
 
     actions = (create_export_job_action,)
+    resource_classes = (models.WinnersResource,)
diff --git a/import_export_celery/admin.py b/import_export_celery/admin.py
index 3e1ea08..787fa17 100644
--- a/import_export_celery/admin.py
+++ b/import_export_celery/admin.py
@@ -1,11 +1,20 @@
 # Copyright (C) 2019 o.s. Auto*Mat
+import json
+import pickle
+from base64 import b64encode
+
+import django
 from django import forms
 from django.conf import settings
 from django.contrib import admin
 from django.core.cache import cache
+from django.shortcuts import redirect
 from django.utils.translation import gettext_lazy as _
+from import_export.admin import ExportMixin
+from import_export.forms import SelectableFieldsExportForm
 
 from . import admin_actions, models
+from .models import ExportJob
 
 
 class JobWithStatusMixin:
@@ -111,3 +120,45 @@ def has_add_permission(self, request, obj=None):
         return False
 
     actions = (admin_actions.run_export_job_action,)
+
+
+class BackgroundJobSelectableFieldsExportForm(SelectableFieldsExportForm):
+    background_job = forms.BooleanField(
+        required=False,
+        initial=False,
+        label="Run export in background, will export all fields",
+    )
+
+
+class BackgroundJobExportMixin(ExportMixin):
+    export_form_class = BackgroundJobSelectableFieldsExportForm
+
+    def _do_file_export(self, file_format, request, queryset, export_form=None):
+        if not export_form.cleaned_data.get("background_job"):
+            return super()._do_file_export(file_format, request, queryset, export_form=export_form)
+        base64_pickle_qs = b64encode(pickle.dumps(queryset.query)).decode("ascii")
+        export_class = self.choose_export_resource_class(export_form, request)
+        model_class = export_class.Meta.model
+        resource_classes = model_class.export_resource_classes()
+        resource = None
+        for resource_key, resource_choice in resource_classes.items():
+            _resource_label, resource_class = resource_choice
+            if resource_class == export_class:
+                resource = resource_key
+                break
+        job = ExportJob(
+            app_label=self.model._meta.app_label,
+            model=self.model._meta.model_name,
+            site_of_origin=request.scheme + "://" + request.get_host(),
+            queryset=json.dumps(
+                {
+                    "queryString": request.GET.urlencode(),
+                    "djangoVersion": django.get_version(),
+                    "query": base64_pickle_qs,
+                }
+            ),
+            resource=resource,
+            format=file_format.CONTENT_TYPE,
+        )
+        job.save()
+        return redirect(f"admin:{job._meta.app_label}_{job._meta.model_name}_change", job.id)
diff --git a/import_export_celery/models/exportjob.py b/import_export_celery/models/exportjob.py
index 66f5691..04cd74f 100644
--- a/import_export_celery/models/exportjob.py
+++ b/import_export_celery/models/exportjob.py
@@ -1,4 +1,7 @@
 # Copyright (C) 2019 o.s. Auto*Mat
+import pickle
+from base64 import b64decode
+
 from django.utils import timezone
 import json
 
@@ -104,14 +107,20 @@ def get_content_type(self):
         return self._content_type
 
     def get_queryset(self):
-        pks = json.loads(self.queryset)
-        # If customised queryset for the model exists
-        # then it'll apply filter on that otherwise it'll
-        # apply filter directly on the model.
-        resource_class = self.get_resource_class()
-        if hasattr(resource_class, "get_export_queryset"):
-            return resource_class().get_export_queryset().filter(pk__in=pks)
-        return self.get_content_type().model_class().objects.filter(pk__in=pks)
+        queryset_spec = json.loads(self.queryset)
+        if isinstance(queryset_spec, list):
+            # If customised queryset for the model exists
+            # then it'll apply filter on that otherwise it'll
+            # apply filter directly on the model.
+            resource_class = self.get_resource_class()
+            if hasattr(resource_class, "get_export_queryset"):
+                return resource_class().get_export_queryset().filter(pk__in=queryset_spec)
+            return self.get_content_type().model_class().objects.filter(pk__in=queryset_spec)
+        elif isinstance(queryset_spec, dict):
+            query = pickle.loads(b64decode(queryset_spec["query"]))  # noqa: S301
+            queryset = self.get_content_type().model_class().objects.all()
+            queryset.query = query
+            return queryset
 
     def get_resource_choices(self):
         return [