From 82b9d8faf092caea6aff8bc164db09fd325df808 Mon Sep 17 00:00:00 2001
From: Jeremy Howard <github@jhoward.fastmail.fm>
Date: Mon, 3 Jun 2024 18:45:42 +1000
Subject: [PATCH] fixes #5

---
 fastlite/kw.py  | 15 ++++++++++++---
 nbs/index.ipynb | 10 +++++++++-
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/fastlite/kw.py b/fastlite/kw.py
index b4c96e7..1e5c511 100644
--- a/fastlite/kw.py
+++ b/fastlite/kw.py
@@ -12,6 +12,15 @@ def xtra(self:Table, **kwargs):
     "Set `xtra_id`"
     self.xtra_id = kwargs
 
+@patch
+def get_last(self:Table, as_cls:bool=True):
+    row = first(self.rows_where('_rowid_=?', (self.last_rowid,)))
+    vals = [row[pk] for pk in self.pks]
+    self.last_pk = vals[0] if len(vals)==1 else vals
+    if as_cls and hasattr(self,'cls'): row = self.cls(**row)
+    return row
+
+
 @patch
 def get(self:Table, pk_values: list|tuple|str|int, as_cls:bool=True)->Any:
     if not isinstance(pk_values, (list, tuple)): pk_values = [pk_values]
@@ -90,7 +99,7 @@ def update(self:Table, updates: dict|None=None, pk_values: list|tuple|str|int|fl
     updates = {**updates, **kwargs, **xtra}
     if not pk_values: pk_values = [updates[o] for o in self.pks]
     self._orig_update(pk_values, updates=updates, alter=alter, conversions=conversions)
-    return self.get(self.last_pk)
+    return self.get_last()
 
 
 @patch
@@ -144,7 +153,7 @@ def insert(
         record=record, pk=pk, foreign_keys=foreign_keys, column_order=column_order, not_null=not_null,
         defaults=defaults, hash_id=hash_id, hash_id_columns=hash_id_columns, alter=alter, ignore=ignore,
         replace=replace, extracts=extracts, conversions=conversions, columns=columns, strict=strict)
-    return self.get(self.last_pk)
+    return self.get_last()
 
 
 @patch
@@ -172,7 +181,7 @@ def upsert(
         record=record, pk=pk, foreign_keys=foreign_keys, column_order=column_order, not_null=not_null,
         defaults=defaults, hash_id=hash_id, hash_id_columns=hash_id_columns, alter=alter,
         extracts=extracts, conversions=conversions, columns=columns, strict=strict)
-    return self.get(self.last_pk)
+    return self.get_last()
 
 
 @patch
diff --git a/nbs/index.ipynb b/nbs/index.ipynb
index d0304ea..9309ed7 100644
--- a/nbs/index.ipynb
+++ b/nbs/index.ipynb
@@ -47,7 +47,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sqlite_utils import Database\n",
+    "from sqlite_utils import Database"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "from fastlite import *\n",
     "from fastcore.utils import *\n",
     "from fastcore.net import urlsave"