Browse Source

Refactored CSV export logic

Jeremy Stretch 7 years ago
parent
commit
59dcbce417

+ 5 - 6
netbox/circuits/models.py

@@ -9,7 +9,6 @@ from dcim.fields import ASNField
 from extras.models import CustomFieldModel, CustomFieldValue
 from tenancy.models import Tenant
 from utilities.models import CreatedUpdatedModel
-from utilities.utils import csv_format
 from .constants import *
 
 
@@ -41,13 +40,13 @@ class Provider(CreatedUpdatedModel, CustomFieldModel):
         return reverse('circuits:provider', args=[self.slug])
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.slug,
             self.asn,
             self.account,
             self.portal_url,
-        ])
+        )
 
 
 @python_2_unicode_compatible
@@ -99,15 +98,15 @@ class Circuit(CreatedUpdatedModel, CustomFieldModel):
         return reverse('circuits:circuit', args=[self.pk])
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.cid,
             self.provider.name,
             self.type.name,
             self.tenant.name if self.tenant else None,
-            self.install_date.isoformat() if self.install_date else None,
+            self.install_date,
             self.commit_rate,
             self.description,
-        ])
+        )
 
     def _get_termination(self, side):
         for ct in self.terminations.all():

+ 22 - 23
netbox/dcim/models.py

@@ -22,7 +22,6 @@ from tenancy.models import Tenant
 from utilities.fields import ColorField, NullableCharField
 from utilities.managers import NaturalOrderByManager
 from utilities.models import CreatedUpdatedModel
-from utilities.utils import csv_format
 from .constants import *
 from .fields import ASNField, MACAddressField
 from .querysets import InterfaceQuerySet
@@ -57,11 +56,11 @@ class Region(MPTTModel):
         return "{}?region={}".format(reverse('dcim:site_list'), self.slug)
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.slug,
             self.parent.name if self.parent else None,
-        ])
+        )
 
 
 #
@@ -111,7 +110,7 @@ class Site(CreatedUpdatedModel, CustomFieldModel):
         return reverse('dcim:site', args=[self.slug])
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.slug,
             self.region.name if self.region else None,
@@ -121,7 +120,7 @@ class Site(CreatedUpdatedModel, CustomFieldModel):
             self.contact_name,
             self.contact_phone,
             self.contact_email,
-        ])
+        )
 
     @property
     def count_prefixes(self):
@@ -182,11 +181,11 @@ class RackGroup(models.Model):
         return "{}?group_id={}".format(reverse('dcim:rack_list'), self.pk)
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.site,
             self.name,
             self.slug,
-        ])
+        )
 
 
 @python_2_unicode_compatible
@@ -292,7 +291,7 @@ class Rack(CreatedUpdatedModel, CustomFieldModel):
             Device.objects.filter(rack=self).update(site_id=self.site.pk)
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.site.name,
             self.group.name if self.group else None,
             self.name,
@@ -304,7 +303,7 @@ class Rack(CreatedUpdatedModel, CustomFieldModel):
             self.width,
             self.u_height,
             self.desc_units,
-        ])
+        )
 
     @property
     def units(self):
@@ -493,10 +492,10 @@ class Manufacturer(models.Model):
         return "{}?manufacturer={}".format(reverse('dcim:devicetype_list'), self.slug)
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.slug,
-        ])
+        )
 
 
 @python_2_unicode_compatible
@@ -562,7 +561,7 @@ class DeviceType(models.Model, CustomFieldModel):
         return reverse('dcim:devicetype', args=[self.pk])
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.manufacturer.name,
             self.model,
             self.slug,
@@ -574,7 +573,7 @@ class DeviceType(models.Model, CustomFieldModel):
             self.is_network_device,
             self.get_subdevice_role_display() if self.subdevice_role else None,
             self.get_interface_ordering_display(),
-        ])
+        )
 
     def clean(self):
 
@@ -989,7 +988,7 @@ class Device(CreatedUpdatedModel, CustomFieldModel):
         Device.objects.filter(parent_bay__device=self).update(site=self.site, rack=self.rack)
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name or '',
             self.device_role.name,
             self.tenant.name if self.tenant else None,
