From 2cfc29931d299cc9527586788c3663f092d4d416 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 9 Jul 2019 07:36:37 -0700
Subject: [PATCH 1/2] WIP

---
 sqlite_utils/db.py | 49 ++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 45 insertions(+), 4 deletions(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 3a44cb4..9354e49 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -79,6 +79,10 @@ class BadPrimaryKey(Exception):
     pass
 
 
+class RowNotFound(Exception):
+    pass
+
+
 class Database:
     def __init__(self, filename_or_conn):
         if isinstance(filename_or_conn, str):
@@ -780,12 +784,26 @@ class Table:
                 self.last_pk = None
                 # self.last_rowid will be 0 if a "INSERT OR IGNORE" happened
                 if (hash_id or pk) and self.last_rowid:
-                    self.last_pk = self.db.conn.execute(
-                        "select [{}] from [{}] where rowid = ?".format(
-                            hash_id or pk, self.name
+                    compound_pk = False
+                    if hash_id:
+                        pk_cols = [hash_id]
+                    elif isinstance(pk, (list, tuple)):
+                        pk_cols = pk
+                        compound_pk = True
+                    else:
+                        pk_cols = [pk]
+                    result = self.db.conn.execute(
+                        "select {} from [{}] where rowid = ?".format(
+                            ", ".join("[{}]".format(pk_col) for pk_col in pk_cols),
+                            self.name,
                         ),
                         (self.last_rowid,),
-                    ).fetchone()[0]
+                    ).fetchone()
+                    if compound_pk:
+                        self.last_pk = list(result)
+                    else:
+                        self.last_pk = result[0]
+
         return self
 
     def upsert(
@@ -836,6 +854,29 @@ class Table:
             upsert=True,
         )
 
+    def update(self, pk, updates=None):
+        updates = updates or {}
+        row = self.get(pk)
+
+    def get(self, pk):
+        "pk is the primary key value, or a tuple/list for compound pks"
+        if not isinstance(pk, (list, tuple)):
+            pk = [pk]
+        # Confirm pk has the correct number of values
+        if len(pk) != len(self.pks):
+            raise RowNotFound(
+                "Primary key value needs to match ({})".format(
+                    ", ".join(map(str, self.pks))
+                )
+            )
+        sql = "SELECT * FROM {} WHERE {}".format(
+            self.name, " AND ".join("[{}] = ?".format(pk) for pk in self.pks)
+        )
+        rows = self.db.execute_returning_dicts(sql, pk)
+        if len(rows) != 1:
+            raise RowNotFound("Row not found for {}".format(repr(pk)))
+        return rows[0]
+
     def add_missing_columns(self, records):
         needed_columns = self.detect_column_types(records)
         current_columns = self.columns_dict

From facec8df62e49b11a593dd89ae8ed6a3c1cb53ca Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 9 Jul 2019 07:44:47 -0700
Subject: [PATCH 2/2] WIP tests

---
 tests/test_get.py    | 22 ++++++++++++++++++++++
 tests/test_update.py |  9 +++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 tests/test_get.py
 create mode 100644 tests/test_update.py

diff --git a/tests/test_get.py b/tests/test_get.py
new file mode 100644
index 0000000..0b2d714
--- /dev/null
+++ b/tests/test_get.py
@@ -0,0 +1,22 @@
+from sqlite_utils.db import Database, RowNotFound
+import pytest
+
+
+def test_get_single_pk(fresh_db):
+    cleo = {"id": 1, "name": "Cleo", "age": 4}
+    table = fresh_db["dogs"].insert(cleo, pk="id")
+    with pytest.raises(RowNotFound):
+        table.get(2)
+    with pytest.raises(RowNotFound):
+        table.get(None)
+    assert cleo == table.get(1)
+
+
+def test_get_compound_pk(fresh_db):
+    cleo = {"id1": 1, "id2": 1, "name": "Cleo", "age": 4}
+    table = fresh_db["dogs"].insert(cleo, pk=("id1", "id2"))
+    with pytest.raises(RowNotFound):
+        table.get(2)
+    with pytest.raises(RowNotFound):
+        table.get([2, 1])
+    assert cleo == table.get([1, 1])
diff --git a/tests/test_update.py b/tests/test_update.py
new file mode 100644
index 0000000..45018da
--- /dev/null
+++ b/tests/test_update.py
@@ -0,0 +1,9 @@
+from sqlite_utils.db import Database
+
+
+def test_update(fresh_db):
+    cleo = {"id": 1, "name": "Cleo", "age": 4}
+    fresh_db["dogs"].insert(cleo, pk="id")
+    assert [cleo] == list(fresh_db["dogs"].rows)
+    fresh_db["dogs"].update(1, {"age": 5})
+    assert [{"id": 1, "name": "Cleo", "age": 5}] == list(fresh_db["dogs"].rows)