From fcfcae21e67cc15090942b1d2a47b5f016279337 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 5 Jul 2019 17:05:56 -0700
Subject: [PATCH] extra_template_vars plugin hook (#542)

* extra_template_vars plugin hook

Closes #541

* Workaround for cwd bug

Based on https://github.com/pytest-dev/pytest/issues/1235#issuecomment-175295691
---
 datasette/hookspecs.py              |  5 ++
 datasette/views/base.py             | 25 ++++++++-
 datasette/views/index.py            | 11 ++--
 datasette/views/special.py          |  6 +-
 docs/plugins.rst                    | 86 +++++++++++++++++++++++++----
 tests/conftest.py                   | 15 +++++
 tests/fixtures.py                   | 23 ++++++++
 tests/test_plugins.py               | 26 +++++++++
 tests/test_templates/show_json.html |  8 +++
 9 files changed, 186 insertions(+), 19 deletions(-)
 create mode 100644 tests/test_templates/show_json.html

diff --git a/datasette/hookspecs.py b/datasette/hookspecs.py
index 42adaae8..3c6726b7 100644
--- a/datasette/hookspecs.py
+++ b/datasette/hookspecs.py
@@ -35,6 +35,11 @@ def extra_body_script(template, database, table, view_name, datasette):
     "Extra JavaScript code to be included in <script> at bottom of body"
 
 
+@hookspec
+def extra_template_vars(template, database, table, view_name, request, datasette):
+    "Extra template variables to be made available to the template - can return dict or callable or awaitable"
+
+
 @hookspec
 def publish_subcommand(publish):
     "Subcommands for 'datasette publish'"
diff --git a/datasette/views/base.py b/datasette/views/base.py
index 7acb7304..2c14be57 100644
--- a/datasette/views/base.py
+++ b/datasette/views/base.py
@@ -102,7 +102,7 @@ class BaseView(AsgiView):
     def database_color(self, database):
         return "ff0000"
 
-    def render(self, templates, **context):
+    async def render(self, templates, request, context):
         template = self.ds.jinja_env.select_template(templates)
         select_templates = [
             "{}{}".format("*" if template_name == template.name else "", template_name)
@@ -118,6 +118,26 @@ class BaseView(AsgiView):
             datasette=self.ds,
         ):
             body_scripts.append(jinja2.Markup(script))
+
+        extra_template_vars = {}
+        # pylint: disable=no-member
+        for extra_vars in pm.hook.extra_template_vars(
+            template=template.name,
+            database=context.get("database"),
+            table=context.get("table"),
+            view_name=self.name,
+            request=request,
+            datasette=self.ds,
+        ):
+            if callable(extra_vars):
+                extra_vars = extra_vars()
+            if asyncio.iscoroutine(extra_vars):
+                extra_vars = await extra_vars
+            assert isinstance(extra_vars, dict), "extra_vars is of type {}".format(
+                type(extra_vars)
+            )
+            extra_template_vars.update(extra_vars)
+
         return Response.html(
             template.render(
                 {
@@ -137,6 +157,7 @@ class BaseView(AsgiView):
                         "database_url": self.database_url,
                         "database_color": self.database_color,
                     },
+                    **extra_template_vars,
                 }
             )
         )
@@ -471,7 +492,7 @@ class DataView(BaseView):
             }
             if "metadata" not in context:
                 context["metadata"] = self.ds.metadata
-            r = self.render(templates, **context)
+            r = await self.render(templates, request=request, context=context)
             r.status = status_code
 
         ttl = request.args.get("_ttl", None)
diff --git a/datasette/views/index.py b/datasette/views/index.py
index 2c1c017a..fddb04d9 100644
--- a/datasette/views/index.py
+++ b/datasette/views/index.py
@@ -109,9 +109,12 @@ class IndexView(BaseView):
                 headers=headers,
             )
         else:
-            return self.render(
+            return await self.render(
                 ["index.html"],
-                databases=databases,
-                metadata=self.ds.metadata(),
-                datasette_version=__version__,
+                request=request,
+                context={
+                    "databases": databases,
+                    "metadata": self.ds.metadata(),
+                    "datasette_version": __version__,
+                },
             )
diff --git a/datasette/views/special.py b/datasette/views/special.py
index c4976bb2..45e948f6 100644
--- a/datasette/views/special.py
+++ b/datasette/views/special.py
@@ -24,4 +24,8 @@ class JsonDataView(BaseView):
             )
 
         else:
