diff --git a/docs/python-api.rst b/docs/python-api.rst index d515642..8717963 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -3347,6 +3347,26 @@ If you want to deliberately replace the registered function with a new implement def reverse_string(s): return s[::-1] +To remove a registered function, pass the callable to ``db.register_function()`` with ``deregister=True``: + +.. code-block:: python + + db.register_function(reverse_string, deregister=True) + +This removes the SQL function with the callable's name and number of arguments, and clears its registration from the ``Database`` instance. You can register it again without ``replace=True``: + +.. code-block:: python + + db.register_function(reverse_string) + +If the function was registered with ``name=``, use that same SQL name when removing it. SQL function names are case-insensitive for ASCII letters. Other names and overloads with a different number of arguments are left in place: + +.. code-block:: python + + db.register_function(reverse_string, name="rev", deregister=True) + +Removing a function that is not registered has no effect. Deregistration does not restore any earlier function that was replaced. SQLite raises an ``OperationalError`` if you try to invoke the removed function; the error message can differ between SQLite implementations. + Exceptions that occur inside a user-defined function default to returning the following error:: Unexpected error: user-defined function raised exception diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index c011d9b..f5b0e88 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -796,6 +796,7 @@ class Database: deterministic: bool = False, replace: bool = False, name: str | None = None, + deregister: bool = False, ) -> Callable[[Callable], Callable] | None: """ ``fn`` will be made available as a function within SQL, with the same name and number @@ -817,11 +818,28 @@ class Database: :param deterministic: set ``True`` for functions that always returns the same output for a given input :param replace: set ``True`` to replace an existing function with the same name - otherwise throw an error :param name: name of the SQLite function - if not specified, the Python function name will be used + :param deregister: set ``True`` to remove the function with this name and number of arguments """ def register(fn: Callable) -> Callable: fn_name = name or fn.__name__ # type: ignore arity = len(inspect.signature(fn).parameters) + if deregister: + self.conn.create_function(fn_name, arity, None) + # SQLite function names are case-insensitive for ASCII letters. + fold_ascii = str.maketrans( + "ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz" + ) + self._registered_functions.difference_update( + (registered_name, registered_arity) + for registered_name, registered_arity in tuple( + self._registered_functions + ) + if registered_arity == arity + and registered_name.translate(fold_ascii) + == fn_name.translate(fold_ascii) + ) + return fn if not replace and (fn_name, arity) in self._registered_functions: return fn kwargs: dict[str, bool] = {} diff --git a/tests/test_register_function.py b/tests/test_register_function.py index 63f0570..b7a5c95 100644 --- a/tests/test_register_function.py +++ b/tests/test_register_function.py @@ -104,3 +104,111 @@ def test_register_function_replace(fresh_db): return "two" assert "two" == fresh_db.execute("select one()").fetchone()[0] + + +@pytest.mark.parametrize("deterministic", (False, True)) +def test_deregister_function(fresh_db, deterministic): + calls = [] + + def temporary(value): + calls.append(value) + return value * 2 + + fresh_db.register_function(temporary, deterministic=deterministic) + assert fresh_db.execute("select temporary(3)").fetchone()[0] == 6 + assert fresh_db.register_function(temporary, deregister=True) is None + with pytest.raises(sqlite3.OperationalError): + fresh_db.execute("select temporary(4)").fetchone() + assert calls == [3] + + fresh_db.register_function(temporary) + assert fresh_db.execute("select temporary(5)").fetchone()[0] == 10 + assert calls == [3, 5] + + +@pytest.mark.parametrize("removal_name", ("custom_name", "CUSTOM_NAME")) +def test_deregister_function_custom_name(fresh_db, removal_name): + def temporary(value): + return value * 2 + + fresh_db.register_function(temporary) + fresh_db.register_function(temporary, name="custom_name") + fresh_db.register_function(temporary, name=removal_name, deregister=True) + with pytest.raises(sqlite3.OperationalError): + fresh_db.execute("select custom_name(3)").fetchone() + assert fresh_db.execute("select temporary(3)").fetchone()[0] == 6 + fresh_db.register_function(temporary, name="custom_name") + assert fresh_db.execute("select custom_name(4)").fetchone()[0] == 8 + + +def test_deregister_function_preserves_other_arity(fresh_db): + def one_arg(value): + return value + 1 + + def two_args(left, right): + return left + right + + fresh_db.register_function(one_arg, name="temporary") + fresh_db.register_function(two_args, name="temporary") + fresh_db.register_function(one_arg, name="temporary", deregister=True) + with pytest.raises(sqlite3.OperationalError): + fresh_db.execute("select temporary(1)").fetchone() + assert fresh_db.execute("select temporary(1, 2)").fetchone()[0] == 3 + + +def test_deregister_function_preserves_non_ascii_names(fresh_db): + def temporary(value): + return value * 2 + + fresh_db.register_function(temporary, name="FÜ") + fresh_db.register_function(temporary, name="fü") + fresh_db.register_function(temporary, name="fÜ", deregister=True) + with pytest.raises(sqlite3.OperationalError): + fresh_db.execute('select "FÜ"(1)').fetchone() + assert fresh_db.execute('select "fü"(2)').fetchone()[0] == 4 + # A surviving registration must still suppress replacement by default. + fresh_db.register_function(lambda value: 0, name="fü") + assert fresh_db.execute('select "fü"(2)').fetchone()[0] == 4 + + +def test_deregister_function_idempotent(fresh_db): + def temporary(): + return 1 + + # Removing an unregistered function or removing it twice is harmless. + fresh_db.register_function(temporary, deregister=True) + fresh_db.register_function(temporary) + fresh_db.register_function(temporary, deregister=True) + fresh_db.register_function(temporary, deregister=True) + fresh_db.register_function(temporary) + assert fresh_db.execute("select temporary()").fetchone()[0] == 1 + + +def test_deregister_function_failure_preserves_cache(fresh_db): + def temporary(value): + return value * 2 + + fresh_db.register_function(temporary) + original_conn = fresh_db.conn + fresh_db.conn = MagicMock() + fresh_db.conn.create_function.side_effect = sqlite3.OperationalError("busy") + try: + with pytest.raises(sqlite3.OperationalError, match="busy"): + fresh_db.register_function(temporary, deregister=True) + fresh_db.conn.create_function.assert_called_once_with("temporary", 1, None) + finally: + fresh_db.conn = original_conn + fresh_db.register_function(lambda value: 0, name="temporary") + assert fresh_db.execute("select temporary(3)").fetchone()[0] == 6 + + +def test_deregister_function_decorator(fresh_db): + def temporary(): + return 1 + + fresh_db.register_function(temporary) + assert fresh_db.register_function(deregister=True)(temporary) is temporary + with pytest.raises(sqlite3.OperationalError): + fresh_db.execute("select temporary()").fetchone() + fresh_db.register_function()(temporary) + assert fresh_db.execute("select temporary()").fetchone()[0] == 1