diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index c0f794c..af059f1 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -119,6 +119,8 @@ See :ref:`cli_query`. escaped strings -r, --raw Raw output, first column of first row -p, --param ... Named :parameters for SQL query + --functions TEXT Python code defining one or more custom SQL + functions --load-extension TEXT SQLite extensions to load -h, --help Show this message and exit. diff --git a/docs/cli.rst b/docs/cli.rst index 8cb2c56..1632e2f 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -250,6 +250,26 @@ If you execute an ``UPDATE``, ``INSERT`` or ``DELETE`` query the command will re $ sqlite-utils dogs.db "update dogs set age = 5 where name = 'Cleo'" [{"rows_affected": 1}] +.. _cli_query_functions: + +Defining custom SQL functions +----------------------------- + +You can use the ``--functions`` option to pass a block of Python code that defines additional functions which can then be called by your SQL query. + +This example defines a function which extracts the domain from a URL:: + + $ sqlite-utils query dogs.db "select url, domain(url) from urls" --functions ' + from urllib.parse import urlparse + + def domain(url): + return urlparse(url).netloc + ' + +Every callable object defined in the block will be registered as a SQL function with the same name, with the exception of functions with names that begin with an underscore. + +.. _cli_query_extensions: + SQLite extensions ----------------- diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 43e76fa..1ac8d78 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1633,6 +1633,9 @@ def drop_view(path, view, ignore, load_extension): type=(str, str), help="Named :parameters for SQL query", ) +@click.option( + "--functions", help="Python code defining one or more custom SQL functions" +) @load_extension_option def query( path, @@ -1649,6 +1652,7 @@ def query( raw, param, load_extension, + functions, ): """Execute SQL query and return the results as JSON @@ -1665,6 +1669,19 @@ def query( _load_extensions(db, load_extension) db.register_fts4_bm25() + # Register any Python functions as SQL functions: + if functions: + sqlite3.enable_callback_tracebacks(True) + globals = {} + try: + exec(functions, globals) + except SyntaxError as ex: + raise click.ClickException("Error in functions definition: {}".format(ex)) + # Register all callables in the locals dict: + for name, value in globals.items(): + if callable(value) and not name.startswith("_"): + db.register_function(value, name=name) + _execute_query( db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 4eadb88..057f129 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -683,6 +683,11 @@ _one_query = "select id, name, age from dogs where id = 1" (_one_query, ["--nl"], '{"id": 1, "name": "Cleo", "age": 4}'), (_one_query, ["--arrays"], '[[1, "Cleo", 4]]'), (_one_query, ["--arrays", "--nl"], '[1, "Cleo", 4]'), + ( + "select id, dog(age) from dogs", + ["--functions", "def dog(i):\n return i * 7"], + '[{"id": 1, "dog(age)": 28},\n {"id": 2, "dog(age)": 14}]', + ), ], ) def test_query_json(db_path, sql, args, expected): @@ -700,11 +705,72 @@ def test_query_json(db_path, sql, args, expected): def test_query_json_empty(db_path): result = CliRunner().invoke( - cli.cli, [db_path, "select * from sqlite_master where 0"] + cli.cli, + [db_path, "select * from sqlite_master where 0"], ) assert result.output.strip() == "[]" +def test_query_invalid_function(db_path): + result = CliRunner().invoke( + cli.cli, [db_path, "select bad()", "--functions", "def invalid_python"] + ) + assert result.exit_code == 1 + assert ( + result.output.strip() + == "Error: Error in functions definition: invalid syntax (, line 1)" + ) + + +TEST_FUNCTIONS = """ +def zero(): + return 0 + +def one(a): + return a + +def _two(a, b): + return a + b + +def two(a, b): + return _two(a, b) +""" + + +def test_query_complex_function(db_path): + result = CliRunner().invoke( + cli.cli, + [ + db_path, + "select zero(), one(1), two(1, 2)", + "--functions", + TEST_FUNCTIONS, + ], + ) + assert result.exit_code == 0 + assert json.loads(result.output.strip()) == [ + {"zero()": 0, "one(1)": 1, "two(1, 2)": 3} + ] + + +def test_hidden_functions_are_hidden(db_path): + result = CliRunner().invoke( + cli.cli, + [ + db_path, + "select name from pragma_function_list()", + "--functions", + TEST_FUNCTIONS, + ], + ) + assert result.exit_code == 0 + functions = {r["name"] for r in json.loads(result.output.strip())} + assert "zero" in functions + assert "one" in functions + assert "two" in functions + assert "_two" not in functions + + LOREM_IPSUM_COMPRESSED = ( b"x\x9c\xed\xd1\xcdq\x03!\x0c\x05\xe0\xbb\xabP\x01\x1eW\x91\xdc|M\x01\n\xc8\x8e" b"f\xf83H\x1e\x97\x1f\x91M\x8e\xe9\xe0\xdd\x96\x05\x84\xf4\xbek\x9fRI\xc7\xf2J"