Skip to content

Commit

Permalink
add BackgroundJobExportMixin
Browse files Browse the repository at this point in the history
  • Loading branch information
sky-code committed Sep 30, 2024
1 parent 7a9ec54 commit d092726
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 10 deletions.
6 changes: 4 additions & 2 deletions example/winners/admin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,)
51 changes: 51 additions & 0 deletions import_export_celery/admin.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -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)
25 changes: 17 additions & 8 deletions import_export_celery/models/exportjob.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
# Copyright (C) 2019 o.s. Auto*Mat
import pickle
from base64 import b64decode

from django.utils import timezone
import json

Expand Down Expand Up @@ -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 [
Expand Down

0 comments on commit d092726

Please sign in to comment.