@@ -1004,7 +1003,7 @@ class Device(CreatedUpdatedModel, CustomFieldModel):
             self.rack.name if self.rack else None,
             self.position,
             self.get_face_display(),
-        ])
+        )
 
     @property
     def display_name(self):
@@ -1078,13 +1077,13 @@ class ConsolePort(models.Model):
 
     # Used for connections export
     def to_csv(self):
-        return csv_format([
+        return (
             self.cs_port.device.identifier if self.cs_port else None,
             self.cs_port.name if self.cs_port else None,
             self.device.identifier,
             self.name,
             self.get_connection_status_display(),
-        ])
+        )
 
 
 #
@@ -1155,13 +1154,13 @@ class PowerPort(models.Model):
 
     # Used for connections export
     def to_csv(self):
-        return csv_format([
+        return (
             self.power_outlet.device.identifier if self.power_outlet else None,
             self.power_outlet.name if self.power_outlet else None,
             self.device.identifier,
             self.name,
             self.get_connection_status_display(),
-        ])
+        )
 
 
 #
@@ -1384,13 +1383,13 @@ class InterfaceConnection(models.Model):
 
     # Used for connections export
     def to_csv(self):
-        return csv_format([
+        return (
             self.interface_a.device.identifier,
             self.interface_a.name,
             self.interface_b.device.identifier,
             self.interface_b.name,
             self.get_connection_status_display(),
-        ])
+        )
 
 
 #
@@ -1464,7 +1463,7 @@ class InventoryItem(models.Model):
         return self.name
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.device.name or '{' + self.device.pk + '}',
             self.name,
             self.manufacturer.name if self.manufacturer else None,
@@ -1472,4 +1471,4 @@ class InventoryItem(models.Model):
             self.serial,
             self.asset_tag,
             self.description
-        ])
+        )

+ 10 - 4
netbox/extras/models.py

@@ -223,19 +223,25 @@ class ExportTemplate(models.Model):
     def __str__(self):
         return '{}: {}'.format(self.content_type, self.name)
 
-    def to_response(self, context_dict, filename):
+    def render_to_response(self, queryset):
         """
         Render the template to an HTTP response, delivered as a named file attachment
         """
         template = Template(self.template_code)
         mime_type = 'text/plain' if not self.mime_type else self.mime_type
-        output = template.render(Context(context_dict))
+        output = template.render(Context({'queryset': queryset}))
+
         # Replace CRLF-style line terminators
         output = output.replace('\r\n', '\n')
+
+        # Build the response
         response = HttpResponse(output, content_type=mime_type)
-        if self.file_extension:
-            filename += '.{}'.format(self.file_extension)
+        filename = 'netbox_{}{}'.format(
+            queryset.model._meta.verbose_name_plural,
+            '.{}'.format(self.file_extension) if self.file_extension else ''
+        )
         response['Content-Disposition'] = 'attachment; filename="{}"'.format(filename)
+
         return response
 
 

+ 11 - 12
netbox/ipam/models.py

@@ -14,7 +14,6 @@ from dcim.models import Interface
 from extras.models import CustomFieldModel, CustomFieldValue
 from tenancy.models import Tenant
 from utilities.models import CreatedUpdatedModel
-from utilities.utils import csv_format
 from .constants import *
 from .fields import IPNetworkField, IPAddressField
 from .querysets import PrefixQuerySet
@@ -49,13 +48,13 @@ class VRF(CreatedUpdatedModel, CustomFieldModel):
         return reverse('ipam:vrf', args=[self.pk])
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.rd,
             self.tenant.name if self.tenant else None,
             self.enforce_unique,
             self.description,
-        ])
+        )
 
     @property
     def display_name(self):
@@ -147,12 +146,12 @@ class Aggregate(CreatedUpdatedModel, CustomFieldModel):
         super(Aggregate, self).save(*args, **kwargs)
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.prefix,
             self.rir.name,
-            self.date_added.isoformat() if self.date_added else None,
+            self.date_added,
             self.description,
