Use RETURNING star-powered last_pk

Co-authored-by: Audrey Roy Greenfeld <arg@answer.ai>

tests/test_create.py

@@ -165,9 +165,9 @@ def test_create_table_from_example(fresh_db, example, expected_columns):
 
 def test_create_table_from_example_with_compound_primary_keys(fresh_db):
     record = {"name": "Zhang", "group": "staff", "employee_id": 2}
-    records = fresh_db["people"].insert(record, pk=("group", "employee_id"))
-    assert ["group", "employee_id"] == fresh_db["people"].pks
-    assert record == fresh_db["people"].get(("staff", 2))
+    table = fresh_db["people"].insert(record, pk=("group", "employee_id"))
+    assert ["group", "employee_id"] == table.pks
+    assert record == table.get(("staff", 2))
 
 
 @pytest.mark.parametrize(
@@ -930,13 +930,13 @@ def test_insert_ignore(fresh_db):
 
 def test_insert_hash_id(fresh_db):
     dogs = fresh_db["dogs"]
-    id = dogs.insert({"name": "Cleo", "twitter": "cleopaws"}, hash_id="id")[0]['id']
+    id = dogs.insert({"name": "Cleo", "twitter": "cleopaws"}, hash_id="id").last_pk
     assert "f501265970505d9825d8d9f590bfab3519fb20b1" == id
     assert dogs.count == 1
     # Insert replacing a second time should not create a new row
     id2 = dogs.insert(
         {"name": "Cleo", "twitter": "cleopaws"}, hash_id="id", replace=True
-    )[0]['id']
+    ).last_pk
     assert "f501265970505d9825d8d9f590bfab3519fb20b1" == id2
     assert dogs.count == 1
 
@@ -950,21 +950,19 @@ def test_insert_hash_id_columns(fresh_db, use_table_factory):
         dogs = fresh_db["dogs"]
         insert_kwargs = dict(hash_id_columns=("name", "twitter"))
 
-    dogs.insert(
+    id = dogs.insert(
         {"name": "Cleo", "twitter": "cleopaws", "age": 5},
         **insert_kwargs,
-    )
-    id = dogs.last_pk
+    ).last_pk
     expected_hash = hash_record({"name": "Cleo", "twitter": "cleopaws"})
     assert id == expected_hash
     assert dogs.count == 1
     # Insert replacing a second time should not create a new row
-    dogs.insert(
+    id2 = dogs.insert(
         {"name": "Cleo", "twitter": "cleopaws", "age": 6},
         **insert_kwargs,
         replace=True,
-    )
-    id2 = dogs.last_pk
+    ).last_pk
     assert id2 == expected_hash
     assert dogs.count == 1
 
@@ -1060,7 +1058,7 @@ def test_cannot_provide_both_filename_and_memory():
 
 
 def test_creates_id_column(fresh_db):
-    last_pk = fresh_db.table("cats", pk="id").insert({"name": "barry"})[0]['id']
+    last_pk = fresh_db.table("cats", pk="id").insert({"name": "barry"}).last_pk
     assert [{"name": "barry", "id": last_pk}] == list(fresh_db["cats"].rows)
 
 

tests/test_delete.py

@@ -1,7 +1,7 @@
 def test_delete_rowid_table(fresh_db):
     table = fresh_db["table"]
     table.insert({"foo": 1, 'id': 1})
-    rowid = table.insert({"foo": 2, 'id': 2})[0]['id']
+    rowid = table.insert({"foo": 2, 'id': 2}).last_pk
     table.delete(rowid)
     assert [{"foo": 1, 'id': 1}] == list(table.rows)
 

tests/test_duplicate.py

@@ -25,7 +25,7 @@ def test_duplicate(fresh_db):
         "datetime_col": str(dt),
     }
     table1 = fresh_db["table1"]
-    row_id = table1.insert(data)[0]['id']
+    row_id = table1.insert(data).last_pk
     # Duplicate table:
     table2 = table1.duplicate("table2")
     # Ensure data integrity:

tests/test_get.py

@@ -5,14 +5,14 @@
 def test_get_rowid(fresh_db):
     dogs = fresh_db["dogs"]
     cleo = {"name": "Cleo", "age": 4}
-    dogs.insert(cleo)
-    assert cleo == dogs.get(dogs.last_rowid)
+    row_id = dogs.insert(cleo).last_rowid
+    assert cleo == dogs.get(row_id)
 
 
 def test_get_primary_key(fresh_db):
     dogs = fresh_db["dogs"]
     cleo = {"name": "Cleo", "age": 4, "id": 5}
-    last_pk = dogs.insert(cleo, pk="id")[0]['id']
+    last_pk = dogs.insert(cleo, pk="id").last_pk
     assert 5 == last_pk
     assert cleo == dogs.get(5)