diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index e97f71d..36dec86 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -456,31 +456,24 @@ class Table: @property def pks(self): - return [column.name for column in self.columns if column.is_pk] + names = [column.name for column in self.columns if column.is_pk] + if not names: + names = ["rowid"] + return names def get(self, pk_values): if not isinstance(pk_values, (list, tuple)): pk_values = [pk_values] pks = self.pks - pk_names = [] - if len(pks) == 0: - # rowid table - pk_names = ["rowid"] - last_pk = pk_values[0] - elif len(pks) == 1: - pk_names = [pks[0]] - last_pk = pk_values[0] - elif len(pks) > 1: - pk_names = pks - last_pk = pk_values - if len(pk_names) != len(pk_values): + last_pk = pk_values[0] if len(pks) == 1 else pk_values + if len(pks) != len(pk_values): raise NotFoundError( "Need {} primary key value{}".format( - len(pk_names), "" if len(pk_names) == 1 else "s" + len(pks), "" if len(pks) == 1 else "s" ) ) - wheres = ["[{}] = ?".format(pk_name) for pk_name in pk_names] + wheres = ["[{}] = ?".format(pk_name) for pk_name in pks] rows = self.rows_where(" and ".join(wheres), pk_values) try: row = list(rows)[0] @@ -788,26 +781,15 @@ class Table: updates = updates or {} if not isinstance(pk_values, (list, tuple)): pk_values = [pk_values] - pks = self.pks - pk_names = [] - if len(pks) == 0: - # rowid table - pk_names = ["rowid"] - last_pk = pk_values[0] - elif len(pks) == 1: - pk_names = [pks[0]] - last_pk = pk_values[0] - elif len(pks) > 1: - pk_names = pks - last_pk = pk_values - assert len(pk_names) == len(pk_values) + # Sanity check that the record exists (raises error if not): + self.get(pk_values) args = [] sets = [] wheres = [] for key, value in updates.items(): sets.append("[{}] = ?".format(key)) args.append(value) - wheres = ["[{}] = ?".format(pk_name) for pk_name in pk_names] + wheres = ["[{}] = ?".format(pk_name) for pk_name in self.pks] args.extend(pk_values) sql = "update [{table}] set {sets} where {wheres}".format( table=self.name, sets=", ".join(sets), wheres=" and ".join(wheres) @@ -816,7 +798,7 @@ class Table: rowcount = self.db.conn.execute(sql, args).rowcount # TODO: Test this works (rolls back) - use better exception: assert rowcount == 1 - self.last_pk = last_pk + self.last_pk = pk_values[0] if len(self.pks) == 1 else pk_values return self def insert( diff --git a/tests/test_introspect.py b/tests/test_introspect.py index 0dd544e..d3ab27a 100644 --- a/tests/test_introspect.py +++ b/tests/test_introspect.py @@ -98,3 +98,11 @@ def test_guess_foreign_table(fresh_db, column, expected_table_guess): fresh_db.create_table("authors", {"name": str}) fresh_db.create_table("genre", {"name": str}) assert expected_table_guess == fresh_db["books"].guess_foreign_table(column) + + +@pytest.mark.parametrize( + "pk,expected", ((None, ["rowid"]), ("id", ["id"]), (["id", "id2"], ["id", "id2"])) +) +def test_pks(fresh_db, pk, expected): + fresh_db["foo"].insert_all([{"id": 1, "id2": 2}], pk=pk) + assert expected == fresh_db["foo"].pks