-            return self.render(["show_json.html"], filename=self.filename, data=data)
+            return await self.render(
+                ["show_json.html"],
+                request=request,
+                context={"filename": self.filename, "data": data},
+            )
diff --git a/docs/plugins.rst b/docs/plugins.rst
index 609fa844..3b3653cc 100644
--- a/docs/plugins.rst
+++ b/docs/plugins.rst
@@ -562,6 +562,8 @@ If the value matches that pattern, the plugin returns an HTML link element:
 extra_body_script(template, database, table, view_name, datasette)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
+Extra JavaScript to be added to a ``<script>`` block at the end of the ``<body>`` element on the page.
+
 ``template`` - string
     The template that is being rendered, e.g. ``database.html``
 
@@ -577,14 +579,74 @@ extra_body_script(template, database, table, view_name, datasette)
 ``datasette`` - Datasette instance
     You can use this to access plugin configuration options via ``datasette.plugin_config(your_plugin_name)``
 
-Extra JavaScript to be added to a ``<script>`` block at the end of the ``<body>`` element on the page.
-
 The ``template``, ``database`` and ``table`` options can be used to return different code depending on which template is being rendered and which database or table are being processed.
 
 The ``datasette`` instance is provided primarily so that you can consult any plugin configuration options that may have been set, using the ``datasette.plugin_config(plugin_name)`` method documented above.
 
 The string that you return from this function will be treated as "safe" for inclusion in a ``<script>`` block directly in the page, so it is up to you to apply any necessary escaping.
 
+
+.. _plugin_hook_extra_template_vars:
+
+extra_template_vars(template, database, table, view_name, request, datasette)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Extra template variables that should be made available in the rendered template context.
+
+``template`` - string
+    The template that is being rendered, e.g. ``database.html``
+
+``database`` - string or None
+    The name of the database, or ``None`` if the page does not correspond to a database (e.g. the root page)
+
+``table`` - string or None
+    The name of the table, or ``None`` if the page does not correct to a table
+
+``view_name`` - string
+    The name of the view being displayed. (`database`, `table`, and `row` are the most important ones.)
+
+``request`` - object
+    The current HTTP request object. ``request.scope`` provides access to the ASGI scope.
+
+``datasette`` - Datasette instance
+    You can use this to access plugin configuration options via ``datasette.plugin_config(your_plugin_name)``
+
+This hook can return one of three different types:
+
+Dictionary
+    If you return a dictionary its keys and values will be merged into the template context.
+
+Function that returns a dictionary
+    If you return a function it will be executed. If it returns a dictionary those values will will be merged into the template context.
+
+Function that returns an awaitable function that returns a dictionary
+    You can also return a function which returns an awaitable function which returns a dictionary. This means you can execute additional SQL queries using ``datasette.execute()``.
+
+Here's an example plugin that returns an authentication object from the ASGI scope:
+
+.. code-block:: python
+
+    @hookimpl
+    def extra_template_vars(request):
+        return {
+            "auth": request.scope.get("auth")
+        }
+
+And here's an example which returns the current version of SQLite:
+
+.. code-block:: python
+
+    @hookimpl
+    def extra_template_vars(datasette):
+        async def inner():
+            first_db = list(datasette.databases.keys())[0]
+            return {
+                "sqlite_version": (
+                    await datasette.execute(first_db, "select sqlite_version()")
+                ).rows[0][0]
+            }
+        return inner
+
 .. _plugin_register_output_renderer:
 
 register_output_renderer(datasette)
@@ -597,12 +659,12 @@ Allows the plugin to register a new output renderer, to output data in a custom
 
 .. code-block:: python
 
-        @hookimpl
-        def register_output_renderer(datasette):
-            return {
-                'extension': 'test',
-                'callback': render_test
-            }
+    @hookimpl
+    def register_output_renderer(datasette):
+        return {
+            'extension': 'test',
+            'callback': render_test
+        }
 
 This will register `render_test` to be called when paths with the extension `.test` (for example `/database.test`, `/database/table.test`, or `/database/table/row.test`) are requested. When a request is received, the callback function is called with three positional arguments:
 
