From 6964d67ce160e42202583d47177761e1a7093482 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sat, 31 Jul 2021 21:33:00 -0700 Subject: [PATCH] Began implementation of sqlite-utils convert, refs #251 --- sqlite_utils/cli.py | 206 ++++++++++++++++++++++++++ sqlite_utils/utils.py | 14 ++ tests/test_convert.py | 331 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 551 insertions(+) create mode 100644 tests/test_convert.py diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index b168ec0..56af19e 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -17,6 +17,7 @@ import tabulate from .utils import ( file_progress, find_spatialite, + progressbar, sqlite3, decode_base64_values, rows_from_file, @@ -1903,6 +1904,211 @@ def analyze_tables( click.echo(details) +@cli.command(name="convert") +@click.argument( + "db_path", + type=click.Path(file_okay=True, dir_okay=False, allow_dash=False), + required=True, +) +@click.argument("table", type=str) +@click.argument("columns", type=str, nargs=-1, required=True) +@click.argument("code", type=str) +@click.option( + "--import", "imports", type=str, multiple=True, help="Python modules to import" +) +@click.option( + "--dry-run", is_flag=True, help="Show results of running this against first 10 rows" +) +@click.option( + "--multi", is_flag=True, help="Populate columns for keys in returned dictionary" +) +@click.option("--output", help="Optional separate column to populate with the output") +@click.option( + "--output-type", + help="Column type to use for the output column", + default="text", + type=click.Choice(["integer", "float", "blob", "text"]), +) +@click.option("--drop", is_flag=True, help="Drop original column afterwards") +@click.option("-s", "--silent", is_flag=True, help="Don't show a progress bar") +def lambda_( + db_path, + table, + columns, + code, + imports, + dry_run, + multi, + output, + output_type, + drop, + silent, +): + """ + Convert columns using Python code you supply. For example: + + \b + $ sqlite-utils convert my.db mytable mycolumn + --code='"\\n".join(textwrap.wrap(value, 10))' + --import=textwrap + + "value" is a variable with the column value to be converted. + """ + if output is not None and len(columns) > 1: + raise click.ClickException("Cannot use --output with more than one column") + if multi and len(columns) > 1: + raise click.ClickException("Cannot use --multi with more than one column") + # If single line and no 'return', add the return + if "\n" not in code and not code.strip().startswith("return "): + code = "return {}".format(code) + # Compile the code into a function body called fn(value) + new_code = ["def fn(value):"] + for line in code.split("\n"): + new_code.append(" {}".format(line)) + code_o = compile("\n".join(new_code), "", "exec") + locals = {} + globals = {} + for import_ in imports: + globals[import_] = __import__(import_) + exec(code_o, globals, locals) + fn = locals["fn"] + if dry_run: + # Pull first 20 values for first column and preview them + db = sqlite3.connect(db_path) + db.create_function("preview_transform", 1, lambda v: fn(v) if v else v) + sql = """ + select + [{column}] as value, + preview_transform([{column}]) as preview + from [{table}] limit 10 + """.format( + column=columns[0], table=table + ) + for row in db.execute(sql).fetchall(): + print(row[0]) + print(" --- becomes:") + print(row[1]) + print() + elif multi: + _transform_multi(db_path, table, columns[0], fn, drop, silent) + else: + _transform(db_path, table, columns, fn, output, output_type, drop, silent) + + +def _transform(db_path, table, columns, fn, output, output_type, drop, silent): + db = sqlite_utils.Database(db_path) + count_sql = "select count(*) from [{}]".format(table) + todo_count = list(db.execute(count_sql).fetchall())[0][0] * len(columns) + + if drop and not output: + raise click.ClickException("--drop can only be used with --output or --multi") + + if output is not None: + if output not in db[table].columns_dict: + db[table].add_column(output, output_type or "text") + + with progressbar(length=todo_count, silent=silent) as bar: + + def transform_value(v): + bar.update(1) + if not v: + return v + return fn(v) + + db.register_function(transform_value) + sql = "update [{table}] set {sets};".format( + table=table, + sets=", ".join( + [ + "[{output_column}] = transform_value([{column}])".format( + output_column=output or column, column=column + ) + for column in columns + ] + ), + ) + with db.conn: + db.execute(sql) + if drop: + db[table].transform(drop=columns) + + +def _transform_multi(db_path, table, column, fn, drop, silent): + db = sqlite_utils.Database(db_path) + # First we execute the function + pk_to_values = {} + new_column_types = {} + pks = [column.name for column in db[table].columns if column.is_pk] + if not pks: + pks = ["rowid"] + + with progressbar( + length=db[table].count, silent=silent, label="1: Evaluating" + ) as bar: + for row in db[table].rows_where( + select=", ".join( + "[{}]".format(column_name) for column_name in (pks + [column]) + ) + ): + row_pk = tuple(row[pk] for pk in pks) + if len(row_pk) == 1: + row_pk = row_pk[0] + values = fn(row[column]) + if values is not None and not isinstance(values, dict): + raise click.ClickException( + "With --multi code must return a Python dictionary - returned {}".format( + repr(values) + ) + ) + if values: + for key, value in values.items(): + new_column_types.setdefault(key, set()).add(type(value)) + pk_to_values[row_pk] = values + bar.update(1) + + # Add any new columns + columns_to_create = _suggest_column_types(new_column_types) + for column_name, column_type in columns_to_create.items(): + if column_name not in db[table].columns_dict: + db[table].add_column(column_name, column_type) + + # Run the updates + with progressbar(length=db[table].count, silent=silent, label="2: Updating") as bar: + with db.conn: + for pk, updates in pk_to_values.items(): + db[table].update(pk, updates) + bar.update(1) + if drop: + db[table].transform(drop=(column,)) + + +def _suggest_column_types(all_column_types): + column_types = {} + for key, types in all_column_types.items(): + # Ignore null values if at least one other type present: + if len(types) > 1: + types.discard(None.__class__) + if {None.__class__} == types: + t = str + elif len(types) == 1: + t = list(types)[0] + # But if it's a subclass of list / tuple / dict, use str + # instead as we will be storing it as JSON in the table + for superclass in (list, tuple, dict): + if issubclass(t, superclass): + t = str + elif {int, bool}.issuperset(types): + t = int + elif {int, float, bool}.issuperset(types): + t = float + elif {bytes, str}.issuperset(types): + t = bytes + else: + t = str + column_types[key] = t + return column_types + + def _render_common(title, values): if values is None: return "" diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 4f3c819..977dd79 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -254,3 +254,17 @@ class ValueTracker: not_these.append(name) for key in not_these: del self.couldbe[key] + + +class NullProgressBar: + def update(self, value): + pass + + +@contextlib.contextmanager +def progressbar(silent=False, **kwargs): + if silent: + yield NullProgressBar() + else: + with click.progressbar(**kwargs) as bar: + yield bar diff --git a/tests/test_convert.py b/tests/test_convert.py new file mode 100644 index 0000000..acaf555 --- /dev/null +++ b/tests/test_convert.py @@ -0,0 +1,331 @@ +from click.testing import CliRunner +from sqlite_utils import cli +import sqlite_utils +import textwrap +import pathlib +import pytest + + +@pytest.fixture +def test_db_and_path(fresh_db_and_path): + db, db_path = fresh_db_and_path + db["example"].insert_all( + [ + {"id": 1, "dt": "5th October 2019 12:04"}, + {"id": 2, "dt": "6th October 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ], + pk="id", + ) + return db, db_path + + +@pytest.fixture +def fresh_db_and_path(tmpdir): + db_path = str(pathlib.Path(tmpdir) / "data.db") + db = sqlite_utils.Database(db_path) + return db, db_path + + +@pytest.mark.parametrize( + "code", + [ + "return value.replace('October', 'Spooktober')", + # Return is optional: + "value.replace('October', 'Spooktober')", + ], +) +def test_convert_single_line(test_db_and_path, code): + db, db_path = test_db_and_path + result = CliRunner().invoke(cli.cli, ["convert", db_path, "example", "dt", code]) + assert 0 == result.exit_code, result.output + assert [ + {"id": 1, "dt": "5th Spooktober 2019 12:04"}, + {"id": 2, "dt": "6th Spooktober 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] == list(db["example"].rows) + + +def test_convert_multiple_lines(test_db_and_path): + db, db_path = test_db_and_path + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "v = value.replace('October', 'Spooktober')\nreturn v.upper()", + ], + ) + assert 0 == result.exit_code, result.output + assert [ + {"id": 1, "dt": "5TH SPOOKTOBER 2019 12:04"}, + {"id": 2, "dt": "6TH SPOOKTOBER 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] == list(db["example"].rows) + + +def test_convert_import(test_db_and_path): + db, db_path = test_db_and_path + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "return re.sub('O..', 'OXX', value)", + "--import", + "re", + ], + ) + assert 0 == result.exit_code, result.output + assert [ + {"id": 1, "dt": "5th OXXober 2019 12:04"}, + {"id": 2, "dt": "6th OXXober 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] == list(db["example"].rows) + + +def test_convert_dryrun(test_db_and_path): + db, db_path = test_db_and_path + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "return re.sub('O..', 'OXX', value)", + "--import", + "re", + "--dry-run", + ], + ) + assert result.exit_code == 0 + assert result.output.strip() == ( + "5th October 2019 12:04\n" + " --- becomes:\n" + "5th OXXober 2019 12:04\n" + "\n" + "6th October 2019 00:05:06\n" + " --- becomes:\n" + "6th OXXober 2019 00:05:06\n" + "\n" + "\n" + " --- becomes:\n" + "\n" + "\n" + "None\n" + " --- becomes:\n" + "None" + ) + # But it should not have actually modified the table data + 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}, + ] + + +@pytest.mark.parametrize("drop", (True, False)) +def test_convert_output_column(test_db_and_path, drop): + db, db_path = test_db_and_path + args = [ + "convert", + db_path, + "example", + "dt", + "value.replace('October', 'Spooktober')", + "--output", + "newcol", + ] + if drop: + args += ["--drop"] + result = CliRunner().invoke(cli.cli, args) + assert 0 == result.exit_code, result.output + expected = [ + { + "id": 1, + "dt": "5th October 2019 12:04", + "newcol": "5th Spooktober 2019 12:04", + }, + { + "id": 2, + "dt": "6th October 2019 00:05:06", + "newcol": "6th Spooktober 2019 00:05:06", + }, + {"id": 3, "dt": "", "newcol": ""}, + {"id": 4, "dt": None, "newcol": None}, + ] + if drop: + for row in expected: + del row["dt"] + assert list(db["example"].rows) == expected + + +@pytest.mark.parametrize( + "output_type,expected", + ( + ("text", [(1, "1"), (2, "2"), (3, "3"), (4, "4")]), + ("float", [(1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0)]), + ("integer", [(1, 1), (2, 2), (3, 3), (4, 4)]), + (None, [(1, "1"), (2, "2"), (3, "3"), (4, "4")]), + ), +) +def test_convert_output_column_output_type(test_db_and_path, output_type, expected): + db, db_path = test_db_and_path + args = [ + "convert", + db_path, + "example", + "id", + "value", + "--output", + "new_id", + ] + if output_type: + args += ["--output-type", output_type] + result = CliRunner().invoke( + cli.cli, + args, + ) + assert 0 == result.exit_code, result.output + assert expected == list(db.execute("select id, new_id from example")) + + +@pytest.mark.parametrize( + "options,expected_error", + [ + ( + [ + "dt", + "id", + "value.replace('October', 'Spooktober')", + "--output", + "newcol", + ], + "Cannot use --output with more than one column", + ), + ( + [ + "dt", + "value.replace('October', 'Spooktober')", + "--output", + "newcol", + "--output-type", + "invalid", + ], + "Error: Invalid value for '--output-type'", + ), + ( + [ + "value.replace('October', 'Spooktober')", + ], + "Missing argument 'COLUMNS...'", + ), + ], +) +def test_convert_output_error(test_db_and_path, options, expected_error): + db_path = test_db_and_path[1] + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + ] + + options, + ) + assert result.exit_code != 0 + assert expected_error in result.output + + +@pytest.mark.parametrize("drop", (True, False)) +def test_convert_multi(fresh_db_and_path, drop): + db, db_path = fresh_db_and_path + db["creatures"].insert_all( + [ + {"id": 1, "name": "Simon"}, + {"id": 2, "name": "Cleo"}, + ], + pk="id", + ) + args = [ + "convert", + db_path, + "creatures", + "name", + "--multi", + '{"upper": value.upper(), "lower": value.lower()}', + ] + if drop: + args += ["--drop"] + result = CliRunner().invoke(cli.cli, args) + assert result.exit_code == 0, result.output + expected = [ + {"id": 1, "name": "Simon", "upper": "SIMON", "lower": "simon"}, + {"id": 2, "name": "Cleo", "upper": "CLEO", "lower": "cleo"}, + ] + if drop: + for row in expected: + del row["name"] + assert list(db["creatures"].rows) == expected + + +def test_convert_multi_complex_column_types(fresh_db_and_path): + db, db_path = fresh_db_and_path + db["rows"].insert_all( + [ + {"id": 1}, + {"id": 2}, + {"id": 3}, + {"id": 4}, + ], + pk="id", + ) + code = textwrap.dedent( + """ + if value == 1: + return {"is_str": "", "is_float": 1.2, "is_int": None} + elif value == 2: + return {"is_float": 1, "is_int": 12} + elif value == 3: + return {"is_bytes": b"blah"} + """ + ) + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "rows", + "id", + "--multi", + code, + ], + ) + assert result.exit_code == 0, result.output + assert list(db["rows"].rows) == [ + {"id": 1, "is_str": "", "is_float": 1.2, "is_int": None, "is_bytes": None}, + {"id": 2, "is_str": None, "is_float": 1.0, "is_int": 12, "is_bytes": None}, + { + "id": 3, + "is_str": None, + "is_float": None, + "is_int": None, + "is_bytes": b"blah", + }, + {"id": 4, "is_str": None, "is_float": None, "is_int": None, "is_bytes": None}, + ] + assert db["rows"].schema == ( + "CREATE TABLE [rows] (\n" + " [id] INTEGER PRIMARY KEY\n" + ", [is_str] TEXT, [is_float] FLOAT, [is_int] INTEGER, [is_bytes] BLOB)" + )