prepare_connection plugin hook

Closes:
- #574

Refs #567

---------

Co-authored-by: Simon Willison <swillison@gmail.com>

Author: asg017

docs/plugins.rst

@@ -86,7 +86,7 @@ See the `LLM plugin documentation <https://llm.datasette.io/en/stable/plugins/tu
 Plugin hooks
 ------------
 
-Plugin hooks allow ``sqlite-utils`` to be customized. There is currently one hook.
+Plugin hooks allow ``sqlite-utils`` to be customized.
 
 .. _plugins_hooks_register_commands:
 
@@ -109,3 +109,29 @@ Example implementation:
             "Say hello world"
             click.echo("Hello world!")
 
+.. _plugins_hooks_prepare_connection:
+
+prepare_connection(conn)
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+This hook is called when a new SQLite database connection is created. You can
+use it to `register custom SQL functions <https://docs.python.org/2/library/sqlite3.html#sqlite3.Connection.create_function>`_,
+aggregates and collations. For example:
+
+.. code-block:: python
+
+    import click
+    import sqlite_utils
+
+    @sqlite_utils.hookimpl
+    def prepare_connection(conn):
+        conn.create_function(
+          "hello", 1, lambda name: f"Hello, {name}!"
+        )
+
+This registers a SQL function called ``hello`` which takes a single
+argument and can be called like this:
+
+.. code-block:: sql
+
+    select hello("world"); -- "Hello, world!"

sqlite_utils/__init__.py

@@ -1,6 +1,6 @@
-from .db import Database
 from .utils import suggest_column_types
 from .hookspecs import hookimpl
 from .hookspecs import hookspec
+from .db import Database
 
 __all__ = ["Database", "suggest_column_types", "hookimpl", "hookspec"]

sqlite_utils/db.py

@@ -37,6 +37,7 @@
     Tuple,
 )
 import uuid
+from sqlite_utils.plugins import pm
 
 try:
     from sqlite_dump import iterdump
@@ -342,6 +343,8 @@ def __init__(
         self._registered_functions: set = set()
         self.use_counts_table = use_counts_table
 
+        pm.hook.prepare_connection(conn=self.conn)
+
     def close(self):
         "Close the SQLite connection, and the underlying database file"
         self.conn.close()

sqlite_utils/hookspecs.py

@@ -8,3 +8,8 @@
 @hookspec
 def register_commands(cli):
     """Register additional CLI commands, e.g. 'sqlite-utils mycommand ...'"""
+
+
+@hookspec
+def prepare_connection(conn):
+    """Modify SQLite connection in some way e.g. register custom SQL functions"""

tests/test_plugins.py

@@ -1,7 +1,7 @@
 from click.testing import CliRunner
 import click
 import importlib
-from sqlite_utils import cli, hookimpl, plugins
+from sqlite_utils import cli, Database, hookimpl, plugins
 
 
 def test_register_commands():
@@ -35,3 +35,45 @@ def hello_world():
         plugins.pm.unregister(name="HelloWorldPlugin")
         importlib.reload(cli)
         assert plugins.get_plugins() == []
+
+
+def test_prepare_connection():
+    importlib.reload(cli)
+    assert plugins.get_plugins() == []
+
+    class HelloFunctionPlugin:
+        __name__ = "HelloFunctionPlugin"
+
+        @hookimpl
+        def prepare_connection(self, conn):
+            conn.create_function("hello", 1, lambda name: f"Hello, {name}!")
+
+    db = Database(memory=True)
+    functions = db.execute(
+        "select distinct name from pragma_function_list order by 1"
+    ).fetchall()
+    assert "hello" not in functions
+
+    try:
+        plugins.pm.register(HelloFunctionPlugin(), name="HelloFunctionPlugin")
+
+        assert plugins.get_plugins() == [
+            {"name": "HelloFunctionPlugin", "hooks": ["prepare_connection"]}
+        ]
+
+        db = Database(memory=True)
+
+        functions = [
+            row[0]
+            for row in db.execute(
+                "select distinct name from pragma_function_list order by 1"
+            ).fetchall()
+        ]
+        assert "hello" in functions
+
+        result = db.execute('select hello("world")').fetchone()[0]
+        assert result == "Hello, world!"
+
+    finally:
+        plugins.pm.unregister(name="HelloFunctionPlugin")
+        assert plugins.get_plugins() == []