diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index 9e14c1ab..adbb1d1e 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -131,7 +131,7 @@ See :ref:`cli_query`. -r, --raw Raw output, first column of first row --raw-lines Raw output, first column of each row -p, --param ... Named :parameters for SQL query - --functions TEXT Python code defining one or more custom SQL + --functions TEXT Python code or file path defining custom SQL functions --load-extension TEXT Path to SQLite extension, with optional :entrypoint @@ -174,7 +174,7 @@ See :ref:`cli_memory`. sqlite-utils memory animals.csv --schema Options: - --functions TEXT Python code defining one or more custom SQL + --functions TEXT Python code or file path defining custom SQL functions --attach ... Additional databases to attach - specify alias and filepath @@ -374,7 +374,7 @@ See :ref:`cli_bulk`. Options: --batch-size INTEGER Commit every X records - --functions TEXT Python code defining one or more custom SQL functions + --functions TEXT Python code or file path defining custom SQL functions --flatten Flatten nested JSON objects, so {"a": {"b": 1}} becomes {"a_b": 1} --nl Expect newline-delimited JSON diff --git a/docs/cli.rst b/docs/cli.rst index 88756ba9..501cba6f 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -368,6 +368,22 @@ This example defines a function which extracts the domain from a URL: 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. +You can also pass the path to a Python file containing function definitions: + +.. code-block:: bash + + sqlite-utils query sites.db "select url, domain(url) from urls" --functions functions.py + +The ``--functions`` option can be used multiple times to load functions from multiple sources: + +.. code-block:: bash + + sqlite-utils query sites.db "select url, domain(url), extract_path(url) from urls" \ + --functions domain_funcs.py \ + --functions 'def extract_path(url): + from urllib.parse import urlparse + return urlparse(url).path' + .. _cli_query_extensions: SQLite extensions diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 81548bcd..66e5e27a 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1375,7 +1375,9 @@ def upsert( @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" + "--functions", + help="Python code or file path defining custom SQL functions", + multiple=True, ) @import_options @load_extension_option @@ -1764,7 +1766,9 @@ def drop_view(path, view, ignore, load_extension): help="Named :parameters for SQL query", ) @click.option( - "--functions", help="Python code defining one or more custom SQL functions" + "--functions", + help="Python code or file path defining custom SQL functions", + multiple=True, ) @load_extension_option def query( @@ -1828,7 +1832,9 @@ def query( ) @click.argument("sql") @click.option( - "--functions", help="Python code defining one or more custom SQL functions" + "--functions", + help="Python code or file path defining custom SQL functions", + multiple=True, ) @click.option( "--attach", @@ -3290,6 +3296,13 @@ def _load_extensions(db, load_extension): def _register_functions(db, functions): # Register any Python functions as SQL functions: + # Check if this is a file path + if "\n" not in functions and functions.endswith(".py"): + try: + functions = pathlib.Path(functions).read_text() + except FileNotFoundError: + raise click.ClickException("File not found: {}".format(functions)) + sqlite3.enable_callback_tracebacks(True) globals = {} try: @@ -3308,9 +3321,10 @@ def _value_or_none(value): return value -def _maybe_register_functions(db, functions): - functions = _value_or_none(functions) - if isinstance(functions, (bytes, bytearray)): - functions = functions.decode("utf-8") - if isinstance(functions, str) and functions.strip(): - _register_functions(db, functions) +def _maybe_register_functions(db, functions_list): + if not functions_list: + return + for functions in functions_list: + functions = _value_or_none(functions) + if isinstance(functions, str) and functions.strip(): + _register_functions(db, functions) diff --git a/tests/test_cli.py b/tests/test_cli.py index bcf3e983..092a6fc5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -806,6 +806,77 @@ def test_hidden_functions_are_hidden(db_path): assert "_two" not in functions +def test_query_functions_from_file(db_path, tmp_path): + # Create a temporary file with function definitions + functions_file = tmp_path / "my_functions.py" + functions_file.write_text(TEST_FUNCTIONS) + + result = CliRunner().invoke( + cli.cli, + [ + db_path, + "select zero(), one(1), two(1, 2)", + "--functions", + str(functions_file), + ], + ) + assert result.exit_code == 0 + assert json.loads(result.output.strip()) == [ + {"zero()": 0, "one(1)": 1, "two(1, 2)": 3} + ] + + +def test_query_functions_file_not_found(db_path): + result = CliRunner().invoke( + cli.cli, + [ + db_path, + "select zero()", + "--functions", + "nonexistent.py", + ], + ) + assert result.exit_code == 1 + assert "File not found: nonexistent.py" in result.output + + +def test_query_functions_multiple_invocations(db_path): + # Test using --functions multiple times + result = CliRunner().invoke( + cli.cli, + [ + db_path, + "select triple(2), quadruple(2)", + "--functions", + "def triple(x):\n return x * 3", + "--functions", + "def quadruple(x):\n return x * 4", + ], + ) + assert result.exit_code == 0 + assert json.loads(result.output.strip()) == [{"triple(2)": 6, "quadruple(2)": 8}] + + +def test_query_functions_file_and_inline(db_path, tmp_path): + # Test combining file and inline code + functions_file = tmp_path / "file_funcs.py" + functions_file.write_text("def triple(x):\n return x * 3") + + result = CliRunner().invoke( + cli.cli, + [ + db_path, + "select triple(2), quadruple(2)", + "--functions", + str(functions_file), + "--functions", + "def quadruple(x):\n return x * 4", + ], + ) + assert result.exit_code == 0 + assert json.loads(result.output.strip()) == [{"triple(2)": 6, "quadruple(2)": 8}] + + 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" diff --git a/tests/test_cli_bulk.py b/tests/test_cli_bulk.py index 909ed096..514f4acc 100644 --- a/tests/test_cli_bulk.py +++ b/tests/test_cli_bulk.py @@ -45,6 +45,32 @@ def test_cli_bulk(test_db_and_path): ] == list(db["example"].rows) +def test_cli_bulk_multiple_functions(test_db_and_path): + db, db_path = test_db_and_path + result = CliRunner().invoke( + cli.cli, + [ + "bulk", + db_path, + "insert into example (id, name) values (:id, myupper(mylower(:name)))", + "-", + "--nl", + "--functions", + "myupper = lambda s: s.upper()", + "--functions", + "mylower = lambda s: s.lower()", + ], + input='{"id": 3, "name": "ThReE"}\n{"id": 4, "name": "FoUr"}\n', + ) + assert result.exit_code == 0, result.output + assert [ + {"id": 1, "name": "One"}, + {"id": 2, "name": "Two"}, + {"id": 3, "name": "THREE"}, + {"id": 4, "name": "FOUR"}, + ] == list(db["example"].rows) + + def test_cli_bulk_batch_size(test_db_and_path): db, db_path = test_db_and_path proc = subprocess.Popen( diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py index 572a13ec..8483963b 100644 --- a/tests/test_cli_memory.py +++ b/tests/test_cli_memory.py @@ -307,6 +307,22 @@ def test_memory_functions(): assert result.output.strip() == '[{"hello()": "Hello"}]' +def test_memory_functions_multiple(): + result = CliRunner().invoke( + cli.cli, + [ + "memory", + "select triple(2), quadruple(2)", + "--functions", + "def triple(x):\n return x * 3", + "--functions", + "def quadruple(x):\n return x * 4", + ], + ) + assert result.exit_code == 0 + assert result.output.strip() == '[{"triple(2)": 6, "quadruple(2)": 8}]' + + def test_memory_return_db(tmpdir): # https://github.com/simonw/sqlite-utils/issues/643 from sqlite_utils.cli import cli