Load extension entrypoint support, refs #470

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
@@ -21,6 +22,17 @@ def _supports_pragma_function_list():
     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",
     (
@@ -2284,3 +2296,33 @@ 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.asyncio
+@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")),
+    ),
+)
+async 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