prepare_connection plugin hook

Closes:
- #574

Refs #567

---------

Co-authored-by: Simon Willison <swillison@gmail.com>
This commit is contained in:
Alex Garcia 2023-07-22 15:59:08 -07:00 committed by GitHub
commit 3f80a02698
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 79 additions and 3 deletions

View file

@ -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() == []