JAVA-324: Supporing object class (but not internal classes) for findAndModify

Author: jyemin

src/main/com/mongodb/DBCollection.java

@@ -378,12 +378,43 @@ public DBObject findAndModify(DBObject query, DBObject fields, DBObject sort, bo
             throw new MongoException("FindAndModify: Remove cannot be mixed with the Update, or returnNew params!");
 
         CommandResult res = this._db.command( cmd );
-        if (res.ok() || res.getErrorMessage().equals( "No matching object found" ))
-            return (DBObject) res.get( "value" );
+        if (res.ok() || res.getErrorMessage().equals( "No matching object found" )) {
+            return replaceWithObjectClass((DBObject) res.get( "value" ));
+        }
         res.throwOnError();
         return null;
     }
 
+    /**
+     * Doesn't yet handle internal classes properly, so this method only does something if object class is set but
+     * no internal classes are set.
+     *
+     * @param oldObj  the original value from the command result
+     * @return replaced object if necessary, or oldObj
+     */
+    private DBObject replaceWithObjectClass(DBObject oldObj) {
+        if (oldObj == null || getObjectClass() == null &  _internalClass.isEmpty()) {
+            return oldObj;
+        }
+
+        DBObject newObj = instantiateObjectClassInstance();
+
+        for (String key : oldObj.keySet()) {
+            newObj.put(key, oldObj.get(key));
+        }
+        return newObj;
+    }
+
+    private DBObject instantiateObjectClassInstance() {
+        try {
+            return (DBObject) getObjectClass().newInstance();
+        } catch (InstantiationException e) {
+            throw new MongoInternalException("can't create instance of type " + getObjectClass(), e);
+        } catch (IllegalAccessException e) {
+            throw new MongoInternalException("can't create instance of type " + getObjectClass(), e);
+        }
+    }
+
 
     /**
      * calls {@link DBCollection#findAndModify(com.mongodb.DBObject, com.mongodb.DBObject, com.mongodb.DBObject, boolean, com.mongodb.DBObject, boolean, boolean)}

src/test/com/mongodb/ReflectionTest.java

@@ -27,11 +27,6 @@ public class ReflectionTest extends TestCase {
     public static class Person extends ReflectionDBObject {
 
         public Person(){
-
-        }
-
-        Person( String name ){
-            _name = name;
         }
 
         public String getName(){
@@ -55,16 +50,17 @@ public ReflectionTest()
     @Test
     public void test1()
         throws MongoException {
-        DBCollection c = _db.getCollection( "persen.test1" );
+        DBCollection c = _db.getCollection( "person.test1" );
         c.drop();
         c.setObjectClass( Person.class );
 
-        Person p = new Person( "eliot" );
+        Person p = new Person();
+        p.setName( "eliot" );
         c.save( p );
 
         DBObject out = c.findOne();
         assertEquals( "eliot" , out.get( "Name" ) );
-        assertTrue( out instanceof Person , "didn't come out as Person" );
+        assertEquals(Person.class, out.getClass());
     }
 
     public static class Outer extends ReflectionDBObject {
@@ -105,12 +101,54 @@ public void test2()
 
         DBObject out = c.findOne();
         assertEquals( "eliot" , out.get( "Name" ) );
-        assertTrue( out instanceof Outer , "didn't come out as Person" );        
+        assertEquals(Outer.class, out.getClass());
         o = (Outer)out;
         assertEquals( "eliot" , o.getName() );
         assertEquals( 17 , o.getInner().getNumber() );
     }
 
+    static class Process extends ReflectionDBObject {
+
+        public Process() {}
+
+        public String getName() {
+            return name;
+        }
+
+        public void setName(String name) {
+            this.name = name;
+        }
+
+        public int getStatus() {
+            return status;
+        }
+
+        public void setStatus(int status) {
+            this.status = status;
+        }
+
+        String name;
+        int status;
+    }
+
+    @Test
+    public void testFindAndModify() {
+        DBCollection c = _db.getCollection( "findAndModify" );
+        c.drop();
+        c.setObjectClass( Process.class );
+
+        Process p = new Process();
+        p.setName("test");
+        p.setStatus(0);
+        c.save(p, WriteConcern.SAFE);
+
+        DBObject obj = c.findAndModify(new BasicDBObject(), new BasicDBObject("$set", new BasicDBObject("status", 1)));
+        assertEquals(Process.class, obj.getClass());
+        Process pModified = (Process) obj;
+        assertEquals(0, pModified.getStatus());
+        assertEquals("test", pModified.getName());
+    }
+
     final DB _db;
 
     public static void main( String args[] )