From 3f80a026983d3e634f05a46f2a6da162b5139dd9 Mon Sep 17 00:00:00 2001
From: Alex Garcia <alexsebastian.garcia@gmail.com>
Date: Sat, 22 Jul 2023 15:59:08 -0700
Subject: [PATCH] prepare_connection plugin hook

Closes:
- #574

Refs #567

---------

Co-authored-by: Simon Willison <swillison@gmail.com>
---
 docs/plugins.rst          | 28 ++++++++++++++++++++++++-
 sqlite_utils/__init__.py  |  2 +-
 sqlite_utils/db.py        |  3 +++
 sqlite_utils/hookspecs.py |  5 +++++
 tests/test_plugins.py     | 44 ++++++++++++++++++++++++++++++++++++++-
 5 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/docs/plugins.rst b/docs/plugins.rst
index 924e7d5..83b6688 100644
--- a/docs/plugins.rst
+++ b/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!"
diff --git a/sqlite_utils/__init__.py b/sqlite_utils/__init__.py
index a55c2d2..b8046f6 100644
--- a/sqlite_utils/__init__.py
+++ b/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"]
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 663a184..945458f 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -37,6 +37,7 @@ from typing import (
     Tuple,
 )
 import uuid
+from sqlite_utils.plugins import pm
 
 try:
     from sqlite_dump import iterdump
@@ -342,6 +343,8 @@ class Database:
         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()
diff --git a/sqlite_utils/hookspecs.py b/sqlite_utils/hookspecs.py
index cf05f7d..83466be 100644
--- a/sqlite_utils/hookspecs.py
+++ b/sqlite_utils/hookspecs.py
@@ -8,3 +8,8 @@ hookimpl = HookimplMarker("sqlite_utils")
 @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"""
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index dcb9e7b..10c4bb5 100644
--- a/tests/test_plugins.py
+++ b/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 test_register_commands():
         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() == []