Coverage for src/backend/InvenTree/data_exporter/mixins.py: 23%
201 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 17:47 +0000
1"""Mixin classes for the exporter app."""
3from collections import OrderedDict
4from typing import Any
6from django.core.exceptions import ValidationError
7from django.core.files.base import ContentFile
8from django.utils.translation import gettext_lazy as _
10import structlog
11import tablib
12from rest_framework import serializers
13from rest_framework.response import Response
14from taggit.serializers import TagListSerializerField
16import data_exporter.serializers
17import data_exporter.tasks
18import InvenTree.exceptions
19from common.models import DataOutput
20from InvenTree.helpers import str2bool
21from InvenTree.tasks import offload_task
22from plugin import PluginMixinEnum, registry
24logger = structlog.get_logger('inventree')
27class DataExportSerializerMixin:
28 """Mixin class for adding data export functionality to a DRF serializer.
30 Provides generic functionality to take the output of a serializer and export it to a file.
32 Attributes:
33 export_only_fields: List of field names which are only used during data export
34 export_exclude_fields: List of field names which are excluded during data export
35 export_child_fields: List of child fields which are exported (using dot notation)
36 """
38 export_only_fields = []
39 export_exclude_fields = []
40 export_child_fields = []
42 def get_export_only_fields(self, **kwargs) -> list:
43 """Return the list of field names which are only used during data export."""
44 return self.export_only_fields
46 def get_export_exclude_fields(self, **kwargs) -> list:
47 """Return the list of field names which are excluded during data export."""
48 return self.export_exclude_fields
50 def __init__(self, *args, **kwargs):
51 """Initialise the DataExportSerializerMixin.
53 Determine if the serializer is being used for data export,
54 and if so, adjust the serializer fields accordingly.
55 """
56 self._exporting_data = exporting = kwargs.pop('exporting', False)
58 super().__init__(*args, **kwargs)
60 # Cache the request object
61 self.request = self.context.get('request')
63 if exporting: 63 ↛ 65line 63 didn't jump to line 65 because the condition on line 63 was never true
64 # Exclude fields which are not required for data export
65 for field in self.get_export_exclude_fields(**kwargs):
66 self.fields.pop(field, None)
67 else:
68 # Exclude fields which are only used for data export
69 for field in self.get_export_only_fields(**kwargs):
70 self.fields.pop(field, None)
72 def get_exportable_fields(self) -> dict:
73 """Return a dict of fields which can be exported against this serializer instance.
75 Note: Any fields which should be excluded from export have already been removed
77 Returns:
78 dict: A dictionary of field names and field objects
79 """
80 fields = {}
82 if meta := getattr(self, 'Meta', None):
83 write_only_fields = getattr(meta, 'write_only_fields', [])
84 else:
85 write_only_fields = []
87 for name, field in self.fields.items():
88 # Skip write-only fields
89 if getattr(field, 'write_only', False) or name in write_only_fields:
90 continue
92 # Skip tags fields
93 # TODO: Implement tag field export support
94 if issubclass(field.__class__, TagListSerializerField):
95 continue
97 # Top-level serializer fields can be exported with dot notation
98 if issubclass(field.__class__, serializers.Serializer):
99 fields.update(self.get_child_fields(name, field))
100 continue
102 fields[name] = field
104 return fields
106 def get_child_fields(self, field_name: str, field) -> dict:
107 """Return a dictionary of child fields for a given field.
109 Only child fields which match the 'export_child_fields' list will be returned.
110 """
111 child_fields = {}
113 if sub_fields := getattr(field, 'fields', None):
114 for sub_name, sub_field in sub_fields.items():
115 name = f'{field_name}.{sub_name}'
117 if name in self.export_child_fields:
118 sub_field.parent_field = field
119 child_fields[name] = sub_field
121 return child_fields
123 @classmethod
124 def arrange_export_headers(cls, headers: list) -> list:
125 """Optional method to arrange the export headers.
127 By default, the headers are returned in the order they are provided.
128 """
129 return headers
131 def get_nested_value(self, row: dict, key: str) -> Any:
132 """Get a nested value from a dictionary.
134 This method allows for dot notation to access nested fields.
136 Arguments:
137 row: The dictionary to extract the value from
138 key: The key to extract
140 Returns:
141 any: The extracted value
142 """
143 keys = key.split('.')
145 value = row
147 for key in keys:
148 if not value:
149 break
151 if not key:
152 continue
154 value = value.get(key, None)
156 return value
158 def generate_headers(self) -> OrderedDict:
159 """Generate a list of default headers for the exported data.
161 Returns an ordered dict of field names and their corresponding human-readable labels.
163 e.g.
165 {
166 'id': 'ID',
167 'name': 'Name',
168 ...
169 }
171 """
172 fields = self.get_exportable_fields()
173 field_names = self.arrange_export_headers(list(fields.keys()))
175 headers = OrderedDict()
177 for field_name in field_names:
178 field = fields[field_name]
180 label = getattr(field, 'label', field_name)
182 if parent := getattr(field, 'parent_field', None):
183 label = f'{parent.label}.{label}'
185 headers[field_name] = label
187 return headers
189 def export_to_file(self, data, headers: OrderedDict, file_format):
190 """Export the queryset to a file in the specified format.
192 Arguments:
193 queryset: The queryset to export
194 data: The serialized dataset to export
195 headers: The headers to use for the exported data {field: label}
196 file_format: The file format to export to
198 Returns:
199 File object containing the exported data
200 """
201 field_names = list(headers.keys())
202 field_headers = list(headers.values())
204 # Create a new dataset with the provided header labels
205 dataset = tablib.Dataset(headers=field_headers)
207 for row in data:
208 dataset.append([self.get_nested_value(row, f) for f in field_names])
210 return dataset.export(file_format)
213class DataExportViewMixin:
214 """An API view mixin for directly exporting selected data.
216 To perform a data export against an API endpoint which inherits from this mixin,
217 perform a GET request with 'export=True'.
219 This will run validation against the DataExportOptionsSerializer.
221 Once the export options have been validated, a new DataOutput object will be created,
222 and this will be returned to the client (including a download link to the exported file).
223 """
225 def is_exporting(self) -> bool:
226 """Determine if the view is currently exporting data."""
227 if request := getattr(self, 'request', None): 227 ↛ 232line 227 didn't jump to line 232 because the condition on line 227 was always true
228 return str2bool(
229 request.data.get('export') or request.query_params.get('export')
230 )
232 return False
234 def get_plugin(self, plugin_slug=None):
235 """Return the plugin instance associated with the export request.
237 Arguments:
238 plugin_slug: The slug of the plugin to use for exporting the data (optional)
239 """
240 PLUGIN_KEY = 'export_plugin'
242 if not plugin_slug:
243 if request := getattr(self, 'request', None):
244 plugin_slug = request.data.get(PLUGIN_KEY) or request.query_params.get(
245 PLUGIN_KEY
246 )
248 if plugin_slug:
249 return registry.get_plugin(
250 plugin_slug, active=True, with_mixin=PluginMixinEnum.EXPORTER
251 )
253 return None
255 def get_serializer(self, *args, **kwargs):
256 """Return the serializer instance for the view.
258 - Only applies for OPTIONS or GET requests
259 - OPTIONS requests to determine plugin serializer options
260 - GET request to perform the data export
261 - If the view is exporting data, return the DataExportOptionsSerializer.
262 - Otherwise, return the default serializer.
263 """
264 exporting = kwargs.pop('exporting', None)
266 if exporting is None: 266 ↛ 270line 266 didn't jump to line 270 because the condition on line 266 was always true
267 method = str(getattr(self.request, 'method', '')).lower()
268 exporting = method in ['options', 'get'] and self.is_exporting()
270 if exporting: 270 ↛ 272line 270 didn't jump to line 272 because the condition on line 270 was never true
271 # Override kwargs when initializing the DataExportOptionsSerializer
272 export_kwargs = {
273 'plugin': self.get_plugin(),
274 'request': self.request,
275 'data': kwargs.get('data'),
276 'context': kwargs.get('context'),
277 }
279 # Get the base model associated with this view
280 try:
281 serializer_class = self.get_serializer_class()
282 export_kwargs['serializer_class'] = serializer_class
283 export_kwargs['model_class'] = serializer_class.Meta.model
284 export_kwargs['view_class'] = self.__class__
285 except AttributeError:
286 # If the serializer class is not available, set to None
287 export_kwargs['serializer_class'] = None
288 export_kwargs['model_class'] = None
289 export_kwargs['view_class'] = None
291 return data_exporter.serializers.DataExportOptionsSerializer(
292 *args, **export_kwargs
293 )
294 else:
295 return super().get_serializer(*args, **kwargs)
297 def export_data(
298 self,
299 export_plugin,
300 export_format: str,
301 export_context: dict,
302 output: DataOutput,
303 ):
304 """Export the data in the specified format.
306 Arguments:
307 export_plugin: The plugin instance to use for exporting the data. If not provided, the default exporter is used
308 export_format: The file format to export the data in
309 export_context: Additional context data to pass to the plugin
310 output: The DataOutput object to write to
312 - By default, uses the provided serializer to generate the data, and return it as a file download.
313 - If a plugin is specified, the plugin can be used to augment or replace the export functionality.
314 """
315 if export_plugin is None:
316 from plugin.registry import registry
318 export_plugin = registry.get_plugin('inventree-exporter')
320 # Get the base serializer class for the view
321 serializer_class = self.get_serializer_class()
323 if not issubclass(serializer_class, DataExportSerializerMixin):
324 raise ValidationError(
325 'Serializer class must inherit from DataExportSerializerMixin'
326 )
328 export_error = _('Error occurred during data export')
330 context = self.get_serializer_context()
332 # Perform initial filtering of the queryset, based on the query parameters
333 queryset = self.filter_queryset(self.get_queryset())
335 # Perform additional filtering, as per the provided plugin
336 try:
337 queryset = export_plugin.filter_queryset(queryset)
338 except Exception:
339 InvenTree.exceptions.log_error('filter_queryset', plugin=export_plugin.slug)
340 raise ValidationError(export_error)
342 # Update the output instance with the total number of items to export
343 output.total = queryset.count()
344 output.save()
345 request = context.get('request', None)
347 if request:
348 query_params = getattr(request, 'query_params', {})
349 context.update(**query_params)
350 context['request'] = request
352 data = None
353 serializer = serializer_class(context=context, exporting=True)
354 serializer.initial_data = queryset
356 # Construct 'default' headers (note: may be overridden by plugin)
357 headers = serializer.generate_headers()
359 # Generate a filename for the exported data (implemented by the plugin)
360 try:
361 filename = export_plugin.generate_filename(
362 serializer_class.Meta.model, export_format
363 )
364 except Exception as e:
365 InvenTree.exceptions.log_error(
366 'generate_filename', plugin=export_plugin.slug
367 )
369 output.mark_failure(error=str(e))
371 raise ValidationError(export_error)
373 # The provided plugin is responsible for exporting the data
374 # The returned data *must* be a list of dict objects
375 try:
376 data = export_plugin.export_data(
377 queryset,
378 serializer_class,
379 headers,
380 export_context,
381 output,
382 serializer_context=context,
383 )
385 except Exception as e:
386 InvenTree.exceptions.log_error('export_data', plugin=export_plugin.slug)
388 # Log the error against the output object
389 output.mark_failure(error=str(e))
391 raise ValidationError(export_error)
393 if not isinstance(data, list):
394 raise ValidationError(
395 _('Data export plugin returned incorrect data format')
396 )
398 # Augment / update the headers (if required)
399 if hasattr(export_plugin, 'update_headers'):
400 try:
401 headers = export_plugin.update_headers(headers, export_context)
402 except Exception as e:
403 InvenTree.exceptions.log_error(
404 'update_headers', plugin=export_plugin.slug
405 )
407 output.mark_failure(error=str(e))
409 raise ValidationError(export_error)
411 # Now, export the data to file
412 try:
413 datafile = serializer.export_to_file(data, headers, export_format)
414 except Exception as e:
415 InvenTree.exceptions.log_error('export_to_file', plugin=export_plugin.slug)
416 output.mark_failure(error=str(e))
417 raise ValidationError(_('Error occurred during data export'))
419 # Update the output object with the exported data
420 output.mark_complete(output=ContentFile(datafile, filename))
422 def get(self, request, *args, **kwargs):
423 """Override the GET method to determine export options."""
424 from common.serializers import DataOutputSerializer
426 # If we are not exporting data, return the default response
427 if self.is_exporting(): 427 ↛ 431line 427 didn't jump to line 431 because the condition on line 427 was never true
428 # Determine if the export options are valid
430 # Extract the export options from the provided query parameters
431 export_options = {}
433 for key in request.query_params:
434 if key.startswith('export_'):
435 export_options[key] = request.query_params.get(key)
437 # Construct the options serializer with the provided data
438 serializer = self.get_serializer(exporting=True, data=export_options)
440 serializer.is_valid(raise_exception=True)
441 serializer_data = serializer.validated_data
443 export_format = serializer_data.pop('export_format', 'csv')
444 plugin_slug = serializer_data.pop('export_plugin', 'inventree-exporter')
445 export_plugin = self.get_plugin(plugin_slug)
447 export_context = {}
449 # Also run the data against the plugin serializer
450 if export_plugin:
451 if hasattr(export_plugin, 'get_export_options_serializer'):
452 if plugin_serializer := export_plugin.get_export_options_serializer(
453 data=export_options
454 ):
455 plugin_serializer.is_valid(raise_exception=True)
456 export_context = plugin_serializer.validated_data
458 user = getattr(request, 'user', None)
460 # Add in extra context data for the plugin
461 export_context['user'] = user
463 # Create an output object to export against
464 output = DataOutput.objects.create(
465 user=user if user and user.is_authenticated else None,
466 total=0, # Note: this should get updated by the export task
467 progress=0,
468 complete=False,
469 output_type=DataOutput.DataOutputTypes.EXPORT,
470 plugin=export_plugin.slug,
471 output=None,
472 )
474 # Offload the export task to a background worker
475 # This is to avoid blocking the web server
476 # Note: The export task will loop back and call the 'export_data' method on this class
477 offload_task(
478 data_exporter.tasks.export_data,
479 self.__class__,
480 request.user.id,
481 request.query_params,
482 plugin_slug,
483 export_format,
484 export_context,
485 output.id,
486 group='exporter',
487 )
489 output.refresh_from_db()
491 # Return a response to the frontend
492 return Response(DataOutputSerializer(output).data, status=200)
494 return super().get(request, *args, **kwargs)