Merge branch 'main' into create-transform

tests/ext.c

@@ -0,0 +1,48 @@
+/*
+** This file implements a SQLite extension with multiple entrypoints.
+**
+** The default entrypoint, sqlite3_ext_init, has a single function "a".
+** The 1st alternate entrypoint, sqlite3_ext_b_init, has a single function "b".
+** The 2nd alternate entrypoint, sqlite3_ext_c_init, has a single function "c".
+**
+** Compiling instructions:
+**     https://www.sqlite.org/loadext.html#compiling_a_loadable_extension
+**
+*/
+
+#include "sqlite3ext.h"
+
+SQLITE_EXTENSION_INIT1
+
+// SQL function that returns back the value supplied during sqlite3_create_function()
+static void func(sqlite3_context *context, int argc, sqlite3_value **argv) {
+  sqlite3_result_text(context, (char *) sqlite3_user_data(context), -1, SQLITE_STATIC);
+}
+
+
+// The default entrypoint, since it matches the "ext.dylib"/"ext.so" name
+#ifdef _WIN32
+__declspec(dllexport)
+#endif
+int sqlite3_ext_init(sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi) {
+  SQLITE_EXTENSION_INIT2(pApi);
+  return sqlite3_create_function(db, "a", 0, 0, "a", func, 0, 0);
+}
+
+// Alternate entrypoint #1
+#ifdef _WIN32
+__declspec(dllexport)
+#endif
+int sqlite3_ext_b_init(sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi) {
+  SQLITE_EXTENSION_INIT2(pApi);
+  return sqlite3_create_function(db, "b", 0, 0, "b", func, 0, 0);
+}
+
+// Alternate entrypoint #2
+#ifdef _WIN32
+__declspec(dllexport)
+#endif
+int sqlite3_ext_c_init(sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi) {
+  SQLITE_EXTENSION_INIT2(pApi);
+  return sqlite3_create_function(db, "c", 0, 0, "c", func, 0, 0);
+}

tests/test_cli.py

@@ -1,6 +1,7 @@
 from sqlite_utils import cli, Database
 from sqlite_utils.db import Index, ForeignKey
 from click.testing import CliRunner
+from pathlib import Path
 import subprocess
 import sys
 from unittest import mock
@@ -12,6 +13,26 @@ import textwrap
 from .utils import collapse_whitespace
 
 
+def _supports_pragma_function_list():
+    db = Database(memory=True)
+    try:
+        db.execute("select * from pragma_function_list()")
+    except Exception:
+        return False
+    return True
+
+
+def _has_compiled_ext():
+    for ext in ["dylib", "so", "dll"]:
+        path = Path(__file__).parent / f"ext.{ext}"
+        if path.is_file():
+            return True
+    return False
+
+
+COMPILED_EXTENSION_PATH = str(Path(__file__).parent / "ext")
+
+
 @pytest.mark.parametrize(
     "options",
     (
@@ -683,6 +704,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 +726,76 @@ 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}
+    ]
+
+
+@pytest.mark.skipif(
+    not _supports_pragma_function_list(),
+    reason="Needs SQLite version that supports pragma_function_list()",
+)
+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"
@@ -881,6 +972,15 @@ def test_query_memory_does_not_create_file(tmpdir):
             ["-c", "name", "--where", "id = :id", "--param", "id", "1"],
             '[{"name": "Cleo"}]',
         ),
+        # --order
+        (
+            ["-c", "id", "--order", "id desc", "--limit", "1"],
+            '[{"id": 2}]',
+        ),
+        (
+            ["-c", "id", "--order", "id", "--limit", "1"],
+            '[{"id": 1}]',
+        ),
     ],
 )
 def test_rows(db_path, args, expected):
@@ -2196,3 +2296,32 @@ def test_duplicate_table(tmpdir):
     assert result.exit_code == 0
     assert db["one"].columns_dict == db["two"].columns_dict
     assert list(db["one"].rows) == list(db["two"].rows)
+
+
+@pytest.mark.skipif(not _has_compiled_ext(), reason="Requires compiled ext.c")
+@pytest.mark.parametrize(
+    "entrypoint,should_pass,should_fail",
+    (
+        (None, ("a",), ("b", "c")),
+        ("sqlite3_ext_b_init", ("b"), ("a", "c")),
+        ("sqlite3_ext_c_init", ("c"), ("a", "b")),
+    ),
+)
+def test_load_extension(entrypoint, should_pass, should_fail):
+    ext = COMPILED_EXTENSION_PATH
+    if entrypoint:
+        ext += ":" + entrypoint
+    for func in should_pass:
+        result = CliRunner().invoke(
+            cli.cli,
+            ["memory", "select {}()".format(func), "--load-extension", ext],
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+    for func in should_fail:
+        result = CliRunner().invoke(
+            cli.cli,
+            ["memory", "select {}()".format(func), "--load-extension", ext],
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 1

tests/test_cli_bulk.py

@@ -28,9 +28,11 @@ def test_cli_bulk(test_db_and_path):
         [
             "bulk",
             db_path,
-            "insert into example (id, name) values (:id, :name)",
+            "insert into example (id, name) values (:id, myupper(:name))",
             "-",
             "--nl",
+            "--functions",
+            "myupper = lambda s: s.upper()",
         ],
         input='{"id": 3, "name": "Three"}\n{"id": 4, "name": "Four"}\n',
     )
@@ -38,8 +40,8 @@ def test_cli_bulk(test_db_and_path):
     assert [
         {"id": 1, "name": "One"},
         {"id": 2, "name": "Two"},
-        {"id": 3, "name": "Three"},
-        {"id": 4, "name": "Four"},
+        {"id": 3, "name": "THREE"},
+        {"id": 4, "name": "FOUR"},
     ] == list(db["example"].rows)
 
 

tests/test_cli_convert.py

@@ -606,3 +606,23 @@ def test_convert_hyphen_workaround(fresh_db_and_path):
     assert list(db["names"].rows) == [
         {"id": 1, "name": "-"},
     ]
+
+
+def test_convert_initialization_pattern(fresh_db_and_path):
+    db, db_path = fresh_db_and_path
+    db["names"].insert_all([{"id": 1, "name": "Cleo"}], pk="id")
+    result = CliRunner().invoke(
+        cli.cli,
+        [
+            "convert",
+            db_path,
+            "names",
+            "name",
+            "-",
+        ],
+        input="import random\nrandom.seed(1)\ndef convert(value):    return random.randint(0, 100)",
+    )
+    assert 0 == result.exit_code, result.output
+    assert list(db["names"].rows) == [
+        {"id": 1, "name": "17"},
+    ]

tests/test_cli_memory.py

@@ -289,3 +289,12 @@ def test_memory_two_files_with_same_stem(tmpdir):
         ");\n"
         "CREATE VIEW t2 AS select * from [data_2];\n"
     )
+
+
+def test_memory_functions():
+    result = CliRunner().invoke(
+        cli.cli,
+        ["memory", "select hello()", "--functions", "hello = lambda: 'Hello'"],
+    )
+    assert result.exit_code == 0
+    assert result.output.strip() == '[{"hello()": "Hello"}]'