From e3639247cd48f7446c0a05f30042588ad05eb678 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 16 Aug 2020 09:50:23 -0700
Subject: [PATCH] Standard arguments for extra_ plugin hooks, closes #939

---
 datasette/app.py           |  50 ++++---
 datasette/hookspecs.py     |   6 +-
 docs/plugin_hooks.rst      | 258 ++++++++++++++++++++-----------------
 tests/plugins/my_plugin.py |  66 ++++++----
 tests/test_plugins.py      |  39 +++++-
 5 files changed, 258 insertions(+), 161 deletions(-)

diff --git a/datasette/app.py b/datasette/app.py
index fa273df0..180ba246 100644
--- a/datasette/app.py
+++ b/datasette/app.py
@@ -709,14 +709,19 @@ class Datasette:
             template = self.jinja_env.select_template(templates)
         body_scripts = []
         # pylint: disable=no-member
-        for script in pm.hook.extra_body_script(
+        for extra_script in pm.hook.extra_body_script(
             template=template.name,
             database=context.get("database"),
             table=context.get("table"),
             view_name=view_name,
+            request=request,
             datasette=self,
         ):
-            body_scripts.append(Markup(script))
+            if callable(extra_script):
+                extra_script = extra_script()
+            if asyncio.iscoroutine(extra_script):
+                extra_script = await extra_script
+            body_scripts.append(Markup(extra_script))
 
         extra_template_vars = {}
         # pylint: disable=no-member
@@ -748,8 +753,12 @@ class Datasette:
                 "body_scripts": body_scripts,
                 "format_bytes": format_bytes,
                 "show_messages": lambda: self._show_messages(request),
-                "extra_css_urls": self._asset_urls("extra_css_urls", template, context),
-                "extra_js_urls": self._asset_urls("extra_js_urls", template, context),
+                "extra_css_urls": await self._asset_urls(
+                    "extra_css_urls", template, context, request, view_name
+                ),
+                "extra_js_urls": await self._asset_urls(
+                    "extra_js_urls", template, context, request, view_name
+                ),
                 "base_url": self.config("base_url"),
                 "csrftoken": request.scope["csrftoken"] if request else lambda: "",
             },
@@ -762,20 +771,26 @@ class Datasette:
 
         return await template.render_async(template_context)
 
-    def _asset_urls(self, key, template, context):
+    async def _asset_urls(self, key, template, context, request, view_name):
         # Flatten list-of-lists from plugins:
         seen_urls = set()
-        for url_or_dict in itertools.chain(
-            itertools.chain.from_iterable(
-                getattr(pm.hook, key)(
-                    template=template.name,
-                    database=context.get("database"),
-                    table=context.get("table"),
-                    datasette=self,
-                )
-            ),
-            (self.metadata(key) or []),
+        collected = []
+        for hook in getattr(pm.hook, key)(
+            template=template.name,
+            database=context.get("database"),
+            table=context.get("table"),
+            datasette=self,
+            view_name=view_name,
+            request=request,
         ):
+            if callable(hook):
+                hook = hook()
+            if asyncio.iscoroutine(hook):
+                hook = await hook
+            collected.extend(hook)
+        collected.extend(self.metadata(key) or [])
+        output = []
+        for url_or_dict in collected:
             if isinstance(url_or_dict, dict):
                 url = url_or_dict["url"]
                 sri = url_or_dict.get("sri")
@@ -786,9 +801,10 @@ class Datasette:
                 continue
             seen_urls.add(url)
             if sri:
-                yield {"url": url, "sri": sri}
+                output.append({"url": url, "sri": sri})
             else:
-                yield {"url": url}
+                output.append({"url": url})
+        return output
 
     def app(self):
         "Returns an ASGI app function that serves the whole of Datasette"
diff --git a/datasette/hookspecs.py b/datasette/hookspecs.py
index 92d321b6..0e9c20cf 100644
--- a/datasette/hookspecs.py
+++ b/datasette/hookspecs.py
@@ -26,17 +26,17 @@ def prepare_jinja2_environment(env):
 
 
 @hookspec
-def extra_css_urls(template, database, table, datasette):
+def extra_css_urls(template, database, table, view_name, request, datasette):
     "Extra CSS URLs added by this plugin"
 
 
 @hookspec
-def extra_js_urls(template, database, table, datasette):
+def extra_js_urls(template, database, table, view_name, request, datasette):
     "Extra JavaScript URLs added by this plugin"
 
 
 @hookspec
-def extra_body_script(template, database, table, view_name, datasette):
+def extra_body_script(template, database, table, view_name, request, datasette):
     "Extra JavaScript code to be included in <script> at bottom of body"
 
 
diff --git a/docs/plugin_hooks.rst b/docs/plugin_hooks.rst
index c84a4097..426134fa 100644
--- a/docs/plugin_hooks.rst
+++ b/docs/plugin_hooks.rst
@@ -82,17 +82,23 @@ You can now use this filter in your custom templates like so::
 
 .. _plugin_hook_extra_css_urls:
 
-extra_css_urls(template, database, table, datasette)
-----------------------------------------------------
+extra_css_urls(template, database, table, view_name, request, datasette)
+------------------------------------------------------------------------
 
 ``template`` - string
     The template that is being rendered, e.g. ``database.html``
 
 ``database`` - string or None
-    The name of the database
+    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
+    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. (``index``, ``database``, ``table``, and ``row`` are the most important ones.)
+
+``request`` - object or None
+    The current HTTP :ref:`internals_request`. This can be ``None`` if the request object is not available.
 
 ``datasette`` - :ref:`internals_datasette`
     You can use this to access plugin configuration options via ``datasette.plugin_config(your_plugin_name)``
@@ -126,17 +132,32 @@ Or a list of dictionaries defining both a URL and an
             'sri': 'sha384-9gVQ4dYFwwWSjIDZnLEWnxCjeSWFphJiwGPXr1jddIhOegiu1FwO5qRGvFXOdJZ4',
         }]
 
