From 5c4d58d1528367c15ec6490024bf2658f251acd3 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 22 Sep 2020 17:02:29 -0700
Subject: [PATCH] Progress bar for "sqlite-utils extract", closes #169

---
 docs/cli.rst        |  3 +++
 sqlite_utils/cli.py | 22 ++++++++++++++++++++--
 sqlite_utils/db.py  |  4 +++-
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/docs/cli.rst b/docs/cli.rst
index 8971fe9..85fba97 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -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'
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index adf33ef..9ff56f5 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -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")
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 2430f0d..e9d5298 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -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: