mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-23 01:14:31 +02:00
Table options can now be passed to constructor OR to insert_all()
If you want to set default options for a table, you can do this:
table = db.table("dogs", pk="id", column_order=["name", "age"])
If you pass those keyword arguments to the .insert/.update/etc
methods they will over-ride the defaults you set on the table.
table = db["dogs"] # This still works too
This commit is contained in:
Simon Willison
parent
152eb2afaf
commit
57e43baece
2 changed files with 142 additions and 74 deletions
|
|
@ -18,6 +18,7 @@ ForeignKey = namedtuple(
|
|||
"ForeignKey", ("table", "column", "other_table", "other_column")
|
||||
)
|
||||
Index = namedtuple("Index", ("seq", "name", "unique", "origin", "partial", "columns"))
|
||||
DEFAULT = object()
|
||||
|
||||
COLUMN_TYPE_MAPPING = {
|
||||
float: "FLOAT",
|
||||
|
|
@ -93,11 +94,14 @@ class Database:
|
|||
self.conn = filename_or_conn
|
||||
|
||||
def __getitem__(self, table_name):
|
||||
return Table(self, table_name)
|
||||
return self.table(table_name)
|
||||
|
||||
def __repr__(self):
|
||||
return "<Database {}>".format(self.conn)
|
||||
|
||||
def table(self, table_name, **kwargs):
|
||||
return Table(self, table_name, **kwargs)
|
||||
|
||||
def escape(self, value):
|
||||
# Normally we would use .execute(sql, [params]) for escaping, but
|
||||
# occasionally that isn't available - most notable when we need
|
||||
|
|
@ -343,10 +347,36 @@ class Database:
|
|||
|
||||
|
||||
class Table:
|
||||
def __init__(self, db, name):
|
||||
def __init__(
|
||||
self,
|
||||
db,
|
||||
name,
|
||||
pk=None,
|
||||
foreign_keys=None,
|
||||
column_order=None,
|
||||
not_null=None,
|
||||
defaults=None,
|
||||
upsert=False,
|
||||
batch_size=100,
|
||||
hash_id=None,
|
||||
alter=False,
|
||||
ignore=False,
|
||||
):
|
||||
self.db = db
|
||||
self.name = name
|
||||
self.exists = self.name in self.db.table_names()
|
||||
self._defaults = dict(
|
||||
pk=pk,
|
||||
foreign_keys=foreign_keys,
|
||||
column_order=column_order,
|
||||
not_null=not_null,
|
||||
defaults=defaults,
|
||||
upsert=upsert,
|
||||
batch_size=batch_size,
|
||||
hash_id=hash_id,
|
||||
alter=alter,
|
||||
ignore=ignore,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return "<Table {}{}>".format(
|
||||
|
|
@ -716,18 +746,21 @@ class Table:
|
|||
)
|
||||
return self.db.conn.execute(sql, (q,)).fetchall()
|
||||
|
||||
def value_or_default(self, key, value):
|
||||
return self._defaults[key] if value is DEFAULT else value
|
||||
|
||||
def insert(
|
||||
self,
|
||||
record,
|
||||
pk=None,
|
||||
foreign_keys=None,
|
||||
column_order=None,
|
||||
not_null=None,
|
||||
defaults=None,
|
||||
upsert=False,
|
||||
hash_id=None,
|
||||
alter=False,
|
||||
ignore=False,
|
||||
pk=DEFAULT,
|
||||
foreign_keys=DEFAULT,
|
||||
column_order=DEFAULT,
|
||||
not_null=DEFAULT,
|
||||
defaults=DEFAULT,
|
||||
upsert=DEFAULT,
|
||||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
ignore=DEFAULT,
|
||||
):
|
||||
return self.insert_all(
|
||||
[record],
|
||||
|
|
@ -745,22 +778,33 @@ class Table:
|
|||
def insert_all(
|
||||
self,
|
||||
records,
|
||||
pk=None,
|
||||
foreign_keys=None,
|
||||
column_order=None,
|
||||
not_null=None,
|
||||
defaults=None,
|
||||
upsert=False,
|
||||
batch_size=100,
|
||||
hash_id=None,
|
||||
alter=False,
|
||||
ignore=False,
|
||||
pk=DEFAULT,
|
||||
foreign_keys=DEFAULT,
|
||||
column_order=DEFAULT,
|
||||
not_null=DEFAULT,
|
||||
defaults=DEFAULT,
|
||||
upsert=DEFAULT,
|
||||
batch_size=DEFAULT,
|
||||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
ignore=DEFAULT,
|
||||
):
|
||||
"""
|
||||
Like .insert() but takes a list of records and ensures that the table
|
||||
that it creates (if table does not exist) has columns for ALL of that
|
||||
data
|
||||
"""
|
||||
pk = self.value_or_default("pk", pk)
|
||||
foreign_keys = self.value_or_default("foreign_keys", foreign_keys)
|
||||
column_order = self.value_or_default("column_order", column_order)
|
||||
not_null = self.value_or_default("not_null", not_null)
|
||||
defaults = self.value_or_default("defaults", defaults)
|
||||
upsert = self.value_or_default("upsert", upsert)
|
||||
batch_size = self.value_or_default("batch_size", batch_size)
|
||||
hash_id = self.value_or_default("hash_id", hash_id)
|
||||
alter = self.value_or_default("alter", alter)
|
||||
ignore = self.value_or_default("ignore", ignore)
|
||||
|
||||
assert not (hash_id and pk), "Use either pk= or hash_id="
|
||||
assert not (
|
||||
ignore and upsert
|
||||
|
|
@ -842,13 +886,13 @@ class Table:
|
|||
def upsert(
|
||||
self,
|
||||
record,
|
||||
pk=None,
|
||||
foreign_keys=None,
|
||||
column_order=None,
|
||||
not_null=None,
|
||||
defaults=None,
|
||||
hash_id=None,
|
||||
alter=False,
|
||||
pk=DEFAULT,
|
||||
foreign_keys=DEFAULT,
|
||||
column_order=DEFAULT,
|
||||
not_null=DEFAULT,
|
||||
defaults=DEFAULT,
|
||||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
):
|
||||
return self.insert(
|
||||
record,
|
||||
|
|
@ -865,14 +909,14 @@ class Table:
|
|||
def upsert_all(
|
||||
self,
|
||||
records,
|
||||
pk=None,
|
||||
foreign_keys=None,
|
||||
column_order=None,
|
||||
not_null=None,
|
||||
defaults=None,
|
||||
batch_size=100,
|
||||
hash_id=None,
|
||||
alter=False,
|
||||
pk=DEFAULT,
|
||||
foreign_keys=DEFAULT,
|
||||
column_order=DEFAULT,
|
||||
not_null=DEFAULT,
|
||||
defaults=DEFAULT,
|
||||
batch_size=DEFAULT,
|
||||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
):
|
||||
return self.insert_all(
|
||||
records,
|
||||
|
|
|
|||
|
|
@ -144,19 +144,22 @@ def test_create_table_from_example_with_compound_primary_keys(fresh_db):
|
|||
assert record == table.get(("staff", 2))
|
||||
|
||||
|
||||
def test_create_table_column_order(fresh_db):
|
||||
fresh_db["table"].insert(
|
||||
collections.OrderedDict(
|
||||
(
|
||||
("zzz", "third"),
|
||||
("abc", "first"),
|
||||
("ccc", "second"),
|
||||
("bbb", "second-to-last"),
|
||||
("aaa", "last"),
|
||||
)
|
||||
),
|
||||
column_order=("abc", "ccc", "zzz"),
|
||||
@pytest.mark.parametrize("use_class_constructor", [True, False])
|
||||
def test_create_table_column_order(fresh_db, use_class_constructor):
|
||||
row = collections.OrderedDict(
|
||||
(
|
||||
("zzz", "third"),
|
||||
("abc", "first"),
|
||||
("ccc", "second"),
|
||||
("bbb", "second-to-last"),
|
||||
("aaa", "last"),
|
||||
)
|
||||
)
|
||||
column_order = ("abc", "ccc", "zzz")
|
||||
if use_class_constructor:
|
||||
fresh_db.table("table", column_order=column_order).insert(row)
|
||||
else:
|
||||
fresh_db["table"].insert(row, column_order=column_order)
|
||||
assert [
|
||||
{"name": "abc", "type": "TEXT"},
|
||||
{"name": "ccc", "type": "TEXT"},
|
||||
|
|
@ -191,21 +194,31 @@ def test_create_table_column_order(fresh_db):
|
|||
(({"one_id": "one"},), AssertionError),
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("use_class_constructor", [True, False])
|
||||
def test_create_table_works_for_m2m_with_only_foreign_keys(
|
||||
fresh_db, foreign_key_specification, expected_exception
|
||||
fresh_db, foreign_key_specification, expected_exception, use_class_constructor
|
||||
):
|
||||
fresh_db["one"].insert({"id": 1}, pk="id")
|
||||
fresh_db["two"].insert({"id": 1}, pk="id")
|
||||
if use_class_constructor:
|
||||
fresh_db.table("one", pk="id").insert({"id": 1})
|
||||
fresh_db.table("two", pk="id").insert({"id": 1})
|
||||
else:
|
||||
fresh_db["one"].insert({"id": 1}, pk="id")
|
||||
fresh_db["two"].insert({"id": 1}, pk="id")
|
||||
|
||||
row = {"one_id": 1, "two_id": 1}
|
||||
|
||||
def do_it():
|
||||
if use_class_constructor:
|
||||
fresh_db.table("m2m", foreign_keys=foreign_key_specification).insert(row)
|
||||
else:
|
||||
fresh_db["m2m"].insert(row, foreign_keys=foreign_key_specification)
|
||||
|
||||
if expected_exception:
|
||||
with pytest.raises(expected_exception):
|
||||
fresh_db["m2m"].insert(
|
||||
{"one_id": 1, "two_id": 1}, foreign_keys=foreign_key_specification
|
||||
)
|
||||
do_it()
|
||||
return
|
||||
else:
|
||||
fresh_db["m2m"].insert(
|
||||
{"one_id": 1, "two_id": 1}, foreign_keys=foreign_key_specification
|
||||
)
|
||||
do_it()
|
||||
assert [
|
||||
{"name": "one_id", "type": "INTEGER"},
|
||||
{"name": "two_id", "type": "INTEGER"},
|
||||
|
|
@ -407,7 +420,10 @@ def test_index_foreign_keys(fresh_db):
|
|||
),
|
||||
],
|
||||
)
|
||||
def test_insert_row_alter_table(fresh_db, extra_data, expected_new_columns):
|
||||
@pytest.mark.parametrize("use_class_constructor", [True, False])
|
||||
def test_insert_row_alter_table(
|
||||
fresh_db, extra_data, expected_new_columns, use_class_constructor
|
||||
):
|
||||
table = fresh_db["books"]
|
||||
table.insert({"title": "Hedgehogs of the world", "author_id": 1})
|
||||
assert [
|
||||
|
|
@ -416,7 +432,10 @@ def test_insert_row_alter_table(fresh_db, extra_data, expected_new_columns):
|
|||
] == [{"name": col.name, "type": col.type} for col in table.columns]
|
||||
record = {"title": "Squirrels of the world", "author_id": 2}
|
||||
record.update(extra_data)
|
||||
fresh_db["books"].insert(record, alter=True)
|
||||
if use_class_constructor:
|
||||
fresh_db.table("books", alter=True).insert(record)
|
||||
else:
|
||||
fresh_db["books"].insert(record, alter=True)
|
||||
assert [
|
||||
{"name": "title", "type": "TEXT"},
|
||||
{"name": "author_id", "type": "INTEGER"},
|
||||
|
|
@ -425,21 +444,26 @@ def test_insert_row_alter_table(fresh_db, extra_data, expected_new_columns):
|
|||
]
|
||||
|
||||
|
||||
def test_upsert_rows_alter_table(fresh_db):
|
||||
table = fresh_db["books"]
|
||||
table.insert({"id": 1, "title": "Hedgehogs of the world", "author_id": 1}, pk="id")
|
||||
table.upsert_all(
|
||||
[
|
||||
{"id": 1, "title": "Hedgehogs of the World", "species": "hedgehogs"},
|
||||
{"id": 2, "title": "Squirrels of the World", "num_species": 200},
|
||||
{
|
||||
"id": 3,
|
||||
"title": "Badgers of the World",
|
||||
"significant_continents": ["Europe", "North America"],
|
||||
},
|
||||
],
|
||||
alter=True,
|
||||
)
|
||||
@pytest.mark.parametrize("use_class_constructor", [True, False])
|
||||
def test_upsert_rows_alter_table(fresh_db, use_class_constructor):
|
||||
first_row = {"id": 1, "title": "Hedgehogs of the world", "author_id": 1}
|
||||
next_rows = [
|
||||
{"id": 1, "title": "Hedgehogs of the World", "species": "hedgehogs"},
|
||||
{"id": 2, "title": "Squirrels of the World", "num_species": 200},
|
||||
{
|
||||
"id": 3,
|
||||
"title": "Badgers of the World",
|
||||
"significant_continents": ["Europe", "North America"],
|
||||
},
|
||||
]
|
||||
if use_class_constructor:
|
||||
table = fresh_db.table("books", pk="id", alter=True)
|
||||
table.insert(first_row)
|
||||
table.upsert_all(next_rows)
|
||||
else:
|
||||
table = fresh_db["books"]
|
||||
table.insert(first_row, pk="id")
|
||||
table.upsert_all(next_rows, alter=True)
|
||||
assert {
|
||||
"author_id": int,
|
||||
"id": int,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue