From adaeedd6a27700c943114463fe222b6c5d26bed2 Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@netboxlabs.com>
Date: Fri, 10 Feb 2023 16:01:41 -0500
Subject: [PATCH] Pre-fetch referenced templates

---
 netbox/extras/models/configs.py | 22 +++++++++++++---------
 netbox/utilities/jinja2.py      | 20 +++++++++++++++-----
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/netbox/extras/models/configs.py b/netbox/extras/models/configs.py
index 8c275d732..d1c90b446 100644
--- a/netbox/extras/models/configs.py
+++ b/netbox/extras/models/configs.py
@@ -230,16 +230,20 @@ class ConfigTemplate(SyncedDataMixin, ExportTemplatesMixin, ChangeLoggedModel):
         Render the contents of the template.
         """
         context = context or {}
-        template_code = self.template_code
-        # output = render_jinja2(template_code, context)
+        template_path = self.data_file.path
 
-        environment = SandboxedEnvironment(
-            loader=ConfigTemplateLoader(data_source=self.data_source)
-        )
+        # Initialize the template loader & cache the base template code
+        loader = ConfigTemplateLoader(data_source=self.data_source)
+        loader.cache_templates({
+            template_path: self.template_code
+        })
+
+        # Initialize the Jinja2 environment
+        environment = SandboxedEnvironment(loader=loader)
         environment.filters.update(get_config().JINJA2_FILTERS)
-        output = environment.from_string(source=template_code).render(**context)
+
+        # Render the template
+        output = environment.get_template(template_path).render(**context)
 
         # Replace CRLF-style line terminators
-        output = output.replace('\r\n', '\n')
-
-        return output
+        return output.replace('\r\n', '\n')
diff --git a/netbox/utilities/jinja2.py b/netbox/utilities/jinja2.py
index 81d11cc85..55eabd461 100644
--- a/netbox/utilities/jinja2.py
+++ b/netbox/utilities/jinja2.py
@@ -1,7 +1,6 @@
-import os
-
 from django.apps import apps
 from jinja2 import BaseLoader, TemplateNotFound
+from jinja2.meta import find_referenced_templates
 
 __all__ = (
     'ConfigTemplateLoader',
@@ -14,14 +13,25 @@ class ConfigTemplateLoader(BaseLoader):
     """
     def __init__(self, data_source):
         self.data_source = data_source
+        self._template_cache = {}
 
     def get_source(self, environment, template):
         DataFile = apps.get_model('core', 'DataFile')
 
+        # Retrieve template content from cache
         try:
-            datafile = DataFile.objects.get(source=self.data_source, path=template)
-            template_source = datafile.data_as_string
-        except DataFile.DoesNotExist:
+            template_source = self._template_cache[template]
+        except KeyError:
             raise TemplateNotFound(template)
 
+        # Find and pre-fetch referenced templates
+        if referenced_templates := find_referenced_templates(environment.parse(template_source)):
+            self.cache_templates({
+                df.path: df.data_as_string for df in
+                DataFile.objects.filter(source=self.data_source, path__in=referenced_templates)
+            })
+
         return template_source, template, True
+
+    def cache_templates(self, templates):
+        self._template_cache.update(templates)