Store list/dict/tuple values as JSON strings

This commit is contained in:
Simon Willison 2018-07-28 15:20:29 -07:00
commit 95bce37ad3
2 changed files with 36 additions and 4 deletions

View file

@ -1,5 +1,6 @@
import sqlite3
from collections import namedtuple
import json
Column = namedtuple(
"Column", ("cid", "name", "type", "notnull", "default_value", "is_pk")
@ -29,7 +30,7 @@ class Database:
foreign_keys = foreign_keys or []
foreign_keys_by_name = {fk[0]: fk for fk in foreign_keys}
extra = ""
columns = ",\n".join(
columns_sql = ",\n".join(
" {col_name} {col_type} {primary_key} {references}".format(
col_name=col_name,
col_type={
@ -52,10 +53,10 @@ class Database:
for col_name, col_type in columns.items()
)
sql = """CREATE TABLE {table} (
{columns}
{columns_sql}
){extra};
""".format(
table=name, columns=columns, extra=extra
table=name, columns_sql=columns_sql, extra=extra
)
self.conn.execute(sql)
return self[name]
@ -108,6 +109,10 @@ class Table:
for key, types in all_column_types.items():
if len(types) == 1:
t = list(types)[0]
# But if it's list / tuple / dict, use str instead as we
# will be storing it as JSON in the table
if t in (list, tuple, dict):
t = str
elif {int, bool}.issuperset(types):
t = int
elif {int, float, bool}.issuperset(types):
@ -154,7 +159,9 @@ class Table:
)
values = []
for record in chunk:
values.extend(record.get(key, None) for key in all_columns)
values.extend(
jsonify_if_needed(record.get(key, None)) for key in all_columns
)
result = self.db.conn.execute(sql, values)
self.db.conn.commit()
return result
@ -169,3 +176,10 @@ class Table:
def chunks(sequence, size):
for i in range(0, len(sequence), size):
yield sequence[i : i + size]
def jsonify_if_needed(value):
if isinstance(value, (dict, list, tuple)):
return json.dumps(value)
else:
return value

View file

@ -1,4 +1,5 @@
from sqlite_utils import db
import json
import sqlite3
import pytest
@ -54,3 +55,20 @@ def test_create_table_works_for_m2m_with_only_foreign_keys(fresh_db):
{"name": "one_id", "type": "INTEGER"},
{"name": "two_id", "type": "INTEGER"},
] == [{"name": col.name, "type": col.type} for col in fresh_db["m2m"].columns]
@pytest.mark.parametrize(
"data_structure",
(
["list with one item"],
["list with", "two items"],
{"dictionary": "simple"},
{"dictionary": {"nested": "complex"}},
[{"list": "of"}, {"two": "dicts"}],
),
)
def test_insert_dictionaries_and_lists_as_json(fresh_db, data_structure):
fresh_db["test"].insert({"id": 1, "data": data_structure}, pk="id")
row = fresh_db.conn.execute("select id, data from test").fetchone()
assert row[0] == 1
assert data_structure == json.loads(row[1])