From d83b2568131f2b1cc01228419bb08c96d843d65d Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Mon, 2 Aug 2021 11:58:05 -0700 Subject: [PATCH] --where and -p options for sqlite-utils convert, closes #304 --- docs/cli.rst | 11 ++++++ sqlite_utils/cli.py | 26 +++++++++++-- tests/test_cli_convert.py | 78 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 111 insertions(+), 4 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index 7bc7a3b..72c8bbf 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -945,6 +945,17 @@ You can specify Python modules that should be imported and made available to you '"\n".join(textwrap.wrap(value, 10))' \ --import=textwrap +The transformation will be applied to every row in the specified table. You can limit that to just rows that match a ``WHERE`` clause using ``--where``:: + + $ sqlite-utils convert content.db articles headline 'value.upper()' \ + --where "headline like '%cat%'" + +You can include named parameters in your where clause and populate them using one or more ``--param`` options:: + + $ sqlite-utils convert content.db articles headline 'value.upper()' \ + --where "headline like :like" \ + --param like '%cat%' + The ``--dry-run`` option will output a preview of the conversion against the first ten rows, without modifying the database. .. _cli_convert_recipes: diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index e1b770f..3f32413 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1959,6 +1959,14 @@ def _generate_convert_help(): @click.option( "--multi", is_flag=True, help="Populate columns for keys in returned dictionary" ) +@click.option("--where", help="Optional where clause") +@click.option( + "-p", + "--param", + multiple=True, + type=(str, str), + help="Named :parameters for where clause", +) @click.option("--output", help="Optional separate column to populate with the output") @click.option( "--output-type", @@ -1976,6 +1984,8 @@ def convert( imports, dry_run, multi, + where, + param, output, output_type, drop, @@ -1992,6 +2002,7 @@ def convert( # If single line and no 'return', add the return if "\n" not in code and not code.strip().startswith("return "): code = "return {}".format(code) + where_args = dict(param) if param else [] # Compile the code into a function body called fn(value) new_code = ["def fn(value):"] for line in code.split("\n"): @@ -2010,20 +2021,29 @@ def convert( select [{column}] as value, preview_transform([{column}]) as preview - from [{table}] limit 10 + from [{table}]{where} limit 10 """.format( - column=columns[0], table=table + column=columns[0], + table=table, + where=" where {}".format(where) if where is not None else "", ) - for row in db.conn.execute(sql).fetchall(): + for row in db.conn.execute(sql, where_args).fetchall(): click.echo(str(row[0])) click.echo(" --- becomes:") click.echo(str(row[1])) click.echo() + count = db[table].count_where( + where=where, + where_args=where_args, + ) + click.echo("Would affect {} row{}".format(count, "" if count == 1 else "s")) else: try: db[table].convert( columns, fn, + where=where, + where_args=where_args, output=output, output_type=output_type, drop=drop, diff --git a/tests/test_cli_convert.py b/tests/test_cli_convert.py index 51634cd..8ce5203 100644 --- a/tests/test_cli_convert.py +++ b/tests/test_cli_convert.py @@ -124,7 +124,8 @@ def test_convert_dryrun(test_db_and_path): "\n" "None\n" " --- becomes:\n" - "None" + "None\n\n" + "Would affect 4 rows" ) # But it should not have actually modified the table data assert list(db["example"].rows) == [ @@ -133,6 +134,27 @@ def test_convert_dryrun(test_db_and_path): {"id": 3, "dt": ""}, {"id": 4, "dt": None}, ] + # Test with a where clause too + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "return re.sub('O..', 'OXX', value)", + "--import", + "re", + "--dry-run", + "--where", + "id = :id", + "-p", + "id", + "4", + ], + ) + assert result.exit_code == 0 + assert result.output.strip().split("\n")[-1] == "Would affect 1 row" @pytest.mark.parametrize("drop", (True, False)) @@ -439,3 +461,57 @@ def test_multi_with_bad_function(test_db_and_path): result = CliRunner().invoke(cli.cli, args) assert result.exit_code == 1, result.output assert "When using --multi code must return a Python dictionary" in result.output + + +def test_convert_where(test_db_and_path): + db, db_path = test_db_and_path + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "str(value).upper()", + "--where", + "id = :id", + "-p", + "id", + 2, + ], + ) + assert result.exit_code == 0, result.output + assert list(db["example"].rows) == [ + {"id": 1, "dt": "5th October 2019 12:04"}, + {"id": 2, "dt": "6TH OCTOBER 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] + + +def test_convert_where_multi(fresh_db_and_path): + db, db_path = fresh_db_and_path + db["names"].insert_all( + [{"id": 1, "name": "Cleo"}, {"id": 2, "name": "Bants"}], pk="id" + ) + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "names", + "name", + '{"upper": value.upper()}', + "--where", + "id = :id", + "-p", + "id", + 2, + "--multi", + ], + ) + assert 0 == result.exit_code, result.output + assert list(db["names"].rows) == [ + {"id": 1, "name": "Cleo", "upper": None}, + {"id": 2, "name": "Bants", "upper": "BANTS"}, + ]