From 19dd077944429c1365b513d80cc71c605ae3bed3 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Fri, 26 Aug 2022 22:55:47 -0700 Subject: [PATCH] Support entrypoints for `--load-extension` (#473) * Entrypoint support, closes #470 --- .github/workflows/test.yml | 4 ++ .gitignore | 3 ++ docs/cli-reference.rst | 82 ++++++++++++++++++++------------------ sqlite_utils/cli.py | 8 +++- tests/ext.c | 48 ++++++++++++++++++++++ tests/test_cli.py | 41 +++++++++++++++++++ 6 files changed, 145 insertions(+), 41 deletions(-) create mode 100644 tests/ext.c diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 75ef594..9303965 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -35,6 +35,10 @@ jobs: - name: Install SpatiaLite if: matrix.os == 'ubuntu-latest' run: sudo apt-get install libsqlite3-mod-spatialite + - name: Build extension for --load-extension test + if: matrix.os == 'ubuntu-latest' + run: |- + (cd tests && gcc ext.c -fPIC -shared -o ext.so && ls -lah) - name: Run tests run: | pytest -v diff --git a/.gitignore b/.gitignore index 32032d5..7947f33 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,6 @@ venv Pipfile Pipfile.lock pyproject.toml +tests/*.dylib +tests/*.so +tests/*.dll diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index c194967..d3951fd 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -121,7 +121,8 @@ See :ref:`cli_query`. -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 + --load-extension TEXT Path to SQLite extension, with optional + :entrypoint -h, --help Show this message and exit. @@ -189,7 +190,8 @@ See :ref:`cli_memory`. --dump Dump SQL for in-memory database --save FILE Save in-memory database to this file --analyze Analyze resulting tables and output results - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional + :entrypoint -h, --help Show this message and exit. @@ -266,7 +268,7 @@ See :ref:`cli_inserting_data`, :ref:`cli_insert_csv_tsv`, :ref:`cli_insert_unstr --default ... Default value that should be set for a column -d, --detect-types Detect types for columns in CSV/TSV data --analyze Run ANALYZE at the end of this operation - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint --silent Do not show progress bar --ignore Ignore records if pk already exists --replace Replace records if pk already exists @@ -320,7 +322,7 @@ See :ref:`cli_upsert`. --default ... Default value that should be set for a column -d, --detect-types Detect types for columns in CSV/TSV data --analyze Run ANALYZE at the end of this operation - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint --silent Do not show progress bar -h, --help Show this message and exit. @@ -364,7 +366,7 @@ See :ref:`cli_bulk`. --sniff Detect delimiter and quote character --no-headers CSV file has no header row --encoding TEXT Character encoding for input, defaults to utf-8 - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -404,7 +406,7 @@ See :ref:`cli_search`. textile, tsv, unsafehtml, youtrack --json-cols Detect JSON cols and output them as JSON, not escaped strings - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -440,7 +442,7 @@ See :ref:`cli_transform_table`. --default-none TEXT Remove default from this column --drop-foreign-key TEXT Drop foreign key constraint for this column --sql Output SQL without executing it - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -465,7 +467,7 @@ See :ref:`cli_extract`. --table TEXT Name of the other table to extract columns to --fk-column TEXT Name of the foreign key column to add to the table --rename ... Rename this column in extracted table - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -487,7 +489,7 @@ See :ref:`cli_schema`. sqlite-utils schema trees.db Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -525,7 +527,7 @@ See :ref:`cli_insert_files`. --text Store file content as TEXT, not BLOB --encoding TEXT Character encoding for input, defaults to utf-8 -s, --silent Don't show a progress bar - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -549,7 +551,7 @@ See :ref:`cli_analyze_tables`. Options: -c, --column TEXT Specific columns to analyze --save Save results to _analyze_tables table - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -656,7 +658,7 @@ See :ref:`cli_tables`. strings --columns Include list of columns for each table --schema Include schema for each table - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -694,7 +696,7 @@ See :ref:`cli_views`. strings --columns Include list of columns for each view --schema Include schema for each view - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -735,7 +737,8 @@ See :ref:`cli_rows`. simple, textile, tsv, unsafehtml, youtrack --json-cols Detect JSON cols and output them as JSON, not escaped strings - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional + :entrypoint -h, --help Show this message and exit. @@ -770,7 +773,7 @@ See :ref:`cli_triggers`. textile, tsv, unsafehtml, youtrack --json-cols Detect JSON cols and output them as JSON, not escaped strings - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -806,7 +809,7 @@ See :ref:`cli_indexes`. textile, tsv, unsafehtml, youtrack --json-cols Detect JSON cols and output them as JSON, not escaped strings - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -830,7 +833,7 @@ See :ref:`cli_create_database`. Options: --enable-wal Enable WAL mode on the created database --init-spatialite Enable SpatiaLite on the created database - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -862,7 +865,7 @@ See :ref:`cli_create_table`. foreign key --ignore If table already exists, do nothing --replace If table already exists, replace it - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -892,7 +895,7 @@ See :ref:`cli_create_index`. --unique Make this a unique index --if-not-exists, --ignore Ignore if index already exists --analyze Run ANALYZE after creating the index - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -920,7 +923,7 @@ See :ref:`cli_fts`. --create-triggers Create triggers to update the FTS tables when the parent table changes. --replace Replace existing FTS configuration if it exists - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -940,7 +943,7 @@ populate-fts sqlite-utils populate-fts chickens.db chickens name Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -960,7 +963,7 @@ rebuild-fts sqlite-utils rebuild-fts chickens.db chickens Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -980,7 +983,7 @@ disable-fts sqlite-utils disable-fts chickens.db chickens Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1004,7 +1007,7 @@ See :ref:`cli_optimize`. Options: --no-vacuum Don't run VACUUM - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1069,7 +1072,7 @@ See :ref:`cli_dump`. sqlite-utils dump chickens.db Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1097,7 +1100,7 @@ See :ref:`cli_add_column`. omitted will automatically use the primary key --not-null-default TEXT Add NOT NULL DEFAULT 'TEXT' constraint --ignore If column already exists, do nothing - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1123,7 +1126,7 @@ See :ref:`cli_add_foreign_key`. Options: --ignore If foreign key already exists, do nothing - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1147,7 +1150,7 @@ See :ref:`cli_add_foreign_keys`. authors country_id countries id Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1169,7 +1172,7 @@ See :ref:`cli_index_foreign_keys`. sqlite-utils index-foreign-keys chickens.db Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1191,7 +1194,7 @@ See :ref:`cli_wal`. sqlite-utils enable-wal chickens.db Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1211,7 +1214,7 @@ disable-wal sqlite-utils disable-wal chickens.db Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1233,7 +1236,7 @@ See :ref:`cli_enable_counts`. sqlite-utils enable-counts chickens.db Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1253,7 +1256,7 @@ reset-counts sqlite-utils reset-counts chickens.db Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1270,7 +1273,7 @@ duplicate Options: --ignore If table does not exist, do nothing - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1293,7 +1296,7 @@ See :ref:`cli_drop_table`. Options: --ignore If table does not exist, do nothing - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1318,7 +1321,7 @@ See :ref:`cli_create_view`. Options: --ignore If view already exists, do nothing --replace If view already exists, replace it - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1341,7 +1344,7 @@ See :ref:`cli_drop_view`. Options: --ignore If view does not exist, do nothing - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -1372,7 +1375,8 @@ See :ref:`cli_spatialite`. --dimensions TEXT Coordinate dimensions. Use XYZ for three- dimensional geometries. --not-null Add a NOT NULL constraint. - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional + :entrypoint -h, --help Show this message and exit. @@ -1394,7 +1398,7 @@ See :ref:`cli_spatialite_indexes`. paths. To load it from a specific path, use --load-extension. Options: - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 16a99cc..c09273d 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -94,7 +94,7 @@ def load_extension_option(fn): return click.option( "--load-extension", multiple=True, - help="SQLite extensions to load", + help="Path to SQLite extension, with optional :entrypoint", )(fn) @@ -2997,7 +2997,11 @@ def _load_extensions(db, load_extension): for ext in load_extension: if ext == "spatialite" and not os.path.exists(ext): ext = find_spatialite() - db.conn.load_extension(ext) + if ":" in ext: + path, _, entrypoint = ext.partition(":") + db.conn.execute("SELECT load_extension(?, ?)", [path, entrypoint]) + else: + db.conn.load_extension(ext) def _register_functions(db, functions): diff --git a/tests/ext.c b/tests/ext.c new file mode 100644 index 0000000..f5b3276 --- /dev/null +++ b/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); +} diff --git a/tests/test_cli.py b/tests/test_cli.py index 9b43cd5..24f36ed 100644 --- a/tests/test_cli.py +++ b/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,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