From 4ab8d46b03a92c68e9694ea7c285d3852ef58530 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 28 Jul 2019 17:51:49 +0300
Subject: [PATCH] Added table.update(pk, ..., alter=True)

---
 sqlite_utils/db.py    | 19 ++++++++++++++-----
 sqlite_utils/utils.py |  5 +++++
 tests/test_create.py  |  2 +-
 tests/test_update.py  | 20 ++++++++++++++++++++
 4 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 36dec86..b842afe 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
@@ -777,7 +777,7 @@ 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):
+    def update(self, pk_values, updates=None, alter=False):
         updates = updates or {}
         if not isinstance(pk_values, (list, tuple)):
             pk_values = [pk_values]
@@ -795,7 +795,16 @@ class Table:
             table=self.name, sets=", ".join(sets), wheres=" and ".join(wheres)
         )
         with self.db.conn:
-            rowcount = self.db.conn.execute(sql, args).rowcount
+            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
@@ -935,8 +944,8 @@ 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)
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_update.py b/tests/test_update.py
index 9fab8a8..2d46ffe 100644
--- a/tests/test_update.py
+++ b/tests/test_update.py
@@ -44,3 +44,23 @@ def test_update_invalid_pk(fresh_db, pk, update_pk):
         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)