+This function can also return an awaitable function, useful if it needs to run any async code:
+
+.. code-block:: python
+
+    from datasette import hookimpl
+
+    @hookimpl
+    def extra_css_urls(datasette):
+        async def inner():
+            db = datasette.get_database()
+            results = await db.execute("select url from css_files")
+            return [r[0] for r in results]
+
+        return inner
+
 Examples: `datasette-cluster-map <https://github.com/simonw/datasette-cluster-map>`_, `datasette-vega <https://github.com/simonw/datasette-vega>`_
 
 .. _plugin_hook_extra_js_urls:
 
-extra_js_urls(template, database, table, datasette)
----------------------------------------------------
+extra_js_urls(template, database, table, view_name, request, datasette)
+-----------------------------------------------------------------------
 
 Same arguments as ``extra_css_urls``.
 
 This works in the same way as ``extra_css_urls()`` but for JavaScript. You can
-return either a list of URLs or a list of dictionaries:
+return a list of URLs, a list of dictionaries or an awaitable function that returns those things:
 
 .. code-block:: python
 
@@ -164,6 +185,120 @@ you have one:
 
 Examples: `datasette-cluster-map <https://github.com/simonw/datasette-cluster-map>`_, `datasette-vega <https://github.com/simonw/datasette-vega>`_
 
