mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-26 02:44:33 +02:00
Began implementation of sqlite-utils convert, refs #251
This commit is contained in:
Simon Willison
parent
c7e8d72be9
commit
6964d67ce1
3 changed files with 551 additions and 0 deletions
|
|
@ -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), "<string>", "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 ""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
331
tests/test_convert.py
Normal file
331
tests/test_convert.py
Normal file
|
|
@ -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)"
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue