Kaynağa Gözat

Address review feedback on exporter responses

- Stream an exporter's iterable only if STREAMING_EXPORTS is enabled
- Return bytearray and memoryview data as regular content
- Raise a clear error when export() is called on a table exporter
Jeremy Stretch 5 gün önce
ebeveyn
işleme
d954ca94af

+ 1 - 1
docs/configuration/miscellaneous.md

@@ -127,7 +127,7 @@ The maximum size (in bytes) of an incoming HTTP request (i.e. `GET` or `POST` da
 
 Default: `False`
 
-When set to `True`, CSV bulk exports are returned as a streaming HTTP response, emitting rows to the client as they are rendered rather than buffering the entire dataset in memory first. This can significantly reduce memory usage and time-to-first-byte for very large exports.
+When set to `True`, CSV bulk exports (and exports from any plugin exporters which return an iterable) are returned as a streaming HTTP response, emitting rows to the client as they are rendered rather than buffering the entire dataset in memory first. This can significantly reduce memory usage and time-to-first-byte for very large exports.
 
 Because streaming responses do not have a `Content-Length` header and defer errors until after the response has begun, this behavior is opt-in.
 

+ 2 - 2
docs/plugins/development/exporters.md

@@ -4,7 +4,7 @@
 
 The "Export" button on object list views offers a choice of export formats. NetBox provides CSV and (for device types and module types) YAML exporters. Plugins can register their own exporters to add formats, such as XLSX or JSON, by subclassing NetBox's `TableExporter` or `Exporter` class.
 
-A `TableExporter` exports the rows of the list view's table. NetBox resolves the table's columns and passes the exporter a list of column headers and an iterator of cell values for each row. Table exporters are offered under both "Current View" (the user's selected columns) and "All Data" (all columns).
+A `TableExporter` exports the rows of the list view's table. NetBox resolves the table's columns and calls the exporter's `export_rows()` method (rather than `export()`) with a list of column headers and an iterator of cell values for each row. Table exporters are offered under both "Current View" (the user's selected columns) and "All Data" (all columns).
 
 ```python title="exporters.py"
 import io
@@ -50,7 +50,7 @@ class VendorDeviceTypeExporter(Exporter):
         ...
 ```
 
-The queryset passed to an exporter has the user's filters applied, and is restricted to objects the user has permission to view. `export()` and `export_rows()` return the exported data as a string or bytes, or as an iterable of chunks, which NetBox streams to the client.
+The queryset passed to an exporter has the user's filters applied, and is restricted to objects the user has permission to view. `export()` and `export_rows()` return the exported data as a string or bytes, or as an iterable of chunks. NetBox streams an iterable to the client if [`STREAMING_EXPORTS`](../../configuration/miscellaneous.md#streaming_exports) is enabled; otherwise, it collects the chunks before responding.
 
 The `register_exporter()` decorator registers an exporter with NetBox. Exporters can be defined anywhere within a plugin, but must be imported during plugin initialization, under the PluginConfig's `ready()` method:
 

+ 8 - 9
netbox/netbox/exporters.py

@@ -1,7 +1,6 @@
 import csv
 from itertools import chain
 
-from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.db.models import FETCH_PEERS
 from django.utils.translation import gettext as _
@@ -50,7 +49,8 @@ class Exporter:
 
     def export(self, *, request, queryset):
         """
-        Return the exported data as str or bytes, or as an iterable of chunks to stream the response.
+        Return the exported data as str or bytes, or as an iterable of chunks (streamed if STREAMING_EXPORTS is
+        enabled).
 
         Args:
             request: The current request
@@ -62,8 +62,11 @@ class Exporter:
 class TableExporter(Exporter):
     """
     An exporter which renders the rows of an object list table. Table exporters are listed under both "Current View"
-    and "All Data".
+    and "All Data". Subclasses implement export_rows() rather than export().
     """
+    def export(self, *, request, queryset):
+        raise NotImplementedError(f"{self.__class__.__name__} is a table exporter; call export_rows() instead")
+
     def export_rows(self, *, request, headers, rows):
         """
         Return the exported data (as for export()) given a list of column headers and an iterator of row values.
@@ -96,8 +99,7 @@ class CSVExporter(TableExporter):
 
     def write_rows(self, headers, rows, delimiter=None):
         """
-        Return the given headers and rows as CSV, either as a string or (if STREAMING_EXPORTS is enabled) as an
-        iterator of lines.
+        Return an iterator of CSV lines for the given headers and rows.
 
         Args:
             headers: A list of column headers
@@ -107,10 +109,7 @@ class CSVExporter(TableExporter):
         if delimiter and delimiter not in CSV_DELIMITERS:
             raise ValueError(_("Invalid delimiter name: {name}").format(name=delimiter))
         writer = csv.writer(_EchoBuffer(), delimiter=CSV_DELIMITERS[delimiter or 'comma'])
-        lines = (writer.writerow(row) for row in chain([headers], rows))
-        if settings.STREAMING_EXPORTS:
-            return lines
-        return ''.join(lines)
+        return (writer.writerow(row) for row in chain([headers], rows))
 
 
 class YAMLExporter(Exporter):

+ 28 - 7
netbox/netbox/tests/test_exporters.py

@@ -1,4 +1,5 @@
 import csv
+from collections.abc import Iterator
 from io import StringIO
 from unittest import skipIf
 
@@ -96,20 +97,19 @@ class CSVExporterTestCase(SimpleTestCase):
 
     def test_write_rows(self):
         data = get_exporter('csv').write_rows(self.headers, iter(self.rows))
-        self.assertEqual(data, 'Name,Count\r\nSite 1,1\r\n"Site, 2",\r\n')
+        self.assertEqual(list(data), ['Name,Count\r\n', 'Site 1,1\r\n', '"Site, 2",\r\n'])
 
     def test_write_rows_delimiter(self):
         data = get_exporter('csv').write_rows(self.headers, iter(self.rows), delimiter='semicolon')
-        self.assertEqual(data, 'Name;Count\r\nSite 1;1\r\nSite, 2;\r\n')
+        self.assertEqual(''.join(data), 'Name;Count\r\nSite 1;1\r\nSite, 2;\r\n')
 
     def test_write_rows_invalid_delimiter(self):
         with self.assertRaises(ValueError):
             get_exporter('csv').write_rows(self.headers, iter(self.rows), delimiter='invalid')
 
-    @override_settings(STREAMING_EXPORTS=True)
-    def test_write_rows_streaming(self):
-        data = get_exporter('csv').write_rows(self.headers, iter(self.rows))
-        self.assertEqual(list(data), ['Name,Count\r\n', 'Site 1,1\r\n', '"Site, 2",\r\n'])
+    def test_export(self):
+        with self.assertRaisesMessage(NotImplementedError, 'call export_rows() instead'):
+            get_exporter('csv').export(request=None, queryset=None)
 
 
 class YAMLExporterTestCase(TestCase):
@@ -197,8 +197,14 @@ class ExportViewTestCase(TestCase):
     def test_export_plugin_exporter(self):
         response = self.client.get(f'{self.url}?export&export_format=dummy_plugin.site_names')
         self.assertHttpStatus(response, 200)
-        self.assertIsInstance(response, StreamingHttpResponse)
+        self.assertNotIsInstance(response, StreamingHttpResponse)
         self.assertEqual(response['Content-Disposition'], 'attachment; filename="netbox_sites.txt"')
+        self.assertEqual(response.content.decode(), 'Site 1\nSite 2\nSite 3\n')
+
+        # Streamed if STREAMING_EXPORTS is enabled
+        with override_settings(STREAMING_EXPORTS=True):
+            response = self.client.get(f'{self.url}?export&export_format=dummy_plugin.site_names')
+        self.assertIsInstance(response, StreamingHttpResponse)
         self.assertEqual(response.getvalue().decode(), 'Site 1\nSite 2\nSite 3\n')
 
         # Not available for the current view
@@ -222,6 +228,21 @@ class ExportViewTestCase(TestCase):
         self.assertNotContains(response, 'export=table&export_format=dummy_plugin.site_names')
         self.assertNotContains(response, 'export_format=yaml')
 
+    def test_export_response(self):
+        view = SiteListView()
+        view.queryset = Site.objects.all()
+        exporter = get_exporter('csv')
+
+        for streaming in (False, True):
+            for data in ('a,b', b'a,b', bytearray(b'a,b'), memoryview(b'a,b'), iter(['a,', 'b'])):
+                with self.subTest(streaming=streaming, type=type(data).__name__):
+                    with override_settings(STREAMING_EXPORTS=streaming):
+                        response = view.get_export_response(exporter, data)
+                    # Only an iterator is streamed, and only if STREAMING_EXPORTS is enabled
+                    expect_stream = streaming and isinstance(data, Iterator)
+                    self.assertEqual(isinstance(response, StreamingHttpResponse), expect_stream)
+                    self.assertEqual(response.getvalue(), b'a,b')
+
     def test_export_table_override(self):
         """
         Overriding the deprecated export_table() still applies to CSV exports.

+ 8 - 4
netbox/netbox/views/generic/bulk_views.py

@@ -5,6 +5,7 @@ from contextlib import contextmanager
 from copy import deepcopy
 from types import SimpleNamespace
 
+from django.conf import settings
 from django.contrib import messages
 from django.contrib.contenttypes.fields import GenericForeignKey, GenericRel
 from django.core.exceptions import FieldDoesNotExist, ImproperlyConfigured, ObjectDoesNotExist, ValidationError
@@ -154,15 +155,18 @@ class ObjectListView(BaseMultiObjectView, ActionsMixin, TableMixin):
 
         Args:
             exporter: The Exporter instance which produced the data
-            data: The exported data, as str or bytes, or as an iterable of chunks to be streamed
+            data: The exported data, as str or bytes, or as an iterable of chunks (streamed if STREAMING_EXPORTS is
+                enabled)
             filename: The name of the file attachment. If None, will be determined automatically from the queryset
                 model name.
         """
         filename = filename or f'netbox_{self.queryset.model._meta.verbose_name_plural}.{exporter.file_extension}'
-        if isinstance(data, (str, bytes)):
-            response = HttpResponse(data, content_type=exporter.content_type)
-        else:
+        if isinstance(data, (bytearray, memoryview)):
+            data = bytes(data)
+        if settings.STREAMING_EXPORTS and not isinstance(data, (str, bytes)):
             response = StreamingHttpResponse(data, content_type=exporter.content_type)
+        else:
+            response = HttpResponse(data, content_type=exporter.content_type)
         response['Content-Disposition'] = content_disposition_header(as_attachment=True, filename=filename)
         return response