-        ])
+        )
 
     def get_utilization(self):
         """
@@ -262,7 +261,7 @@ class Prefix(CreatedUpdatedModel, CustomFieldModel):
         super(Prefix, self).save(*args, **kwargs)
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.prefix,
             self.vrf.rd if self.vrf else None,
             self.tenant.name if self.tenant else None,
@@ -273,7 +272,7 @@ class Prefix(CreatedUpdatedModel, CustomFieldModel):
             self.role.name if self.role else None,
             self.is_pool,
             self.description,
-        ])
+        )
 
     def get_status_class(self):
         return STATUS_CHOICE_CLASSES[self.status]
@@ -461,7 +460,7 @@ class IPAddress(CreatedUpdatedModel, CustomFieldModel):
         else:
             is_primary = False
 
-        return csv_format([
+        return (
             self.address,
             self.vrf.rd if self.vrf else None,
             self.tenant.name if self.tenant else None,
@@ -472,7 +471,7 @@ class IPAddress(CreatedUpdatedModel, CustomFieldModel):
             self.interface.name if self.interface else None,
             is_primary,
             self.description,
-        ])
+        )
 
     @property
     def device(self):
@@ -577,7 +576,7 @@ class VLAN(CreatedUpdatedModel, CustomFieldModel):
             })
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.site.name if self.site else None,
             self.group.name if self.group else None,
             self.vid,
@@ -586,7 +585,7 @@ class VLAN(CreatedUpdatedModel, CustomFieldModel):
             self.get_status_display(),
             self.role.name if self.role else None,
             self.description,
-        ])
+        )
 
     @property
     def display_name(self):

+ 2 - 3
netbox/tenancy/models.py

@@ -7,7 +7,6 @@ from django.utils.encoding import python_2_unicode_compatible
 
 from extras.models import CustomFieldModel, CustomFieldValue
 from utilities.models import CreatedUpdatedModel
-from utilities.utils import csv_format
 
 
 @python_2_unicode_compatible
@@ -53,9 +52,9 @@ class Tenant(CreatedUpdatedModel, CustomFieldModel):
         return reverse('tenancy:tenant', args=[self.slug])
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.slug,
             self.group.name if self.group else None,
             self.description,
-        ])
+        )

+ 61 - 0
netbox/utilities/csv.py

@@ -0,0 +1,61 @@
+from __future__ import unicode_literals
+
+import datetime
+import six
+
+from django.http import HttpResponse
+
+
+def csv_format(data):
+    """
+    Encapsulate any data which contains a comma within double quotes.
+    """
+    csv = []
+    for value in data:
+
+        # Represent None or False with empty string
+        if value in [None, False]:
+            csv.append('')
+            continue
+
+        # Convert dates to ISO format
+        if isinstance(value, (datetime.date, datetime.datetime)):
+            value = value.isoformat()
+
+        # Force conversion to string first so we can check for any commas
+        if not isinstance(value, six.string_types):
+            value = '{}'.format(value)
+
+        # Double-quote the value if it contains a comma
+        if ',' in value:
+            csv.append('"{}"'.format(value))
+        else:
+            csv.append('{}'.format(value))
+
+    return ','.join(csv)
+
+
+def queryset_to_csv(queryset):
+    """
+    Export a queryset of objects as CSV, using the model's to_csv() method.
+    """
+    output = []
+
+    # Start with the column headers
+    headers = ','.join(queryset.model.csv_headers)
+    output.append(headers)
+
+    # Iterate through the queryset
+    for obj in queryset:
+        data = csv_format(obj.to_csv())
+        output.append(data)
+
+    # Build the HTTP response
+    response = HttpResponse(
+        '\n'.join(output),
+        content_type='text/csv'
+    )
+    filename = 'netbox_{}.csv'.format(queryset.model._meta.verbose_name_plural)
+    response['Content-Disposition'] = 'attachment; filename="{}"'.format(filename)
+
+    return response

+ 0 - 27
netbox/utilities/utils.py

@@ -1,32 +1,5 @@
 from __future__ import unicode_literals
 
-import six
-
-
-def csv_format(data):
-    """
-    Encapsulate any data which contains a comma within double quotes.
-    """
-    csv = []
-    for value in data:
-
-        # Represent None or False with empty string
-        if value in [None, False]:
-            csv.append('')
-            continue
-
-        # Force conversion to string first so we can check for any commas
-        if not isinstance(value, six.string_types):
-            value = '{}'.format(value)
-
-        # Double-quote the value if it contains a comma
-        if ',' in value:
-            csv.append('"{}"'.format(value))
-        else:
-            csv.append('{}'.format(value))
-
-    return ','.join(csv)
-
 
 def foreground_color(bg_color):
     """

