jinja2_environment_from_request() plugin hook

Closes #2225

datasette/app.py

@@ -420,21 +420,31 @@ class Datasette:
                 ),
             ]
         )
-        self.jinja_env = Environment(
+        environment = Environment(
             loader=template_loader,
             autoescape=True,
             enable_async=True,
             # undefined=StrictUndefined,
         )
-        self.jinja_env.filters["escape_css_string"] = escape_css_string
-        self.jinja_env.filters["quote_plus"] = urllib.parse.quote_plus
-        self.jinja_env.filters["escape_sqlite"] = escape_sqlite
-        self.jinja_env.filters["to_css_class"] = to_css_class
+        environment.filters["escape_css_string"] = escape_css_string
+        environment.filters["quote_plus"] = urllib.parse.quote_plus
+        self._jinja_env = environment
+        environment.filters["escape_sqlite"] = escape_sqlite
+        environment.filters["to_css_class"] = to_css_class
         self._register_renderers()
         self._permission_checks = collections.deque(maxlen=200)
         self._root_token = secrets.token_hex(32)
         self.client = DatasetteClient(self)
 
+    def get_jinja_environment(self, request: Request = None) -> Environment:
+        environment = self._jinja_env
+        if request:
+            for environment in pm.hook.jinja2_environment_from_request(
+                datasette=self, request=request, env=environment
+            ):
+                pass
+        return environment
+
     def get_permission(self, name_or_abbr: str) -> "Permission":
         """
         Returns a Permission object for the given name or abbreviation. Raises KeyError if not found.
@@ -514,7 +524,7 @@ class Datasette:
                         abbrs[p.abbr] = p
                     self.permissions[p.name] = p
         for hook in pm.hook.prepare_jinja2_environment(
-            env=self.jinja_env, datasette=self
+            env=self._jinja_env, datasette=self
         ):
             await await_me_maybe(hook)
         for hook in pm.hook.startup(datasette=self):
@@ -1218,7 +1228,7 @@ class Datasette:
         else:
             if isinstance(templates, str):
                 templates = [templates]
-            template = self.jinja_env.select_template(templates)
+            template = self.get_jinja_environment(request).select_template(templates)
         if dataclasses.is_dataclass(context):
             context = dataclasses.asdict(context)
         body_scripts = []
@@ -1568,16 +1578,6 @@ class DatasetteRouter:
     def __init__(self, datasette, routes):
         self.ds = datasette
         self.routes = routes or []
-        # Build a list of pages/blah/{name}.html matching expressions
-        pattern_templates = [
-            filepath
-            for filepath in self.ds.jinja_env.list_templates()
-            if "{" in filepath and filepath.startswith("pages/")
-        ]
-        self.page_routes = [
-            (route_pattern_from_filepath(filepath[len("pages/") :]), filepath)
-            for filepath in pattern_templates
-        ]
 
     async def __call__(self, scope, receive, send):
         # Because we care about "foo/bar" v.s. "foo%2Fbar" we decode raw_path ourselves
@@ -1677,13 +1677,24 @@ class DatasetteRouter:
             route_path = request.scope.get("route_path", request.scope["path"])
             # Jinja requires template names to use "/" even on Windows
             template_name = "pages" + route_path + ".html"
+            # Build a list of pages/blah/{name}.html matching expressions
+            environment = self.ds.get_jinja_environment(request)
+            pattern_templates = [
+                filepath
+                for filepath in environment.list_templates()
+                if "{" in filepath and filepath.startswith("pages/")
+            ]
+            page_routes = [
+                (route_pattern_from_filepath(filepath[len("pages/") :]), filepath)
+                for filepath in pattern_templates
+            ]
             try:
-                template = self.ds.jinja_env.select_template([template_name])
+                template = environment.select_template([template_name])
             except TemplateNotFound:
                 template = None
             if template is None:
                 # Try for a pages/blah/{name}.html template match
-                for regex, wildcard_template in self.page_routes:
+                for regex, wildcard_template in page_routes:
                     match = regex.match(route_path)
                     if match is not None:
                         context.update(match.groupdict())

datasette/handle_exception.py

@@ -57,7 +57,8 @@ def handle_exception(datasette, request, exception):
         if request.path.split("?")[0].endswith(".json"):
             return Response.json(info, status=status, headers=headers)
         else:
-            template = datasette.jinja_env.select_template(templates)
+            environment = datasette.get_jinja_environment(request)
+            template = environment.select_template(templates)
             return Response.html(
                 await template.render_async(
                     dict(

datasette/hookspecs.py

@@ -99,6 +99,11 @@ def actors_from_ids(datasette, actor_ids):
     """Returns a dictionary mapping those IDs to actor dictionaries"""
 
 
+@hookspec
+def jinja2_environment_from_request(datasette, request, env):
+    """Return a Jinja2 environment based on the incoming request"""
+
+
 @hookspec
 def filters_from_request(request, database, table, datasette):
     """

datasette/views/base.py

@@ -143,7 +143,8 @@ class BaseView:
 
     async def render(self, templates, request, context=None):
         context = context or {}
-        template = self.ds.jinja_env.select_template(templates)
+        environment = self.ds.get_jinja_environment(request)
+        template = environment.select_template(templates)
         template_context = {
             **context,
             **{

datasette/views/database.py

@@ -143,7 +143,8 @@ class DatabaseView(View):
             datasette.urls.path(path_with_format(request=request, format="json")),
         )
         templates = (f"database-{to_css_class(database)}.html", "database.html")
-        template = datasette.jinja_env.select_template(templates)
+        environment = datasette.get_jinja_environment(request)
+        template = environment.select_template(templates)
         context = {
             **json_data,
             "database_color": db.color,
@@ -594,7 +595,8 @@ class QueryView(View):
                     f"query-{to_css_class(database)}-{to_css_class(canned_query['name'])}.html",
                 )
 
-            template = datasette.jinja_env.select_template(templates)
+            environment = datasette.get_jinja_environment(request)
+            template = environment.select_template(templates)
             alternate_url_json = datasette.absolute_url(
                 request,
                 datasette.urls.path(path_with_format(request=request, format="json")),

datasette/views/table.py

@@ -806,7 +806,8 @@ async def table_view_traced(datasette, request):
             f"table-{to_css_class(resolved.db.name)}-{to_css_class(resolved.table)}.html",
             "table.html",
         ]
-        template = datasette.jinja_env.select_template(templates)
+        environment = datasette.get_jinja_environment(request)
+        template = environment.select_template(templates)
         alternate_url_json = datasette.absolute_url(
             request,
             datasette.urls.path(path_with_format(request=request, format="json")),

docs/plugin_hooks.rst

@@ -1128,6 +1128,48 @@ These IDs could be integers or strings, depending on how the actors used by the
 
 Example: `datasette-remote-actors <https://github.com/datasette/datasette-remote-actors>`_
 
+.. _plugin_hook_jinja2_environment_from_request:
+
+jinja2_environment_from_request(datasette, request, env)
+--------------------------------------------------------
+
+``datasette`` - :ref:`internals_datasette`
+    A Datasette instance.
+
+``request`` - :ref:`internals_request` or ``None``
+    The current HTTP request, if one is available.
+
+``env`` - ``Environment``
+    The Jinja2 environment that will be used to render the current page.
+
+This hook can be used to return a customized `Jinja environment <https://jinja.palletsprojects.com/en/3.0.x/api/#jinja2.Environment>`__ based on the incoming request.
+
+If you want to run a single Datasette instance that serves different content for different domains, you can do so like this:
+
+.. code-block:: python
+
+    from datasette import hookimpl
+    from jinja2 import ChoiceLoader, FileSystemLoader
+
+
+    @hookimpl
+    def jinja2_environment_from_request(request, env):
+        if request and request.host == "www.niche-museums.com":
+            return env.overlay(
+                loader=ChoiceLoader(
+                    [
+                        FileSystemLoader(
+                            "/mnt/niche-museums/templates"
+                        ),
+                        env.loader,
+                    ]
+                ),
+                enable_async=True,
+            )
+        return env
+
+This uses the Jinja `overlay() method <https://jinja.palletsprojects.com/en/3.0.x/api/#jinja2.Environment.overlay>`__ to create a new environment identical to the default environment except for having a different template loader, which first looks in the ``/mnt/niche-museums/templates`` directory before falling back on the default loader.
+
 .. _plugin_hook_filters_from_request:
 
 filters_from_request(request, database, table, datasette)

tests/test_plugins.py

@@ -16,6 +16,7 @@ from datasette.plugins import get_plugins, DEFAULT_PLUGINS, pm
 from datasette.utils.sqlite import sqlite3
 from datasette.utils import CustomRow, StartupError
 from jinja2.environment import Template
+from jinja2 import ChoiceLoader, FileSystemLoader
 import base64
 import importlib
 import json
@@ -563,7 +564,8 @@ async def test_hook_register_output_renderer_can_render(ds_client):
 async def test_hook_prepare_jinja2_environment(ds_client):
     ds_client.ds._HELLO = "HI"
     await ds_client.ds.invoke_startup()
-    template = ds_client.ds.jinja_env.from_string(
+    environment = ds_client.ds.get_jinja_environment(None)
+    template = environment.from_string(
         "Hello there, {{ a|format_numeric }}, {{ a|to_hello }}, {{ b|select_times_three }}",
         {"a": 3412341, "b": 5},
     )
@@ -1294,3 +1296,41 @@ async def test_plugin_is_installed():
 
     finally:
         pm.unregister(name="DummyPlugin")
+
+
+@pytest.mark.asyncio
+async def test_hook_jinja2_environment_from_request(tmpdir):
+    templates = pathlib.Path(tmpdir / "templates")
+    templates.mkdir()
+    (templates / "index.html").write_text("Hello museums!", "utf-8")
+
+    class EnvironmentPlugin:
+        @hookimpl
+        def jinja2_environment_from_request(self, request, env):
+            if request and request.host == "www.niche-museums.com":
+                return env.overlay(
+                    loader=ChoiceLoader(
+                        [
+                            FileSystemLoader(str(templates)),
+                            env.loader,
+                        ]
+                    ),
+                    enable_async=True,
+                )
+            return env
+
+    datasette = Datasette(memory=True)
+
+    try:
+        pm.register(EnvironmentPlugin(), name="EnvironmentPlugin")
+        response = await datasette.client.get("/")
+        assert response.status_code == 200
+        assert "Hello museums!" not in response.text
+        # Try again with the hostname
+        response2 = await datasette.client.get(
+            "/", headers={"host": "www.niche-museums.com"}
+        )
+        assert response2.status_code == 200
+        assert "Hello museums!" in response2.text
+    finally:
+        pm.unregister(name="EnvironmentPlugin")