From a8cdbcfd018a3e42182e3bcc65dfb36e916b08a2 Mon Sep 17 00:00:00 2001
From: TomGrozev <TomGrozev@users.noreply.github.com>
Date: Tue, 15 Dec 2020 15:05:48 +1100
Subject: [PATCH] Add support for custom fields in tables

---
 netbox/utilities/tables.py | 40 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/netbox/utilities/tables.py b/netbox/utilities/tables.py
index bf087f2c9..aa77d485e 100644
--- a/netbox/utilities/tables.py
+++ b/netbox/utilities/tables.py
@@ -1,12 +1,15 @@
 import django_tables2 as tables
 from django.contrib.auth.models import AnonymousUser
 from django.contrib.contenttypes.fields import GenericForeignKey
+from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import FieldDoesNotExist
+from django.db.models import Func, F, Value
 from django.db.models.fields.related import RelatedField
 from django.urls import reverse
 from django.utils.safestring import mark_safe
 from django_tables2.data import TableQuerysetData
 
+from extras.models import CustomField
 
 class BaseTable(tables.Table):
     """
@@ -14,12 +17,25 @@ class BaseTable(tables.Table):
 
     :param user: Personalize table display for the given user (optional). Has no effect if AnonymousUser is passed.
     """
+
     class Meta:
         attrs = {
             'class': 'table table-hover table-headings',
         }
 
     def __init__(self, *args, user=None, **kwargs):
+        # Add custom field columns
+        obj_type = ContentType.objects.get_for_model(self._meta.model)
+        custom_fields = {}
+
+        for cf in CustomField.objects.filter(content_types=obj_type):
+            name = 'cf_{}'.format(cf.name)
+            label = cf.label if cf.label != '' else cf.name
+            self.base_columns[name] = CustomFieldColumn(verbose_name=label)
+            custom_fields[name] = cf
+        self._meta.fields += tuple(custom_fields.keys())
+
+        # Init table
         super().__init__(*args, **kwargs)
 
         # Set default empty_text if none was provided
@@ -57,6 +73,12 @@ class BaseTable(tables.Table):
 
         # Dynamically update the table's QuerySet to ensure related fields are pre-fetched
         if isinstance(self.data, TableQuerysetData):
+            # Extract custom field values
+            cf_fields = {}
+            for key, cf in custom_fields.items():
+                cf_fields[key] = Func(F('custom_field_data'), Value(cf.name), function='jsonb_extract_path_text')
+            self.data.data = self.data.data.annotate(**cf_fields)
+
             prefetch_fields = []
             for column in self.columns:
                 if column.visible:
@@ -267,3 +289,21 @@ class TagColumn(tables.TemplateColumn):
             template_code=self.template_code,
             extra_context={'url_name': url_name}
         )
+
+
+class CustomFieldColumn(tables.Column):
+    """
+    Display custom fields in the appropriate format.
+    """
+
+    def render(self, record, bound_column, value):
+        if value:
+            if isinstance(value, list):
+                template = ''
+                for v in value:
+                    template += f'<span class="label label-default">{v}</span> '
+            else:
+                template = value
+            return mark_safe(template)
+        return self.default
+