diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index 9a025ba..c194967 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -349,6 +349,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 diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 24c8c03..16a99cc 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -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)) 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)