sqlite-utils query --functions option, refs #471

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 <TEXT TEXT>...  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.
 

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
 -----------------
 

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
     )

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 (<string>, 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"