Add 'Drafted' interface and extend the InsertOperation to recognize that so as to be able to persist the mutated set of keys only yet return a pojo with values matching the draft, not just the mutations.

Changed {in,up}sert() to return the instance of the object.

pom.xml

@@ -5,7 +5,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>net.helenus</groupId>
     <artifactId>helenus-core</artifactId>
-    <version>2.0.14-SNAPSHOT</version>
+    <version>2.0.15-SNAPSHOT</version>
     <packaging>jar</packaging>
 
     <name>helenus</name>

src/main/java/net/helenus/core/HelenusSession.java

@@ -19,6 +19,7 @@ import java.io.Closeable;
 import java.io.PrintStream;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
@@ -31,6 +32,7 @@ import com.google.common.cache.CacheBuilder;
 import com.diffplug.common.base.Errors;
 
 import net.helenus.core.operation.*;
+import net.helenus.core.reflect.Drafted;
 import net.helenus.core.reflect.HelenusPropertyNode;
 import net.helenus.mapping.HelenusEntity;
 import net.helenus.mapping.MappingUtil;
@@ -342,30 +344,36 @@ public final class HelenusSession extends AbstractSessionOperations implements C
 		return new UpdateOperation(this, p, v);
 	}
 
-	public InsertOperation insert() {
+	public <T> InsertOperation<T> insert() { return new InsertOperation<T>(this, true); }
-		return new InsertOperation(this, true);
-	}
 
