diff --git a/docs/python-api.rst b/docs/python-api.rst index 752e813..b1ad43d 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -331,6 +331,30 @@ The function can accept an iterator or generator of rows and will commit them ac You can skip inserting any records that have a primary key that already exists using ``ignore=True``. This works with both ``.insert({...}, ignore=True)`` and ``.insert_all([...], ignore=True)``. +.. _python_api_update: + +Updating a specific record +========================== + +You can update a record by its primary key using ``table.update()``:: + + >>> db = sqlite_utils.Database("dogs.db") + >>> print(db["dogs"].get(1)) + {'id': 1, 'age': 4, 'name': 'Cleo'} + >>> db["dogs"].update(1, {"age": 5}) + >>> print(db["dogs"].get(1)) + {'id': 1, 'age': 5, 'name': 'Cleo'} + +The first argument to ``update()`` is the primary key. This can be a single value, or a tuple if that table has a compound primary key:: + + >>> db["compound_dogs"].update((5, 3), {"name": "Updated"}) + +The second argument is a dictonary of columns that should be updated, along with their new values. + +You can cause any missing columns to be added automatically using ``alter=True``:: + + >>> db["dogs"].update(1, {"breed": "Mutt"}, alter=True) + Upserting data ============== diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 267d0e9..bbc6177 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1,4 +1,4 @@ -from .utils import sqlite3 +from .utils import sqlite3, OperationalError from collections import namedtuple import datetime import hashlib @@ -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] @@ -784,6 +777,41 @@ class Table: def value_or_default(self, key, value): return self._defaults[key] if value is DEFAULT else value + def update(self, pk_values, updates=None, alter=False): + updates = updates or {} + if not isinstance(pk_values, (list, tuple)): + pk_values = [pk_values] + # Sanity check that the record exists (raises error if not): + self.get(pk_values) + if not updates: + return self + args = [] + sets = [] + wheres = [] + for key, value in updates.items(): + sets.append("[{}] = ?".format(key)) + args.append(value) + 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) + ) + with self.db.conn: + try: + rowcount = self.db.conn.execute(sql, args).rowcount + except OperationalError as e: + if alter and (" column" in e.args[0]): + # Attempt to add any missing columns, then try again + self.add_missing_columns([updates]) + rowcount = self.db.conn.execute(sql, args).rowcount + else: + raise + + # TODO: Test this works (rolls back) - use better exception: + assert rowcount == 1 + self.last_pk = pk_values[0] if len(self.pks) == 1 else pk_values + return self + def insert( self, record, @@ -918,15 +946,15 @@ class Table: with self.db.conn: try: result = self.db.conn.execute(sql, values) - except sqlite3.OperationalError as e: - if alter and (" has no column " in e.args[0]): + except OperationalError as e: + if alter and (" column" in e.args[0]): # Attempt to add any missing columns, then try again self.add_missing_columns(chunk) result = self.db.conn.execute(sql, values) else: raise self.last_rowid = result.lastrowid - self.last_pk = None + self.last_pk = self.last_rowid # self.last_rowid will be 0 if a "INSERT OR IGNORE" happened if (hash_id or pk) and self.last_rowid: row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0] diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 40b1819..738424d 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -1,4 +1,9 @@ try: import pysqlite3 as sqlite3 + import pysqlite3.dbapi2 + + OperationalError = pysqlite3.dbapi2.OperationalError except ImportError: import sqlite3 + + OperationalError = sqlite3.OperationalError diff --git a/tests/test_create.py b/tests/test_create.py index ca73b2d..b4778f3 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -6,12 +6,12 @@ from sqlite_utils.db import ( NoObviousTable, ForeignKey, ) +from sqlite_utils.utils import sqlite3 import collections import datetime import json import pathlib import pytest -import sqlite3 from .utils import collapse_whitespace 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 diff --git a/tests/test_update.py b/tests/test_update.py new file mode 100644 index 0000000..70f8e9e --- /dev/null +++ b/tests/test_update.py @@ -0,0 +1,77 @@ +from sqlite_utils.db import NotFoundError +import pytest + + +def test_update_rowid_table(fresh_db): + table = fresh_db["table"] + rowid = table.insert({"foo": "bar"}).last_pk + table.update(rowid, {"foo": "baz"}) + assert [{"foo": "baz"}] == list(table.rows) + + +def test_update_pk_table(fresh_db): + table = fresh_db["table"] + pk = table.insert({"foo": "bar", "id": 5}, pk="id").last_pk + assert 5 == pk + table.update(pk, {"foo": "baz"}) + assert [{"id": 5, "foo": "baz"}] == list(table.rows) + + +def test_update_compound_pk_table(fresh_db): + table = fresh_db["table"] + pk = table.insert({"id1": 5, "id2": 3, "v": 1}, pk=("id1", "id2")).last_pk + assert (5, 3) == pk + table.update(pk, {"v": 2}) + assert [{"id1": 5, "id2": 3, "v": 2}] == list(table.rows) + + +@pytest.mark.parametrize( + "pk,update_pk", + ( + (None, 2), + (None, None), + ("id1", None), + ("id1", 4), + (("id1", "id2"), None), + (("id1", "id2"), 4), + (("id1", "id2"), (4, 5)), + ), +) +def test_update_invalid_pk(fresh_db, pk, update_pk): + table = fresh_db["table"] + table.insert({"id1": 5, "id2": 3, "v": 1}, pk=pk).last_pk + with pytest.raises(NotFoundError): + table.update(update_pk, {"v": 2}) + + +def test_update_alter(fresh_db): + table = fresh_db["table"] + rowid = table.insert({"foo": "bar"}).last_pk + table.update(rowid, {"new_col": 1.2}, alter=True) + assert [{"foo": "bar", "new_col": 1.2}] == list(table.rows) + # Let's try adding three cols at once + table.update( + rowid, + {"str_col": "str", "bytes_col": b"\xa0 has bytes", "int_col": -10}, + alter=True, + ) + assert [ + { + "foo": "bar", + "new_col": 1.2, + "str_col": "str", + "bytes_col": b"\xa0 has bytes", + "int_col": -10, + } + ] == list(table.rows) + + +def test_update_with_no_values_sets_last_pk(fresh_db): + table = fresh_db.table("dogs", pk="id") + table.insert_all([{"id": 1, "name": "Cleo"}, {"id": 2, "name": "Pancakes"}]) + table.update(1) + assert 1 == table.last_pk + table.update(2) + assert 2 == table.last_pk + with pytest.raises(NotFoundError): + table.update(3)