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

Closes #251, closes #302.
This commit is contained in:
Simon Willison 2021-08-01 21:47:39 -07:00 committed by GitHub
commit 5ec6686153
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 1093 additions and 9 deletions

View file

@ -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,