-	public InsertOperation insert(Object pojo) {
+    public <T> InsertOperation<T> insert(Object pojo) { return this.<T>insert(pojo, null); }
+
+    public <T> InsertOperation<T> insert(Drafted draft) { return this.<T>insert(draft.build(), draft.mutated()); }
+
+    public <T> InsertOperation<T> insert(Object pojo, Set<String> mutations) {
 		Objects.requireNonNull(pojo, "pojo is empty");
 
 		Class<?> iface = MappingUtil.getMappingInterface(pojo);
 		HelenusEntity entity = Helenus.entity(iface);
 
-		return new InsertOperation(this, entity, pojo, true);
+		return new InsertOperation<T>(this, entity, pojo, mutations, true);
 	}
 
-	public InsertOperation upsert() {
+	public <T> InsertOperation<T> upsert() {
-		return new InsertOperation(this, false);
+		return new InsertOperation<T>(this, false);
 	}
 
-	public InsertOperation upsert(Object pojo) {
+    public <T> InsertOperation<T> upsert(Drafted draft) { return this.<T>upsert(draft.build(), draft.mutated()); }
+
+    public <T> InsertOperation<T> upsert(Object pojo) { return this.<T>upsert(pojo, null); }
+
+	public <T> InsertOperation<T> upsert(Object pojo, Set<String> mutations) {
 		Objects.requireNonNull(pojo, "pojo is empty");
 
 		Class<?> iface = MappingUtil.getMappingInterface(pojo);
 		HelenusEntity entity = Helenus.entity(iface);
 
-		return new InsertOperation(this, entity, pojo, false);
+		return new InsertOperation<T>(this, entity, pojo, mutations, false);
 	}
 
 	public DeleteOperation delete() {

src/main/java/net/helenus/core/operation/InsertOperation.java

@@ -21,8 +21,8 @@ import com.datastax.driver.core.querybuilder.Insert;
 import com.datastax.driver.core.querybuilder.QueryBuilder;
 import net.helenus.core.AbstractSessionOperations;
 import net.helenus.core.Getter;
+import net.helenus.core.Helenus;
 import net.helenus.core.reflect.HelenusPropertyNode;
-import net.helenus.core.reflect.MapExportable;
 import net.helenus.mapping.HelenusEntity;
 import net.helenus.mapping.HelenusProperty;
 import net.helenus.mapping.MappingUtil;
@@ -32,12 +32,13 @@ import net.helenus.support.HelenusMappingException;
 
 import java.util.*;
 
-public final class InsertOperation extends AbstractOperation<ResultSet, InsertOperation> {
+public final class InsertOperation<T> extends AbstractOperation<T, InsertOperation<T>> {
 
     private HelenusEntity entity;
 
     private final List<Fun.Tuple2<HelenusPropertyNode, Object>> values = new ArrayList<Fun.Tuple2<HelenusPropertyNode, Object>>();
     private boolean ifNotExists;
+    private Object pojo;
 
     private int[] ttl;
     private long[] timestamp;
@@ -48,21 +49,15 @@ public final class InsertOperation extends AbstractOperation<ResultSet, InsertOp
         this.ifNotExists = ifNotExists;
     }
 
-    public InsertOperation(AbstractSessionOperations sessionOperations, HelenusEntity entity, Object pojo,
+    public InsertOperation(AbstractSessionOperations sessionOperations, HelenusEntity entity,
-            boolean ifNotExists) {
+            Object pojo, Set<String> mutations, boolean ifNotExists) {
         super(sessionOperations);
 
         this.entity = entity;
+        this.pojo = pojo;
         this.ifNotExists = ifNotExists;
         Collection<HelenusProperty> properties = entity.getOrderedProperties();
-        Set<String> keys = null;
+        Set<String> keys = (mutations == null) ?  null : mutations;
-
-        if (pojo instanceof MapExportable) {
-            keys = ((MapExportable) pojo).mutated();
-            if (keys == null) {
-                keys = ((MapExportable) pojo).toMap().keySet();
-            }
-        }
 
         for (HelenusProperty prop : properties) {
 
@@ -82,17 +77,17 @@ public final class InsertOperation extends AbstractOperation<ResultSet, InsertOp
 
     }
 
-    public InsertOperation ifNotExists() {
+    public InsertOperation<T> ifNotExists() {
         this.ifNotExists = true;
         return this;
     }
 
-    public InsertOperation ifNotExists(boolean enable) {
+    public InsertOperation<T> ifNotExists(boolean enable) {
         this.ifNotExists = enable;
         return this;
     }
 
-    public <V> InsertOperation value(Getter<V> getter, V val) {
+    public <V> InsertOperation<T> value(Getter<V> getter, V val) {
 
         Objects.requireNonNull(getter, "getter is empty");
 
@@ -141,17 +136,26 @@ public final class InsertOperation extends AbstractOperation<ResultSet, InsertOp
     }
 
     @Override
-    public ResultSet transform(ResultSet resultSet) {
+    public T transform(ResultSet resultSet) {
-        return resultSet;
+        if (pojo != null && ((T) pojo).getClass().isAssignableFrom(ResultSet.class)) {
+            return (T) pojo;
+        } else {
+            if (values.size() > 0) {
+                Map<String, Object> backingMap = new HashMap<String, Object>(values.size());
+                values.forEach(t -> backingMap.put(t._1.getProperty().getPropertyName(), t._2));
+                pojo = Helenus.map(values.get(0)._1.getEntity().getMappingInterface(), backingMap);
+            }
+        }
+        return (T) pojo;
     }
 
-    public InsertOperation usingTtl(int ttl) {
+    public InsertOperation<T> usingTtl(int ttl) {
         this.ttl = new int[1];
         this.ttl[0] = ttl;
         return this;
     }
 
-    public InsertOperation usingTimestamp(long timestamp) {
+    public InsertOperation<T> usingTimestamp(long timestamp) {
         this.timestamp = new long[1];
         this.timestamp[0] = timestamp;
         return this;

src/main/java/net/helenus/core/reflect/Drafted.java

@@ -0,0 +1,11 @@
+package net.helenus.core.reflect;
+
+import java.util.Set;
+
+public interface Drafted<T> extends MapExportable {
+
+    Set<String> mutated();
+
+    T build();
+
+}

src/main/java/net/helenus/core/reflect/MapExportable.java

@@ -16,7 +16,6 @@
 package net.helenus.core.reflect;
 
 import java.util.Map;
-import java.util.Set;
 
 public interface MapExportable {
 
@@ -24,6 +23,4 @@ public interface MapExportable {
 
 	Map<String, Object> toMap();
 
-	default Set<String> mutated() { return null; }
-
 }

src/main/java/net/helenus/core/reflect/MapperInvocationHandler.java

@@ -101,8 +101,6 @@ public class MapperInvocationHandler<E> implements InvocationHandler {
 
 				return type.getDefaultValue();
 
-			}  else if (returnType.isEnum()) {
-				throw new HelenusException("missing default type for enum user type " + returnType);
 			}
 
         }

src/test/java/net/helenus/test/integration/core/prepared/PreparedStatementTest.java

@@ -18,6 +18,7 @@ package net.helenus.test.integration.core.prepared;
 import static net.helenus.core.Query.marker;
 
 import java.math.BigDecimal;
+import java.util.Optional;
 
 import net.helenus.core.HelenusSession;
 import net.helenus.core.Query;
@@ -57,7 +58,7 @@ public class PreparedStatementTest extends AbstractEmbeddedCassandraTest {
 		session = Helenus.init(getSession()).showCql().add(Car.class).autoCreateDrop().get();
 		car = Helenus.dsl(Car.class, session.getMetadata());
 
-		insertOp = session.insert()
+		insertOp = session.<ResultSet>insert()
 		  .value(car::make, Query.marker())
 		  .value(car::model, Query.marker())
 		  .value(car::year, 2004)

src/test/java/net/helenus/test/integration/core/simple/InsertPartialTest.java

@@ -48,7 +48,7 @@ public class InsertPartialTest extends AbstractEmbeddedCassandraTest {
         Long id = rnd.nextLong();
         map.put("id", id);
         map.put("age", 5);
-        InsertOperation insert = session.insert(Helenus.map(User.class, map));
+        InsertOperation<User> insert = session.<User>insert(Helenus.map(User.class, map));
         String cql = "INSERT INTO simple_users (id,age) VALUES (" + id.toString() + ",5) IF NOT EXISTS;";
         Assert.assertEquals(cql, insert.cql());
         insert.sync();

src/test/java/net/helenus/test/integration/core/simple/SimpleUserTest.java

@@ -28,6 +28,8 @@ import net.helenus.core.Operator;
 import net.helenus.support.Fun;
 import net.helenus.test.integration.build.AbstractEmbeddedCassandraTest;
 
+import java.util.Optional;
+
 public class SimpleUserTest extends AbstractEmbeddedCassandraTest {
 
 	static User user;
@@ -164,6 +166,12 @@ public class SimpleUserTest extends AbstractEmbeddedCassandraTest {
 		Assert.assertEquals(Integer.valueOf(35), u.age());
 
 		// INSERT
+		User greg = session.<User>insert()
+				.value(user::name, "greg")
+				.value(user::age, 44)
+				.value(user::type, UserType.USER)
+                .value(user::id, 1234L)
+                .sync();
 
 		session.update()
 			.set(user::name, null)

src/test/java/net/helenus/test/integration/core/usertype/UserDefinedTypeTest.java

@@ -143,8 +143,7 @@ public class UserDefinedTypeTest extends AbstractEmbeddedCassandraTest {
 		Assert.assertNull(actual.country());
 		Assert.assertEquals(0, actual.zip());
 
-		// INSERT
+		// INSERT using UPDATE
-
 		session.update().set(account::address, null).where(account::id, Query.eq(123L)).sync();
 
 		Address adrNull = session.select(account::address).where(account::id, Query.eq(123L)).sync().findFirst().get()._1;
@@ -200,8 +199,7 @@ public class UserDefinedTypeTest extends AbstractEmbeddedCassandraTest {
 		Assert.assertEquals(addressNoMapping.getString("city"), found.getString("city"));
 
 
-		// INSERT
+		// INSERT using UPDATE
-
 		session.update().set(account::addressNoMapping, null).where(account::id, Query.eq(777L)).sync();
 
 		found = session.select(account::addressNoMapping).where(account::id, Query.eq(777L)).sync().findFirst().get()._1;

src/test/java/net/helenus/test/unit/core/dsl/Account.java

@@ -16,7 +16,10 @@
 package net.helenus.test.unit.core.dsl;
 
 import java.util.Date;
+import java.util.Map;
+import java.util.Set;
 
+import net.helenus.core.reflect.Drafted;
 import net.helenus.mapping.annotation.*;
 
 @Table
@@ -35,6 +38,21 @@ public interface Account {
 
 	@Transient
 	default Draft draft() { return new Draft(); }
-	class Draft {}
+	class Draft implements Drafted { //TODO
+		@Override
+		public Set<String> mutated() {
+			return null;
+		}
+
+		@Override
+		public Object build() {
+			return null;
+		}
+
+		@Override
+		public Map<String, Object> toMap() {
+			return null;
+		}
+	}
 
 }