+.. _plugin_hook_extra_body_script:
+
+extra_body_script(template, database, table, view_name, request, 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``
+
+``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. (``index``, ``database``, ``table``, and ``row`` are the most important ones.)
+
+``request`` - object or None
+    The current HTTP :ref:`internals_request`. This can be ``None`` if the request object is not available.
+
+``datasette`` - :ref:`internals_datasette`
+    You can use this to access plugin configuration options via ``datasette.plugin_config(your_plugin_name)``
+
+The ``template``, ``database``, ``table`` and ``view_name`` 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.
+
+You can also return an awaitable function that returns a string.
+
+Example: `datasette-cluster-map <https://github.com/simonw/datasette-cluster-map>`_
+
+.. _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. (``index``, ``database``, ``table``, and ``row`` are the most important ones.)
+
+``request`` - object or None
+    The current HTTP :ref:`internals_request`. This can be ``None`` if the request object is not available.
+
+``datasette`` - :ref:`internals_datasette`
+    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.
+
+Datasette runs Jinja2 in `async mode <https://jinja.palletsprojects.com/en/2.10.x/api/#async-support>`__, which means you can add awaitable functions to the template scope and they will be automatically awaited when they are rendered by the template.
+
+Here's an example plugin that adds a ``"user_agent"`` variable to the template context containing the current request's User-Agent header:
+
+.. code-block:: python
+
+    @hookimpl
+    def extra_template_vars(request):
+        return {
+            "user_agent": request.headers.get("user-agent")
+        }
+
+This example returns an awaitable function which adds a list of ``hidden_table_names`` to the context:
+
+.. code-block:: python
+
+    @hookimpl
+    def extra_template_vars(datasette, database):
+        async def hidden_table_names():
+            if database:
+                db = datasette.databases[database]
+                return {"hidden_table_names": await db.hidden_table_names()}
+            else:
+                return {}
+        return hidden_table_names
+
+And here's an example which adds a ``sql_first(sql_query)`` function which executes a SQL statement and returns the first column of the first row of results:
+
+.. code-block:: python
+
+    @hookimpl
+    def extra_template_vars(datasette, database):
+        async def sql_first(sql, dbname=None):
+            dbname = dbname or database or next(iter(datasette.databases.keys()))
+            return (await datasette.execute(dbname, sql)).rows[0][0]
+        return {"sql_first": sql_first}
+
+You can then use the new function in a template like so::
+
+    SQLite version: {{ sql_first("select sqlite_version()") }}
+
+Examples: `datasette-search-all <https://github.com/simonw/datasette-search-all>`_, `datasette-template-sql <https://github.com/simonw/datasette-template-sql>`_
+
 .. _plugin_hook_publish_subcommand:
 
 publish_subcommand(publish)
@@ -293,115 +428,6 @@ If the value matches that pattern, the plugin returns an HTML link element:
 
 Examples: `datasette-render-binary <https://github.com/simonw/datasette-render-binary>`_, `datasette-render-markdown <https://github.com/simonw/datasette-render-markdown>`_
 
-.. _plugin_hook_extra_body_script:
-
-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``
-
-``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. (`index`, `database`, `table`, and `row` are the most important ones.)
-
-``datasette`` - :ref:`internals_datasette`
-    You can use this to access plugin configuration options via ``datasette.plugin_config(your_plugin_name)``
-
-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.
-
-Example: `datasette-cluster-map <https://github.com/simonw/datasette-cluster-map>`_
-
-.. _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. (`index`, `database`, `table`, and `row` are the most important ones.)
-
-``request`` - object
-    The current HTTP :ref:`internals_request`.
-
-``datasette`` - :ref:`internals_datasette`
-    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.
-
-Datasette runs Jinja2 in `async mode <https://jinja.palletsprojects.com/en/2.10.x/api/#async-support>`__, which means you can add awaitable functions to the template scope and they will be automatically awaited when they are rendered by the template.
-
-Here's an example plugin that adds a ``"user_agent"`` variable to the template context containing the current request's User-Agent header:
-
-.. code-block:: python
-
-    @hookimpl
-    def extra_template_vars(request):
-        return {
-            "user_agent": request.headers.get("user-agent")
-        }
-
-This example returns an awaitable function which adds a list of ``hidden_table_names`` to the context:
-
-.. code-block:: python
-
-    @hookimpl
-    def extra_template_vars(datasette, database):
-        async def hidden_table_names():
-            if database:
-                db = datasette.databases[database]
-                return {"hidden_table_names": await db.hidden_table_names()}
-            else:
-                return {}
-        return hidden_table_names
-
-And here's an example which adds a ``sql_first(sql_query)`` function which executes a SQL statement and returns the first column of the first row of results:
-
-.. code-block:: python
-
-    @hookimpl
-    def extra_template_vars(datasette, database):
-        async def sql_first(sql, dbname=None):
-            dbname = dbname or database or next(iter(datasette.databases.keys()))
-            return (await datasette.execute(dbname, sql)).rows[0][0]
-        return {"sql_first": sql_first}
-
-You can then use the new function in a template like so::
-
-    SQLite version: {{ sql_first("select sqlite_version()") }}
-
-Examples: `datasette-search-all <https://github.com/simonw/datasette-search-all>`_, `datasette-template-sql <https://github.com/simonw/datasette-template-sql>`_
-
 .. _plugin_register_output_renderer:
 
 register_output_renderer(datasette)
diff --git a/tests/plugins/my_plugin.py b/tests/plugins/my_plugin.py
index 1870824f..0f6e80f3 100644
--- a/tests/plugins/my_plugin.py
+++ b/tests/plugins/my_plugin.py
@@ -26,16 +26,30 @@ def prepare_connection(conn, database, datasette):
 
 
 @hookimpl
-def extra_css_urls(template, database, table, datasette):
-    return [
-        "https://plugin-example.com/{}/extra-css-urls-demo.css".format(
-            base64.b64encode(
-                json.dumps(
-                    {"template": template, "database": database, "table": table,}
-                ).encode("utf8")
-            ).decode("utf8")
-        )
-    ]
+def extra_css_urls(template, database, table, view_name, request, datasette):
+    async def inner():
+        return [
+            "https://plugin-example.com/{}/extra-css-urls-demo.css".format(
+                base64.b64encode(
+                    json.dumps(
+                        {
+                            "template": template,
+                            "database": database,
+                            "table": table,
+                            "view_name": view_name,
+                            "request_path": request.path
+                            if request is not None
+                            else None,
+                            "added": (
+                                await datasette.get_database().execute("select 3 * 5")
+                            ).first()[0],
+                        }
+                    ).encode("utf8")
+                ).decode("utf8")
+            )
+        ]
+
+    return inner
 
 
 @hookimpl
@@ -47,19 +61,27 @@ def extra_js_urls():
 
 
 @hookimpl
-def extra_body_script(template, database, table, datasette):
-    return "var extra_body_script = {};".format(
-        json.dumps(
-            {
-                "template": template,
-                "database": database,
-                "table": table,
-                "config": datasette.plugin_config(
-                    "name-of-plugin", database=database, table=table,
-                ),
-            }
+def extra_body_script(template, database, table, view_name, request, datasette):
+    async def inner():
+        return "var extra_body_script = {};".format(
+            json.dumps(
+                {
+                    "template": template,
+                    "database": database,
+                    "table": table,
+                    "config": datasette.plugin_config(
+                        "name-of-plugin", database=database, table=table,
+                    ),
+                    "view_name": view_name,
+                    "request_path": request.path if request is not None else None,
+                    "added": (
+                        await datasette.get_database().execute("select 3 * 5")
+                    ).first()[0],
+                }
+            )
         )
-    )
+
+    return inner
 
 
 @hookimpl
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index c9fdf2e8..9831f0cf 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -56,14 +56,38 @@ def test_plugin_prepare_connection_arguments(app_client):
 @pytest.mark.parametrize(
     "path,expected_decoded_object",
     [
-        ("/", {"template": "index.html", "database": None, "table": None}),
+        (
+            "/",
+            {
+                "template": "index.html",
+                "database": None,
+                "table": None,
+                "view_name": "index",
+                "request_path": "/",
+                "added": 15,
+            },
+        ),
         (
             "/fixtures/",
-            {"template": "database.html", "database": "fixtures", "table": None},
+            {
+                "template": "database.html",
+                "database": "fixtures",
+                "table": None,
+                "view_name": "database",
+                "request_path": "/fixtures",
+                "added": 15,
+            },
         ),
         (
             "/fixtures/sortable",
-            {"template": "table.html", "database": "fixtures", "table": "sortable"},
+            {
+                "template": "table.html",
+                "database": "fixtures",
+                "table": "sortable",
+                "view_name": "table",
+                "request_path": "/fixtures/sortable",
+                "added": 15,
+            },
         ),
     ],
 )
@@ -207,6 +231,9 @@ def test_plugin_config_file(app_client):
                 "database": None,
                 "table": None,
                 "config": {"depth": "root"},
+                "view_name": "index",
+                "request_path": "/",
+                "added": 15,
             },
         ),
         (
@@ -216,6 +243,9 @@ def test_plugin_config_file(app_client):
                 "database": "fixtures",
                 "table": None,
                 "config": {"depth": "database"},
+                "view_name": "database",
+                "request_path": "/fixtures",
+                "added": 15,
             },
         ),
         (
@@ -225,6 +255,9 @@ def test_plugin_config_file(app_client):
                 "database": "fixtures",
                 "table": "sortable",
                 "config": {"depth": "table"},
+                "view_name": "table",
+                "request_path": "/fixtures/sortable",
+                "added": 15,
             },
         ),
     ],