Implemented sqlite-utils transform command, closes #164

This commit is contained in:
Simon Willison 2020-09-22 00:46:32 -07:00
commit 752d261229
5 changed files with 309 additions and 39 deletions

View file

@ -887,6 +887,106 @@ def rows(ctx, path, dbtable, nl, arrays, csv, no_headers, table, fmt, json_cols)
)
@cli.command()
@click.argument(
"path",
type=click.Path(file_okay=True, dir_okay=False, allow_dash=False),
required=True,
)
@click.argument("table")
@click.option(
"--type",
type=(str, str),
multiple=True,
help="Change column type to X",
)
@click.option("--drop", type=str, multiple=True, help="Drop this column")
@click.option(
"--rename", type=(str, str), multiple=True, help="Rename this column to X"
)
@click.option("--not-null", type=str, multiple=True, help="Set this column to NOT NULL")
@click.option(
"--not-null-false", type=str, multiple=True, help="Remove NOT NULL from this column"
)
@click.option("--pk", type=str, multiple=True, help="Make this column the primary key")
@click.option(
"--pk-none", is_flag=True, help="Remove primary key (convert to rowid table)"
)
@click.option(
"--default",
type=(str, str),
multiple=True,
help="Set default value for this column",
)
@click.option(
"--default-none", type=str, multiple=True, help="Remove default from this column"
)
@click.option(
"--drop-foreign-key",
type=(str, str, str),
multiple=True,
help="Drop this foreign key constraint",
)
@click.option("--sql", is_flag=True, help="Output SQL without executing it")
def transform(
path,
table,
type,
drop,
rename,
not_null,
not_null_false,
pk,
pk_none,
default,
default_none,
drop_foreign_key,
sql,
):
db = sqlite_utils.Database(path)
types = {}
kwargs = {}
for column, ctype in type:
if ctype.upper() not in VALID_COLUMN_TYPES:
raise click.ClickException(
"column types must be one of {}".format(VALID_COLUMN_TYPES)
)
types[column] = ctype.upper()
not_null_dict = {}
for column in not_null:
not_null_dict[column] = True
for column in not_null_false:
not_null_dict[column] = False
default_dict = {}
for column, value in default:
default_dict[column] = value
for column in default_none:
default_dict[column] = None
kwargs["types"] = types
kwargs["drop"] = set(drop)
kwargs["rename"] = dict(rename)
kwargs["not_null"] = not_null_dict
if pk:
if len(pk) == 1:
kwargs["pk"] = pk[0]
else:
kwargs["pk"] = pk
elif pk_none:
kwargs["pk"] = None
kwargs["defaults"] = default_dict
if drop_foreign_key:
kwargs["drop_foreign_keys"] = drop_foreign_key
if sql:
for line in db[table].transform_sql(**kwargs):
click.echo(line)
else:
db[table].transform(**kwargs)
@cli.command(name="insert-files")
@click.argument(
"path",

View file

@ -731,7 +731,7 @@ class Table(Queryable):
sqls = self.transform_sql(
types=types,
rename=rename,
drop=None,
drop=drop,
pk=pk,
not_null=not_null,
defaults=defaults,
@ -790,7 +790,7 @@ class Table(Queryable):
sqls = []
if should_flip_foreign_keys_pragma:
sqls.append("PRAGMA foreign_keys=OFF")
sqls.append("PRAGMA foreign_keys=OFF;")
if pk is DEFAULT:
pks_renamed = tuple(rename.get(p) or p for p in self.pks)
@ -800,7 +800,12 @@ class Table(Queryable):
pk = pks_renamed
# not_null may be a set or dict, need to convert to a set
create_table_not_null = {c.name for c in self.columns if c.notnull}
create_table_not_null = {
rename.get(c.name) or c.name
for c in self.columns
if c.notnull
if c.name not in drop
}
if isinstance(not_null, dict):
# Remove any columns with a value of False
for key, value in not_null.items():
@ -811,13 +816,16 @@ class Table(Queryable):
else:
create_table_not_null.add(key)
elif isinstance(not_null, set):
create_table_not_null.update(rename.get(k) or k for k in not_null)
create_table_not_null.update((rename.get(k) or k) for k in not_null)
elif not_null is None:
pass
else:
assert False, "not_null must be a dict or a set or None"
# defaults=
create_table_defaults = {
(rename.get(c.name) or c.name): c.default_value
for c in self.columns
if c.default_value is not None
if c.default_value is not None and c.name not in drop
}
if defaults is not None:
create_table_defaults.update(
@ -844,13 +852,14 @@ class Table(Queryable):
foreign_keys=create_table_foreign_keys,
).strip()
)
# Copy across data, respecting any renamed columns
new_cols = []
old_cols = []
for from_, to_ in copy_from_to.items():
old_cols.append(from_)
new_cols.append(to_)
copy_sql = "INSERT INTO [{new_table}] ({new_cols}) SELECT {old_cols} FROM [{old_table}]".format(
copy_sql = "INSERT INTO [{new_table}] ({new_cols})\n SELECT {old_cols} FROM [{old_table}];".format(
new_table=new_table_name,
old_table=self.name,
old_cols=", ".join("[{}]".format(col) for col in old_cols),
@ -858,13 +867,13 @@ class Table(Queryable):
)
sqls.append(copy_sql)
# Drop the old table
sqls.append("DROP TABLE [{}]".format(self.name))
sqls.append("DROP TABLE [{}];".format(self.name))
# Rename the new one
sqls.append("ALTER TABLE [{}] RENAME TO [{}]".format(new_table_name, self.name))
sqls.append("ALTER TABLE [{}] RENAME TO [{}];".format(new_table_name, self.name))
if should_flip_foreign_keys_pragma:
sqls.append("PRAGMA foreign_key_check")
sqls.append("PRAGMA foreign_keys=ON")
sqls.append("PRAGMA foreign_key_check;")
sqls.append("PRAGMA foreign_keys=ON;")
return sqls