@db.register_function(deterministic=True), closes #191

This commit is contained in:
Simon Willison 2020-10-28 14:24:03 -07:00
commit 0789bad8f7
3 changed files with 55 additions and 4 deletions

View file

@ -1703,3 +1703,13 @@ SQLite supports registering custom SQL functions written in Python. The ``@db.re
# This prints "olleh" # This prints "olleh"
If your Python function accepts multiple arguments the SQL version will accept the same number of arguments. If your Python function accepts multiple arguments the SQL version will accept the same number of arguments.
Python 3.8 added the ability to register `deterministic SQLite functions <https://sqlite.org/deterministic.html>`__, allowing you to indicate that a function will return the exact same result for any given inputs and hence allowing SQLite to apply some performance optimizations. You can mark a function as deterministic using ``deterministic=True``, like this:
.. code-block:: python
@db.register_function(deterministic=True)
def reverse_string(s):
return "".join(reversed(list(s)))
If you run this on a version of Python prior to 3.8 your code will still work, but the ``deterministic=True`` parameter will be ignored.

View file

@ -10,6 +10,7 @@ import itertools
import json import json
import os import os
import pathlib import pathlib
import sys
import textwrap import textwrap
import uuid import uuid
@ -149,10 +150,20 @@ class Database:
def __repr__(self): def __repr__(self):
return "<Database {}>".format(self.conn) return "<Database {}>".format(self.conn)
def register_function(self, fn): def register_function(self, fn=None, deterministic=None):
name = fn.__name__ def register(fn):
arity = len(inspect.signature(fn).parameters) name = fn.__name__
self.conn.create_function(name, arity, fn) arity = len(inspect.signature(fn).parameters)
kwargs = {}
if deterministic and sys.version_info >= (3, 8):
kwargs["deterministic"] = True
self.conn.create_function(name, arity, fn, **kwargs)
return fn
if fn is None:
return register
else:
register(fn)
def execute(self, sql, parameters=None): def execute(self, sql, parameters=None):
if self._tracer: if self._tracer:

View file

@ -1,3 +1,8 @@
import pytest
import sys
from unittest.mock import MagicMock
def test_register_function(fresh_db): def test_register_function(fresh_db):
@fresh_db.register_function @fresh_db.register_function
def reverse_string(s): def reverse_string(s):
@ -14,3 +19,28 @@ def test_register_function_multiple_arguments(fresh_db):
result = fresh_db.execute("select a_times_b_plus_c(2, 3, 4)").fetchone()[0] result = fresh_db.execute("select a_times_b_plus_c(2, 3, 4)").fetchone()[0]
assert result == 10 assert result == 10
def test_register_function_deterministic(fresh_db):
@fresh_db.register_function(deterministic=True)
def to_lower(s):
return s.lower()
result = fresh_db.execute("select to_lower('BOB')").fetchone()[0]
assert result == "bob"
@pytest.mark.skipif(
sys.version_info < (3, 8), reason="deterministic=True was added in Python 3.8"
)
def test_register_function_deterministic_registered(fresh_db):
fresh_db.conn = MagicMock()
fresh_db.conn.create_function = MagicMock()
@fresh_db.register_function(deterministic=True)
def to_lower_2(s):
return s.lower()
fresh_db.conn.create_function.assert_called_with(
"to_lower_2", 1, to_lower_2, deterministic=True
)