table.update() method

* Also now set .last_pk to lastrowid for rowid tables
* table.pks introspection now returns ["rowid"] for rowid tables

Closes #35

Author: simonw

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
 ==============
 

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 @@ def rows_where(self, where=None, where_args=None):
 
     @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 @@ def search(self, q):
     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 @@ def insert_all(
             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]

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

tests/test_create.py

@@ -6,12 +6,12 @@
     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
 

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

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)