Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions docs/cli-reference.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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 <TEXT TEXT>... 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
Expand Down Expand Up @@ -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 <TEXT FILE>... Additional databases to attach - specify alias and
filepath
Expand Down Expand Up @@ -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
Expand Down
16 changes: 16 additions & 0 deletions docs/cli.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
32 changes: 23 additions & 9 deletions sqlite_utils/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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:
Expand All @@ -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)
71 changes: 71 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
26 changes: 26 additions & 0 deletions tests/test_cli_bulk.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
16 changes: 16 additions & 0 deletions tests/test_cli_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading