Add pk column if missing from insert

This commit is contained in:
Simon Willison 2019-07-23 13:12:48 +02:00
commit e9d983e561
2 changed files with 7 additions and 0 deletions

View file

@ -224,6 +224,8 @@ class Database:
single_pk = None
if isinstance(pk, str):
single_pk = pk
if pk not in [c[0] for c in column_items]:
column_items.insert(0, (pk, int))
for column_name, column_type in column_items:
column_extras = []
if column_name == single_pk:

View file

@ -734,3 +734,8 @@ def test_cannot_provide_both_filename_and_memory():
AssertionError, match="Either specify a filename_or_conn or pass memory=True"
):
Database("/tmp/foo.db", memory=True)
def test_creates_id_column(fresh_db):
last_pk = fresh_db.table("cats", pk="id").insert({"name": "barry"}).last_pk
assert [{"name": "barry", "id": last_pk}] == list(fresh_db["cats"].rows)