From b9b199b725cf70b5c88b6bb8149351dc06846511 Mon Sep 17 00:00:00 2001
From: Michael Bayne <mdb@threerings.net>
Date: Tue, 26 Sep 2006 19:51:58 +0000
Subject: [PATCH] Update our Modifier's key on an insert (and store) so that it
 can (eventually) properly cache the results of our insert based on the newly
 assigned primary key. Also avoid freakout if an attempt is made to store() a
 non-primary-key using object.

---
 .../jdbc/depot/DepotMarshaller.java           | 10 +++++++---
 .../jdbc/depot/DepotRepository.java           | 20 ++++++++++++++-----
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/src/java/com/samskivert/jdbc/depot/DepotMarshaller.java b/src/java/com/samskivert/jdbc/depot/DepotMarshaller.java
index bad2b01..17ec9d6 100644
--- a/src/java/com/samskivert/jdbc/depot/DepotMarshaller.java
+++ b/src/java/com/samskivert/jdbc/depot/DepotMarshaller.java
@@ -340,24 +340,28 @@ public class DepotMarshaller<T>
     /**
      * Fills in the primary key just assigned to the supplied persistence
      * object by the execution of the results of {@link #createInsert}.
+     *
+     * @return the newly assigned primary key or null if the object does not
+     * use primary keys or this is not the right time to assign the key.
      */
-    public void assignPrimaryKey (
+    public DepotRepository.Key assignPrimaryKey (
             Connection conn, Object po, boolean postFactum)
         throws SQLException
     {
         // if we have no primary key or no generator, then we're done
         if (!hasPrimaryKey() || _keyGenerator == null) {
-            return;
+            return null;
         }
 
         // run this generator either before or after the actual insertion
         if (_keyGenerator.isPostFactum() != postFactum) {
-            return;
+            return null;
         }
 
         try {
             int nextValue = _keyGenerator.nextGeneratedValue(conn);
             _primaryKey.getField().set(po, nextValue);
+            return makePrimaryKey(nextValue);
         } catch (Exception e) {
             String errmsg = "Failed to assign primary key " +
                 "[type=" + _pclass + "]";
diff --git a/src/java/com/samskivert/jdbc/depot/DepotRepository.java b/src/java/com/samskivert/jdbc/depot/DepotRepository.java
index bfc2a3a..a54ccb1 100644
--- a/src/java/com/samskivert/jdbc/depot/DepotRepository.java
+++ b/src/java/com/samskivert/jdbc/depot/DepotRepository.java
@@ -142,11 +142,13 @@ public class DepotRepository
         final DepotMarshaller marsh = getMarshaller(record.getClass());
         return invoke(new Modifier(null) {
             public int invoke (Connection conn) throws SQLException {
-                marsh.assignPrimaryKey(conn, record, false);
+                // update our modifier's key so that it can cache our results
+                updateKey(marsh.assignPrimaryKey(conn, record, false));
                 PreparedStatement stmt = marsh.createInsert(conn, record);
                 try {
                     int mods = stmt.executeUpdate();
-                    marsh.assignPrimaryKey(conn, record, true);
+                    // check again in case we have a post-factum key generator
+                    updateKey(marsh.assignPrimaryKey(conn, record, true));
                     return mods;
                 } finally {
                     stmt.close();
@@ -341,7 +343,8 @@ public class DepotRepository
         throws PersistenceException
     {
         final DepotMarshaller marsh = getMarshaller(record.getClass());
-        return invoke(new Modifier(marsh.getPrimaryKey(record)) {
+        Key key = marsh.hasPrimaryKey() ? marsh.getPrimaryKey(record) : null;
+        return invoke(new Modifier(key) {
             public int invoke (Connection conn) throws SQLException {
                 PreparedStatement stmt = null;
                 try {
@@ -358,10 +361,10 @@ public class DepotRepository
 
                     // if the update modified zero rows or the primary key was
                     // obviously unset, do an insertion
-                    marsh.assignPrimaryKey(conn, record, false);
+                    updateKey(marsh.assignPrimaryKey(conn, record, false));
                     stmt = marsh.createInsert(conn, record);
                     int mods = stmt.executeUpdate();
-                    marsh.assignPrimaryKey(conn, record, true);
+                    updateKey(marsh.assignPrimaryKey(conn, record, true));
                     return mods;
 
                 } finally {
@@ -584,6 +587,13 @@ public class DepotRepository
             _key = key;
         }
 
+        protected void updateKey (Key key)
+        {
+            if (key != null) {
+                _key = key;
+            }
+        }
+
         protected Key _key;
     }