remote/sqlite-utils
1
0
Fork 0
mirror of https://github.com/simonw/sqlite-utils.git synced 2026-07-26 19:04:32 +02:00
Code Issues Projects Releases Packages Wiki Activity

Implemented .convert(multi=True) and refactored CLI to use that method, refs #251, #302

Browse source
This commit is contained in:
Simon Willison 2021-08-01 21:08:40 -07:00
commit 570bee7edd
3 changed files with 90 additions and 125 deletions
Download patch file Download diff file Expand all files Collapse all files

146
sqlite_utils/cli.py
View file

@ -5,7 +5,7 @@ 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
@ -23,6 +23,7 @@ from .utils import (
sqlite3,
decode_base64_values,
rows_from_file,
types_for_column_types,
Format,
TypeTracker,
)
@ -1983,6 +1984,7 @@ def convert(
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:
@ -2005,8 +2007,7 @@ def convert(
fn = locals["fn"]
if dry_run:
# Pull first 20 values for first column and preview them
db = sqlite3.connect(db_path)
db.create_function("preview_transform", 1, lambda v: fn(v) if v else v)
db.conn.create_function("preview_transform", 1, lambda v: fn(v) if v else v)
sql = """
select
[{column}] as value,
@ -2015,129 +2016,28 @@ def convert(
""".format(
column=columns[0], table=table
)
for row in db.execute(sql).fetchall():
print(row[0])
print(" --- becomes:")
print(row[1])
print()
elif multi:
_transform_multi(db_path, table, columns[0], fn, drop, silent)
for row in db.conn.execute(sql).fetchall():
click.echo(str(row[0]))
click.echo(" --- becomes:")
click.echo(str(row[1]))
click.echo()
else:
_transform(db_path, table, columns, fn, output, output_type, drop, silent)
def _transform(db_path, table, columns, fn, output, output_type, drop, silent):
db = sqlite_utils.Database(db_path)
count_sql = "select count(*) from [{}]".format(table)
todo_count = list(db.execute(count_sql).fetchall())[0][0] * len(columns)
if drop and not output:
raise click.ClickException("--drop can only be used with --output or --multi")
if output is not None:
if output not in db[table].columns_dict:
db[table].add_column(output, output_type or "text")
with progressbar(length=todo_count, silent=silent) as bar:
def transform_value(v):
bar.update(1)
if not v:
return v
return fn(v)
db.register_function(transform_value)
sql = "update [{table}] set {sets};".format(
table=table,
sets=", ".join(
[
"[{output_column}] = transform_value([{column}])".format(
output_column=output or column, column=column
)
for column in columns
]
),
)
with db.conn:
db.execute(sql)
if drop:
db[table].transform(drop=columns)
def _transform_multi(db_path, table, column, fn, drop, silent):
db = sqlite_utils.Database(db_path)
# First we execute the function
pk_to_values = {}
new_column_types = {}
pks = [column.name for column in db[table].columns if column.is_pk]
if not pks:
pks = ["rowid"]
with progressbar(
length=db[table].count, silent=silent, label="1: Evaluating"
) as bar:
for row in db[table].rows_where(
select=", ".join(
"[{}]".format(column_name) for column_name in (pks + [column])
try:
db[table].convert(
columns,
fn,
output=output,
output_type=output_type,
drop=drop,
multi=multi,
show_progress=not silent,
)
):
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 click.ClickException(
"With --multi code must return a Python dictionary - returned {}".format(
repr(values)
)
except BadMultiValues as e:
raise click.ClickException(
"When using --multi code must return a Python dictionary - returned: {}".format(
repr(e.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 = _suggest_column_types(new_column_types)
for column_name, column_type in columns_to_create.items():
if column_name not in db[table].columns_dict:
db[table].add_column(column_name, column_type)
# Run the updates
with progressbar(length=db[table].count, silent=silent, label="2: Updating") as bar:
with db.conn:
for pk, updates in pk_to_values.items():
db[table].update(pk, updates)
bar.update(1)
if drop:
db[table].transform(drop=(column,))
def _suggest_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:
types.discard(None.__class__)
if {None.__class__} == types:
t = str
elif len(types) == 1:
t = list(types)[0]
# But if it's a subclass of list / tuple / dict, use str
# instead as we will be storing it as JSON in the table
for superclass in (list, tuple, dict):
if issubclass(t, superclass):
t = str
elif {int, bool}.issuperset(types):
t = int
elif {int, float, bool}.issuperset(types):
t = float
elif {bytes, str}.issuperset(types):
t = bytes
else:
t = str
column_types[key] = t
return column_types
)
def _render_common(title, values):

64
sqlite_utils/db.py
View file

@ -2,11 +2,13 @@ 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 click
import contextlib
import datetime
import decimal
@ -159,6 +161,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,
@ -1709,11 +1718,18 @@ class Table(Queryable):
fn,
output=None,
output_type=None,
show_progress=False,
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
)
todo_count = self.count * len(columns)
if output is not None:
@ -1744,6 +1760,52 @@ class Table(Queryable):
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,

5
sqlite_utils/utils.py
View file

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