Progress bar for "sqlite-utils extract", closes #169

This commit is contained in:
Simon Willison 2020-09-22 17:02:29 -07:00
commit 5c4d58d152
3 changed files with 26 additions and 3 deletions

View file

@ -650,6 +650,9 @@ The command takes the following options:
``--rename <TEXT TEXT>``
Use this option to rename the columns created in the new lookup table.
``--silent``
Don't display the progress bar.
Here's a more complex example that makes use of these options. It converts `this CSV file <https://github.com/wri/global-power-plant-database/blob/232a666653e14d803ab02717efc01cdd437e7601/output_database/global_power_plant_database.csv>`__ full of global power plants into SQLite, then extracts the ``country`` and ``country_long`` columns into a separate ``countries`` table::
wget 'https://github.com/wri/global-power-plant-database/blob/232a6666/output_database/global_power_plant_database.csv?raw=true'

View file

@ -1006,6 +1006,12 @@ def transform(
multiple=True,
help="Rename this column in extracted table",
)
@click.option(
"-s",
"--silent",
is_flag=True,
help="Don't show progress bar",
)
def extract(
path,
table,
@ -1013,12 +1019,24 @@ def extract(
other_table,
fk_column,
rename,
silent,
):
"Extract one or more columns into a separate table"
db = sqlite_utils.Database(path)
db[table].extract(
columns, table=other_table, fk_column=fk_column, rename=dict(rename)
kwargs = dict(
columns=columns,
table=other_table,
fk_column=fk_column,
rename=dict(rename),
)
if silent:
db[table].extract(**kwargs)
else:
with click.progressbar(
length=db[table].count, label="Extracting columns"
) as bar:
kwargs["progress"] = bar.update
db[table].extract(**kwargs)
@cli.command(name="insert-files")

View file

@ -883,7 +883,7 @@ class Table(Queryable):
return sqls
def extract(self, columns, table=None, fk_column=None, rename=None):
def extract(self, columns, table=None, fk_column=None, rename=None, progress=None):
rename = rename or {}
if isinstance(columns, str):
columns = [columns]
@ -905,6 +905,8 @@ class Table(Queryable):
row_pks = tuple(row[pk] for pk in pks)
lookups = {rename.get(column) or column: row[column] for column in columns}
self.update(row_pks, {first_column: lookup_table.lookup(lookups)})
if progress:
progress(1)
fk_column = fk_column or "{}_id".format(table)
# Now rename first_column and change its type to integer, and drop
# any other extracted columns: