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

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"