+ 8 - 17
netbox/utilities/views.py

@@ -10,7 +10,6 @@ from django.core.exceptions import ValidationError
 from django.db import transaction, IntegrityError
 from django.db.models import ProtectedError
 from django.forms import CharField, Form, ModelMultipleChoiceField, MultipleHiddenInput, Textarea, TypedChoiceField
-from django.http import HttpResponse
 from django.shortcuts import get_object_or_404, redirect, render
 from django.template import TemplateSyntaxError
 from django.urls import reverse
@@ -21,6 +20,7 @@ from django.views.generic import View
 from django_tables2 import RequestConfig
 
 from extras.models import CustomField, CustomFieldValue, ExportTemplate, UserAction
+from utilities.csv import queryset_to_csv
 from utilities.forms import BootstrapMixin, CSVDataField
 from .error_handlers import handle_protectederror
 from .forms import ConfirmationForm
@@ -95,24 +95,15 @@ class ObjectListView(View):
             et = get_object_or_404(ExportTemplate, content_type=object_ct, name=request.GET.get('export'))
             queryset = CustomFieldQueryset(self.queryset, custom_fields) if custom_fields else self.queryset
             try:
-                response = et.to_response(context_dict={'queryset': queryset},
-                                          filename='netbox_{}'.format(model._meta.verbose_name_plural))
-                return response
+                return et.render_to_response(queryset)
             except TemplateSyntaxError:
-                messages.error(request, "There was an error rendering the selected export template ({})."
-                               .format(et.name))
-        # Fall back to built-in CSV export
+                messages.error(
+                    request,
+                    "There was an error rendering the selected export template ({}).".format(et.name)
+                )
+        # Fall back to built-in CSV export if no template was specified
         elif 'export' in request.GET and hasattr(model, 'to_csv'):
-            headers = getattr(model, 'csv_headers', None)
-            output = ','.join(headers) + '\n' if headers else ''
-            output += '\n'.join([obj.to_csv() for obj in self.queryset])
-            response = HttpResponse(
-                output,
-                content_type='text/csv'
-            )
-            response['Content-Disposition'] = 'attachment; filename="netbox_{}.csv"'\
-                .format(self.queryset.model._meta.verbose_name_plural)
-            return response
+            return queryset_to_csv(self.queryset)
 
         # Provide a hook to tweak the queryset based on the request immediately prior to rendering the object list
         self.queryset = self.alter_queryset(request)

+ 4 - 5
netbox/virtualization/models.py

@@ -10,7 +10,6 @@ from django.utils.encoding import python_2_unicode_compatible
 from dcim.models import Device
 from extras.models import CustomFieldModel, CustomFieldValue
 from utilities.models import CreatedUpdatedModel
-from utilities.utils import csv_format
 from .constants import STATUS_ACTIVE, STATUS_CHOICES, VM_STATUS_CLASSES
 
 
@@ -135,13 +134,13 @@ class Cluster(CreatedUpdatedModel, CustomFieldModel):
                 })
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.type.name,
             self.group.name if self.group else None,
             self.site.name if self.site else None,
             self.comments,
-        ])
+        )
 
 
 #
@@ -243,7 +242,7 @@ class VirtualMachine(CreatedUpdatedModel, CustomFieldModel):
         return reverse('virtualization:virtualmachine', args=[self.pk])
 
     def to_csv(self):
-        return csv_format([
+        return (
             self.name,
             self.get_status_display(),
             self.cluster.name,
@@ -253,7 +252,7 @@ class VirtualMachine(CreatedUpdatedModel, CustomFieldModel):
             self.memory,
             self.disk,
             self.comments,
-        ])
+        )
 
     def get_status_class(self):
         return VM_STATUS_CLASSES[self.status]