@@ -630,10 +692,10 @@ A simple example of an output renderer callback function:
 
 .. code-block:: python
 
-        def render_test(args, data, view_name):
-            return {
-                'body': 'Hello World'
-            }
+    def render_test(args, data, view_name):
+        return {
+            'body': 'Hello World'
+        }
 
 .. _plugin_register_facet_classes:
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 6cc1f044..aaa8f735 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,3 +1,7 @@
+import os
+import pytest
+
+
 def pytest_configure(config):
     import sys
 
@@ -22,3 +26,14 @@ def move_to_front(items, test_name):
     test = [fn for fn in items if fn.name == test_name]
     if test:
         items.insert(0, items.pop(items.index(test[0])))
+
+
+@pytest.fixture
+def restore_working_directory(tmpdir, request):
+    previous_cwd = os.getcwd()
+    tmpdir.chdir()
+
+    def return_to_previous():
+        os.chdir(previous_cwd)
+
+    request.addfinalizer(return_to_previous)
diff --git a/tests/fixtures.py b/tests/fixtures.py
index db5f06e2..d686142b 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -376,6 +376,16 @@ def render_cell(value, column, table, database, datasette):
             table=table,
         )
     })
+
+
+@hookimpl
+def extra_template_vars(template, database, table, view_name, request, datasette):
+    return {
+        "extra_template_vars": json.dumps({
+            "template": template,
+            "scope_path": request.scope["path"]
+        }, default=lambda b: b.decode("utf8"))
+    }
 """
 
 PLUGIN2 = """
@@ -424,6 +434,19 @@ def render_cell(value, database):
     )
 
 
+@hookimpl
+def extra_template_vars(template, database, table, view_name, request, datasette):
+    async def inner():
+        return {
+            "extra_template_vars_from_awaitable": json.dumps({
+                "template": template,
+                "scope_path": request.scope["path"],
+                "awaitable": True,
+            }, default=lambda b: b.decode("utf8"))
+        }
+    return inner
+
+
 @hookimpl
 def asgi_wrapper(datasette):
     def wrap_with_databases_header(app):
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index 9af2a430..b1c7fd9a 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -3,6 +3,7 @@ from .fixtures import app_client, make_app_client, TEMP_PLUGIN_SECRET_FILE  # no
 import base64
 import json
 import os
+import pathlib
 import re
 import pytest
 import urllib
@@ -188,3 +189,28 @@ def test_plugins_extra_body_script(app_client, path, expected_extra_body_script)
 def test_plugins_asgi_wrapper(app_client):
     response = app_client.get("/fixtures")
     assert "fixtures" == response.headers["x-databases"]
+
+
+def test_plugins_extra_template_vars(restore_working_directory):
+    for client in make_app_client(
+        template_dir=str(pathlib.Path(__file__).parent / "test_templates")
+    ):
+        response = client.get("/-/metadata")
+        assert response.status == 200
+        extra_template_vars = json.loads(
+            Soup(response.body, "html.parser").select("pre.extra_template_vars")[0].text
+        )
+        assert {
+            "template": "show_json.html",
+            "scope_path": "/-/metadata",
+        } == extra_template_vars
+        extra_template_vars_from_awaitable = json.loads(
+            Soup(response.body, "html.parser")
+            .select("pre.extra_template_vars_from_awaitable")[0]
+            .text
+        )
+        assert {
+            "template": "show_json.html",
+            "awaitable": True,
+            "scope_path": "/-/metadata",
+        } == extra_template_vars_from_awaitable
diff --git a/tests/test_templates/show_json.html b/tests/test_templates/show_json.html
new file mode 100644
index 00000000..bbf1bc06
--- /dev/null
+++ b/tests/test_templates/show_json.html
@@ -0,0 +1,8 @@
+{% extends "base.html" %}
+
+{% block content %}
+{{ super() }}
+Test data for extra_template_vars:
+<pre class="extra_template_vars">{{ extra_template_vars|safe }}</pre>
+<pre class="extra_template_vars_from_awaitable">{{ extra_template_vars_from_awaitable|safe }}</pre>
+{% endblock %}