sqlite-utils convert command and db[table].convert(...) method

Closes #251, closes #302.

sqlite_utils/cli.py

@@ -5,8 +5,10 @@ from datetime import datetime
 import hashlib
 import pathlib
 import sqlite_utils
-from sqlite_utils.db import AlterError, DescIndex
+from sqlite_utils.db import AlterError, BadMultiValues, DescIndex
+from sqlite_utils import recipes
 import textwrap
+import inspect
 import io
 import itertools
 import json
@@ -1903,6 +1905,139 @@ def analyze_tables(
         click.echo(details)
 
 
+def _generate_convert_help():
+    help = textwrap.dedent(
+        """
+    Convert columns using Python code you supply. For example:
+
+    \b
+    $ sqlite-utils convert my.db mytable mycolumn \\
+        '"\\n".join(textwrap.wrap(value, 10))' \\
+        --import=textwrap
+
+    "value" is a variable with the column value to be converted.
+
+    The following common operations are available as recipe functions:
+    """
+    ).strip()
+    recipe_names = [
+        n for n in dir(recipes) if not n.startswith("_") and n not in ("json", "parser")
+    ]
+    for name in recipe_names:
+        fn = getattr(recipes, name)
+        help += "\n\nr.{}{}\n\n  {}".format(
+            name, str(inspect.signature(fn)), fn.__doc__
+        )
+    help += "\n\n"
+    help += textwrap.dedent(
+        """
+    You can use these recipes like so:
+
+    \b
+    $ sqlite-utils convert my.db mytable mycolumn \\
+        'r.jsonsplit(value, delimiter=":")'
+    """
+    ).strip()
+    return help
+
+
+@cli.command(help=_generate_convert_help())
+@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 convert(
+    db_path,
+    table,
+    columns,
+    code,
+    imports,
+    dry_run,
+    multi,
+    output,
+    output_type,
+    drop,
+    silent,
+):
+    sqlite3.enable_callback_tracebacks(True)
+    db = sqlite_utils.Database(db_path)
+    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 drop and not (output or multi):
+        raise click.ClickException("--drop can only be used with --output or --multi")
+    # 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 = {"r": recipes, "recipes": recipes}
+    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.conn.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.conn.execute(sql).fetchall():
+            click.echo(str(row[0]))
+            click.echo(" --- becomes:")
+            click.echo(str(row[1]))
+            click.echo()
+    else:
+        try:
+            db[table].convert(
+                columns,
+                fn,
+                output=output,
+                output_type=output_type,
+                drop=drop,
+                multi=multi,
+                show_progress=not silent,
+            )
+        except BadMultiValues as e:
+            raise click.ClickException(
+                "When using --multi code must return a Python dictionary - returned: {}".format(
+                    repr(e.values)
+                )
+            )
+
+
 def _render_common(title, values):
     if values is None:
         return ""

sqlite_utils/db.py

@@ -1,4 +1,11 @@
-from .utils import sqlite3, OperationalError, suggest_column_types, column_affinity
+from .utils import (
+    sqlite3,
+    OperationalError,
+    suggest_column_types,
+    types_for_column_types,
+    column_affinity,
+    progressbar,
+)
 from collections import namedtuple
 from collections.abc import Mapping
 import contextlib
@@ -153,6 +160,13 @@ class DescIndex(str):
     pass
 
 
+class BadMultiValues(Exception):
+    "With multi=True code must return a Python dictionary"
+
+    def __init__(self, values):
+        self.values = values
+
+
 _COUNTS_TABLE_CREATE_SQL = """
 CREATE TABLE IF NOT EXISTS [{}](
    [table] TEXT PRIMARY KEY,
@@ -1697,6 +1711,101 @@ class Table(Queryable):
         self.last_pk = pk_values[0] if len(pks) == 1 else pk_values
         return self
 
+    def convert(
+        self,
+        columns,
+        fn,
+        output=None,
+        output_type=None,
+        drop=False,
+        multi=False,
+        show_progress=False,
+    ):
+        if isinstance(columns, str):
+            columns = [columns]
+
+        if multi:
+            return self._convert_multi(
+                columns[0], fn, drop=drop, show_progress=show_progress
+            )
+
+        if output is not None:
+            assert len(columns) == 1, "output= can only be used with a single column"
+            if output not in self.columns_dict:
+                self.add_column(output, output_type or "text")
+
+        todo_count = self.count * len(columns)
+        with progressbar(length=todo_count, silent=not show_progress) as bar:
+
+            def convert_value(v):
+                bar.update(1)
+                if not v:
+                    return v
+                return fn(v)
+
+            self.db.register_function(convert_value)
+            sql = "update [{table}] set {sets};".format(
+                table=self.name,
+                sets=", ".join(
+                    [
+                        "[{output_column}] = convert_value([{column}])".format(
+                            output_column=output or column, column=column
+                        )
+                        for column in columns
+                    ]
+                ),
+            )
+            with self.db.conn:
+                self.db.execute(sql)
+                if drop:
+                    self.transform(drop=columns)
+        return self
+
+    def _convert_multi(self, column, fn, drop, show_progress):
+        # First we execute the function
+        pk_to_values = {}
+        new_column_types = {}
+        pks = [column.name for column in self.columns if column.is_pk]
+        if not pks:
+            pks = ["rowid"]
+
+        with progressbar(
+            length=self.count, silent=not show_progress, label="1: Evaluating"
+        ) as bar:
+            for row in self.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 BadMultiValues(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 = types_for_column_types(new_column_types)
+        for column_name, column_type in columns_to_create.items():
+            if column_name not in self.columns_dict:
+                self.add_column(column_name, column_type)
+
+        # Run the updates
+        with progressbar(
+            length=self.count, silent=not show_progress, label="2: Updating"
+        ) as bar:
+            with self.db.conn:
+                for pk, updates in pk_to_values.items():
+                    self.update(pk, updates)
+                    bar.update(1)
+                if drop:
+                    self.transform(drop=(column,))
+
     def build_insert_queries_and_params(
         self,
         extracts,

sqlite_utils/recipes.py

@@ -0,0 +1,19 @@
+from dateutil import parser
+import json
+
+
+def parsedate(value, dayfirst=False, yearfirst=False):
+    "Parse a date and convert it to ISO date format: yyyy-mm-dd"
+    return (
+        parser.parse(value, dayfirst=dayfirst, yearfirst=yearfirst).date().isoformat()
+    )
+
+
+def parsedatetime(value, dayfirst=False, yearfirst=False):
+    "Parse a datetime and convert it to ISO datetime format: yyyy-mm-ddTHH:MM:SS"
+    return parser.parse(value, dayfirst=dayfirst, yearfirst=yearfirst).isoformat()
+
+
+def jsonsplit(value, delimiter=",", type=str):
+    'Convert a string like a,b,c into a JSON array ["a", "b", "c"]'
+    return json.dumps([type(s.strip()) for s in value.split(delimiter)])

sqlite_utils/utils.py

@@ -31,8 +31,11 @@ def suggest_column_types(records):
     for record in records:
         for key, value in record.items():
             all_column_types.setdefault(key, set()).add(type(value))
-    column_types = {}
+    return types_for_column_types(all_column_types)
 
+
+def types_for_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:
@@ -254,3 +257,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