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 2f4f5a3..d74855e 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -119,7 +119,10 @@ See :ref:`cli_query`. escaped strings -r, --raw Raw output, first column of first row -p, --param ... Named :parameters for SQL query - --load-extension TEXT SQLite extensions to load + --functions TEXT Python code defining one or more custom SQL + functions + --load-extension TEXT Path to SQLite extension, with optional + :entrypoint -h, --help Show this message and exit. @@ -159,6 +162,8 @@ See :ref:`cli_memory`. sqlite-utils memory animals.csv --schema Options: + --functions TEXT Python code defining one or more custom SQL + functions --attach ... Additional databases to attach - specify alias and filepath --flatten Flatten nested JSON objects, so {"foo": {"bar": @@ -185,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. @@ -262,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 @@ -316,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. @@ -345,6 +351,7 @@ See :ref:`cli_bulk`. Options: --batch-size INTEGER Commit every X records + --functions TEXT Python code defining one or more custom SQL functions --flatten Flatten nested JSON objects, so {"a": {"b": 1}} becomes {"a_b": 1} --nl Expect newline-delimited JSON @@ -359,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. @@ -399,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. @@ -435,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. @@ -460,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. @@ -482,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. @@ -520,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. @@ -544,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. @@ -651,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. @@ -689,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. @@ -713,6 +720,7 @@ See :ref:`cli_rows`. Options: -c, --column TEXT Columns to return --where TEXT Optional where clause + -o, --order TEXT Order by ('column' or 'column desc') -p, --param ... Named :parameters for where clause --limit INTEGER Number of rows to return - defaults to everything --offset INTEGER SQL offset to use @@ -729,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. @@ -764,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. @@ -800,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. @@ -824,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. @@ -857,7 +866,7 @@ See :ref:`cli_create_table`. --ignore If table already exists, do nothing --replace If table already exists, replace it --transform If table already exists, try to transform the schema - --load-extension TEXT SQLite extensions to load + --load-extension TEXT Path to SQLite extension, with optional :entrypoint -h, --help Show this message and exit. @@ -887,7 +896,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. @@ -915,7 +924,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. @@ -935,7 +944,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. @@ -955,7 +964,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. @@ -975,7 +984,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. @@ -999,7 +1008,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. @@ -1064,7 +1073,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. @@ -1092,7 +1101,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. @@ -1118,7 +1127,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. @@ -1142,7 +1151,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. @@ -1164,7 +1173,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. @@ -1186,7 +1195,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. @@ -1206,7 +1215,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. @@ -1228,7 +1237,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. @@ -1248,7 +1257,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. @@ -1265,7 +1274,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. @@ -1288,7 +1297,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. @@ -1313,7 +1322,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. @@ -1336,7 +1345,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. @@ -1367,7 +1376,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. @@ -1389,7 +1399,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/docs/cli.rst b/docs/cli.rst index d2d1812..3c5c595 100644 --- a/docs/cli.rst +++ b/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 ----------------- @@ -467,6 +487,8 @@ Or pass named parameters using ``--where`` in combination with ``-p``:: $ sqlite-utils rows dogs.db dogs -c name --where 'name = :name' -p name Cleo [{"name": "Cleo"}] +You can define a sort order using ``--order column`` or ``--order 'column desc'``. + Use ``--limit N`` to only return the first ``N`` rows. Use ``--offset N`` to return rows starting from the specified offset. .. note:: @@ -1337,12 +1359,9 @@ The following example adds a new ``score`` column, then updates it to list a ran random.seed(10) def convert(value): - global random return random.random() ' -Note the ``global random`` line here. Due to the way the tool compiles Python code, this is necessary to ensure the ``random`` module is available within the ``convert()`` function. If you were to omit this you would see a ``NameError: name 'random' is not defined`` error. - .. _cli_convert_recipes: sqlite-utils convert recipes diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index fe5a402..c51b101 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) @@ -926,9 +926,12 @@ def insert_upsert_implementation( load_extension=None, silent=False, bulk_sql=None, + functions=None, ): db = sqlite_utils.Database(path) _load_extensions(db, load_extension) + if functions: + _register_functions(db, functions) if (delimiter or quotechar or sniff or no_headers) and not tsv: csv = True if (nl + csv + tsv) >= 2: @@ -1305,6 +1308,9 @@ def upsert( @click.argument("sql") @click.argument("file", type=click.File("rb"), required=True) @click.option("--batch-size", type=int, default=100, help="Commit every X records") +@click.option( + "--functions", help="Python code defining one or more custom SQL functions" +) @import_options @load_extension_option def bulk( @@ -1312,6 +1318,7 @@ def bulk( sql, file, batch_size, + functions, flatten, nl, csv, @@ -1368,6 +1375,7 @@ def bulk( load_extension=load_extension, silent=False, bulk_sql=sql, + functions=functions, ) except (OperationalError, sqlite3.IntegrityError) as e: raise click.ClickException(str(e)) @@ -1655,6 +1663,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, @@ -1671,6 +1682,7 @@ def query( raw, param, load_extension, + functions, ): """Execute SQL query and return the results as JSON @@ -1687,6 +1699,9 @@ def query( _load_extensions(db, load_extension) db.register_fts4_bm25() + if functions: + _register_functions(db, functions) + _execute_query( db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols ) @@ -1700,6 +1715,9 @@ def query( nargs=-1, ) @click.argument("sql") +@click.option( + "--functions", help="Python code defining one or more custom SQL functions" +) @click.option( "--attach", type=(str, click.Path(file_okay=True, dir_okay=False, allow_dash=False)), @@ -1746,6 +1764,7 @@ def query( def memory( paths, sql, + functions, attach, flatten, nl, @@ -1859,6 +1878,9 @@ def memory( _load_extensions(db, load_extension) db.register_fts4_bm25() + if functions: + _register_functions(db, functions) + _execute_query( db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols ) @@ -2007,6 +2029,7 @@ def search( @click.argument("dbtable") @click.option("-c", "--column", type=str, multiple=True, help="Columns to return") @click.option("--where", help="Optional where clause") +@click.option("-o", "--order", type=str, help="Order by ('column' or 'column desc')") @click.option( "-p", "--param", @@ -2033,6 +2056,7 @@ def rows( dbtable, column, where, + order, param, limit, offset, @@ -2059,6 +2083,8 @@ def rows( sql = "select {} from [{}]".format(columns, dbtable) if where: sql += " where " + where + if order: + sql += " order by " + order if limit: sql += " limit {}".format(limit) if offset: @@ -2993,4 +3019,22 @@ 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): + # Register any Python functions as SQL 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) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 49b1f68..179cea9 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1842,7 +1842,7 @@ class Table(Queryable): magic_lookup_column=magic_lookup_column, lookup_table=table, where=" AND ".join( - "[{table}].[{column}] = [{lookup_table}].[{lookup_column}]".format( + "[{table}].[{column}] IS [{lookup_table}].[{lookup_column}]".format( table=self.name, lookup_table=table, column=column, diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 64bba50..c0b7bf1 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -432,12 +432,11 @@ def progressbar(*args, **kwargs): def _compile_code(code, imports, variable="value"): - locals = {} globals = {"r": recipes, "recipes": recipes} # If user defined a convert() function, return that try: - exec(code, globals, locals) - return locals["convert"] + exec(code, globals) + return globals["convert"] except (AttributeError, SyntaxError, NameError, KeyError, TypeError): pass @@ -464,8 +463,8 @@ def _compile_code(code, imports, variable="value"): for import_ in imports: globals[import_.split(".")[0]] = __import__(import_) - exec(code_o, globals, locals) - return locals["fn"] + exec(code_o, globals) + return globals["fn"] def chunks(sequence: Iterable, size: int) -> Iterable[Iterable]: 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 061a578..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 @@ -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 (, 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 diff --git a/tests/test_cli_bulk.py b/tests/test_cli_bulk.py index 6d01952..909ed09 100644 --- a/tests/test_cli_bulk.py +++ b/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) diff --git a/tests/test_cli_convert.py b/tests/test_cli_convert.py index 10b4563..d26ab8c 100644 --- a/tests/test_cli_convert.py +++ b/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"}, + ] diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py index aee74e7..bb50d23 100644 --- a/tests/test_cli_memory.py +++ b/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"}]'