#595 - Refactor - persistence context root type handling

src/main/java/com/avaje/ebean/bean/PersistenceContext.java

@@ -12,7 +12,7 @@ public interface PersistenceContext {
   /**
    * Put the entity bean into the PersistenceContext.
    */
-  void put(Object id, Object bean);
+  void put(Class<?> rootType, Object id, Object bean);
 
   /**
    * Put the entity bean into the PersistenceContext if one is not already
@@ -22,7 +22,7 @@ public interface PersistenceContext {
    * returns null.
    * </p>
    */
-  Object putIfAbsent(Object id, Object bean);
+  Object putIfAbsent(Class<?> rootType, Object id, Object bean);
 
   /**
    * Return an object given its type and unique id.

src/main/java/com/avaje/ebean/bean/PersistenceContextUtil.java

@@ -0,0 +1,21 @@
+package com.avaje.ebean.bean;
+
+import javax.persistence.Entity;
+
+/**
+ * Utility to find the root bean type.
+ */
+public class PersistenceContextUtil {
+
+  /**
+   * Find and return the root bean type for the given class.
+   */
+  public static Class<?> root(Class<?> beanType) {
+    Class<?> parent = beanType.getSuperclass();
+    while (parent != null && parent.isAnnotationPresent(Entity.class)) {
+      beanType = parent;
+      parent = parent.getSuperclass();
+    }
+    return beanType;
+  }
+}

src/main/java/com/avaje/ebeaninternal/server/core/DefaultBeanLoader.java

@@ -126,7 +126,7 @@ public class DefaultBeanLoader {
 
     if (pc == null) {
       pc = new DefaultPersistenceContext();
-      pc.put(parentId, parentBean);
+      parentDesc.contextPut(pc, parentId, parentBean);
     }
 
     boolean useManyIdCache = beanCollection != null && parentDesc.isManyPropCaching();
@@ -263,7 +263,7 @@ public class DefaultBeanLoader {
     if (pc == null) {
       // a reference with no existing persistenceContext
       pc = new DefaultPersistenceContext();
-      pc.put(id, bean);
+      desc.contextPut(pc, id, bean);
       ebi.setPersistenceContext(pc);
     }
 

src/main/java/com/avaje/ebeaninternal/server/core/DefaultServer.java

@@ -616,9 +616,8 @@ public final class DefaultServer implements SpiServer, SpiEbeanServer {
         ref = desc.createReference(null, id);
       }
 
-      if (ctx != null && (ref instanceof EntityBean)) {
-        // Not putting a vanilla reference in the persistence context
-        ctx.put(id, ref);
+      if (ctx != null) {
+        desc.contextPut(ctx, id, ref);
       }
     }
     return (T) ref;

src/main/java/com/avaje/ebeaninternal/server/core/OrmQueryRequest.java

@@ -189,7 +189,7 @@ public final class OrmQueryRequest<T> extends BeanRequest implements BeanQueryRe
    */
   public void persistenceContextAdd(EntityBean bean) {
     Object id = beanDescriptor.getId(bean);
-    persistenceContext.put(id, bean);
+    beanDescriptor.contextPut(persistenceContext, id, bean);
   }
 
   /**

src/main/java/com/avaje/ebeaninternal/server/deploy/AssocOneHelp.java

@@ -0,0 +1,39 @@
+package com.avaje.ebeaninternal.server.deploy;
+
+import com.avaje.ebean.bean.EntityBean;
+import com.avaje.ebeaninternal.server.query.SqlJoinType;
+
+import java.sql.SQLException;
+
+/**
+ * Local interface to handle Embedded, Reference and Reference Exported
+ * cases.
+ */
+abstract class AssocOneHelp {
+
+  /**
+   * Effectively skip reading (the jdbc resultSet as already in the persistence context etc).
+   */
+  abstract void loadIgnore(DbReadContext ctx);
+
+  /**
+   * Read and return the bean.
+   */
+  abstract Object read(DbReadContext ctx) throws SQLException;
+
+  /**
+   * Read setting values into the bean.
+   */
+  abstract Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException;
+
+  /**
+   * Append to the select clause.
+   */
+  abstract void appendSelect(DbSqlContext ctx, boolean subQuery);
+
+  /**
+   * Append to the from clause.
+   */
+  abstract void appendFrom(DbSqlContext ctx, SqlJoinType joinType);
+
+}

src/main/java/com/avaje/ebeaninternal/server/deploy/AssocOneHelpEmbedded.java

@@ -0,0 +1,69 @@
+package com.avaje.ebeaninternal.server.deploy;
+
+import com.avaje.ebean.bean.EntityBean;
+import com.avaje.ebeaninternal.server.query.SqlJoinType;
+
+import java.sql.SQLException;
+
+/**
+ * Helper for Embedded BeanPropertyAssocOne.
+ */
+final class AssocOneHelpEmbedded extends AssocOneHelp {
+
+  private BeanPropertyAssocOne beanPropertyAssocOne;
+
+  public AssocOneHelpEmbedded(BeanPropertyAssocOne beanPropertyAssocOne) {
+    this.beanPropertyAssocOne = beanPropertyAssocOne;
+  }
+
+  void loadIgnore(DbReadContext ctx) {
+    for (int i = 0; i < beanPropertyAssocOne.embeddedProps.length; i++) {
+      beanPropertyAssocOne.embeddedProps[i].loadIgnore(ctx);
+    }
+  }
+
+  @Override
+  Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException {
+    Object dbVal = read(ctx);
+    if (bean != null) {
+      // set back to the parent bean
+      beanPropertyAssocOne.setValue(bean, dbVal);
+      ctx.propagateState(dbVal);
+      return dbVal;
+
+    } else {
+      return null;
+    }
+  }
+
+  @Override
+  Object read(DbReadContext ctx) throws SQLException {
+
+    EntityBean embeddedBean = beanPropertyAssocOne.targetDescriptor.createEntityBean();
+
+    boolean notNull = false;
+    for (int i = 0; i < beanPropertyAssocOne.embeddedProps.length; i++) {
+      Object value = beanPropertyAssocOne.embeddedProps[i].readSet(ctx, embeddedBean);
+      if (value != null) {
+        notNull = true;
+      }
+    }
+    if (notNull) {
+      ctx.propagateState(embeddedBean);
+      return embeddedBean;
+    } else {
+      return null;
+    }
+  }
+
+  @Override
+  void appendFrom(DbSqlContext ctx, SqlJoinType joinType) {
+  }
+
+  @Override
+  void appendSelect(DbSqlContext ctx, boolean subQuery) {
+    for (int i = 0; i < beanPropertyAssocOne.embeddedProps.length; i++) {
+      beanPropertyAssocOne.embeddedProps[i].appendSelect(ctx, subQuery);
+    }
+  }
+}

src/main/java/com/avaje/ebeaninternal/server/deploy/AssocOneHelpReference.java

@@ -0,0 +1,120 @@
+package com.avaje.ebeaninternal.server.deploy;
+
+import com.avaje.ebean.bean.EntityBean;
+import com.avaje.ebean.bean.EntityBeanIntercept;
+import com.avaje.ebean.bean.PersistenceContextUtil;
+import com.avaje.ebeaninternal.server.query.SqlJoinType;
+
+import java.sql.SQLException;
+
+/**
+ * Helper for BeanPropertyAssocOne imported reference - this is the common case.
+ */
+final class AssocOneHelpReference extends AssocOneHelp {
+
+  private BeanPropertyAssocOne beanPropertyAssocOne;
+
+  AssocOneHelpReference(BeanPropertyAssocOne beanPropertyAssocOne) {
+    this.beanPropertyAssocOne = beanPropertyAssocOne;
+  }
+
+  @Override
+  void loadIgnore(DbReadContext ctx) {
+    beanPropertyAssocOne.targetIdBinder.loadIgnore(ctx);
+    if (beanPropertyAssocOne.targetInheritInfo != null) {
+      ctx.getDataReader().incrementPos(1);
+    }
+  }
+
+  @Override
+  Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException {
+    Object val = read(ctx);
+    if (bean != null) {
+      beanPropertyAssocOne.setValue(bean, val);
+      ctx.propagateState(val);
+    }
+    return val;
+  }
+
+  /**
+   * Read and set a Reference bean.
+   */
+  @Override
+  Object read(DbReadContext ctx) throws SQLException {
+
+    BeanDescriptor<?> rowDescriptor = null;
+    Class<?> rowType = beanPropertyAssocOne.targetType;
+    if (beanPropertyAssocOne.targetInheritInfo != null) {
+      // read discriminator to determine the type
+      InheritInfo rowInheritInfo = beanPropertyAssocOne.targetInheritInfo.readType(ctx);
+      if (rowInheritInfo != null) {
+        rowType = rowInheritInfo.getType();
+        rowDescriptor = rowInheritInfo.getBeanDescriptor();
+      }
+    }
+
+    // read the foreign key column(s)
+    Object id = beanPropertyAssocOne.targetIdBinder.read(ctx);
+    if (id == null) {
+      return null;
+    }
+
+    // check transaction context to see if it already exists
+    Object existing = ctx.getPersistenceContext().get(rowType, id);
+
+    if (existing != null) {
+      return existing;
+    }
+
+    Boolean readOnly = ctx.isReadOnly();
+    Object ref;
+    if (beanPropertyAssocOne.targetInheritInfo != null) {
+      // for inheritance hierarchy create the correct type for this row...
+      ref = rowDescriptor.createReference(readOnly, id);
+    } else {
+      ref = beanPropertyAssocOne.targetDescriptor.createReference(readOnly, id);
+    }
+
+    Class<?> rootType = PersistenceContextUtil.root(ref.getClass());
+    Object existingBean = ctx.getPersistenceContext().putIfAbsent(rootType, id, ref);
+    if (existingBean != null) {
+      // advanced case when we use multiple concurrent threads to
+      // build a single object graph, and another thread has since
+      // loaded a matching bean so we will use that instead.
+      ref = existingBean;
+
+    } else {
+      EntityBeanIntercept ebi = ((EntityBean) ref)._ebean_getIntercept();
+      if (Boolean.TRUE.equals(ctx.isReadOnly())) {
+        ebi.setReadOnly(true);
+      }
+      ctx.register(beanPropertyAssocOne.name, ebi);
+    }
+
+    return ref;
+  }
+
+  @Override
+  void appendFrom(DbSqlContext ctx, SqlJoinType joinType) {
+    if (beanPropertyAssocOne.targetInheritInfo != null) {
+      // add join to support the discriminator column
+      String relativePrefix = ctx.getRelativePrefix(beanPropertyAssocOne.name);
+      beanPropertyAssocOne.tableJoin.addJoin(joinType, relativePrefix, ctx);
+    }
+  }
+
+  /**
+   * Append columns for foreign key columns.
+   */
+  @Override
+  void appendSelect(DbSqlContext ctx, boolean subQuery) {
+
+    if (!subQuery && beanPropertyAssocOne.targetInheritInfo != null) {
+      // add discriminator column
+      String relativePrefix = ctx.getRelativePrefix(beanPropertyAssocOne.getName());
+      String tableAlias = ctx.getTableAlias(relativePrefix);
+      ctx.appendColumn(tableAlias, beanPropertyAssocOne.targetInheritInfo.getDiscriminatorColumn());
+    }
+    beanPropertyAssocOne.importedId.sqlAppend(ctx);
+  }
+}

src/main/java/com/avaje/ebeaninternal/server/deploy/AssocOneHelpReferenceExported.java

@@ -0,0 +1,91 @@
+package com.avaje.ebeaninternal.server.deploy;
+
+import com.avaje.ebean.bean.EntityBean;
+import com.avaje.ebean.bean.EntityBeanIntercept;
+import com.avaje.ebean.bean.PersistenceContext;
+import com.avaje.ebeaninternal.server.deploy.id.IdBinder;
+import com.avaje.ebeaninternal.server.query.SqlJoinType;
+
+import java.sql.SQLException;
+
+/**
+ * Helper for BeanPropertyAssocOne for OneToOne exported reference - not so common.
+ */
+final class AssocOneHelpReferenceExported extends AssocOneHelp {
+
+  private BeanPropertyAssocOne beanPropertyAssocOne;
+
+  public AssocOneHelpReferenceExported(BeanPropertyAssocOne beanPropertyAssocOne) {
+    this.beanPropertyAssocOne = beanPropertyAssocOne;
+  }
+
+  @Override
+  void loadIgnore(DbReadContext ctx) {
+    beanPropertyAssocOne.targetDescriptor.getIdBinder().loadIgnore(ctx);
+  }
+
+  /**
+   * Read and set a Reference bean.
+   */
+  @Override
+  Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException {
+
+    Object dbVal = read(ctx);
+    if (bean != null) {
+      beanPropertyAssocOne.setValue(bean, dbVal);
+      ctx.propagateState(dbVal);
+    }
+    return dbVal;
+  }
+
+  @Override
+  Object read(DbReadContext ctx) throws SQLException {
+
+    // TODO: Support for Inheritance hierarchy on exported OneToOne ?
+    IdBinder idBinder = beanPropertyAssocOne.targetDescriptor.getIdBinder();
+    Object id = idBinder.read(ctx);
+    if (id == null) {
+      return null;
+    }
+
+    PersistenceContext persistCtx = ctx.getPersistenceContext();
+    Object existing = persistCtx.get(beanPropertyAssocOne.targetType, id);
+
+    if (existing != null) {
+      return existing;
+    }
+    Object ref = beanPropertyAssocOne.targetDescriptor.createReference(ctx.isReadOnly(), id);
+
+    EntityBeanIntercept ebi = ((EntityBean) ref)._ebean_getIntercept();
+    if (Boolean.TRUE.equals(ctx.isReadOnly())) {
+      ebi.setReadOnly(true);
+    }
+    beanPropertyAssocOne.targetDescriptor.contextPut(persistCtx, id, ref);
+    ctx.register(beanPropertyAssocOne.name, ebi);
+    return ref;
+  }
+
+  /**
+   * Append columns for foreign key columns.
+   */
+  @Override
+  void appendSelect(DbSqlContext ctx, boolean subQuery) {
+
+    // set appropriate tableAlias for the exported id columns
+
+    String relativePrefix = ctx.getRelativePrefix(beanPropertyAssocOne.getName());
+    ctx.pushTableAlias(relativePrefix);
+
+    IdBinder idBinder = beanPropertyAssocOne.targetDescriptor.getIdBinder();
+    idBinder.appendSelect(ctx, subQuery);
+
+    ctx.popTableAlias();
+  }
+
+  @Override
+  void appendFrom(DbSqlContext ctx, SqlJoinType joinType) {
+
+    String relativePrefix = ctx.getRelativePrefix(beanPropertyAssocOne.getName());
+    beanPropertyAssocOne.tableJoin.addJoin(joinType, relativePrefix, ctx);
+  }
+}

src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptor.java

@@ -12,6 +12,7 @@ import com.avaje.ebean.bean.BeanCollection;
 import com.avaje.ebean.bean.EntityBean;
 import com.avaje.ebean.bean.EntityBeanIntercept;
 import com.avaje.ebean.bean.PersistenceContext;
+import com.avaje.ebean.bean.PersistenceContextUtil;
 import com.avaje.ebean.config.EncryptKey;
 import com.avaje.ebean.config.ServerConfig;
 import com.avaje.ebean.config.dbplatform.IdGenerator;
@@ -33,7 +34,6 @@ import com.avaje.ebean.plugin.BeanDocType;
 import com.avaje.ebean.plugin.BeanType;
 import com.avaje.ebean.plugin.ExpressionPath;
 import com.avaje.ebean.plugin.Property;
-import com.avaje.ebean.text.json.JsonReadOptions;
 import com.avaje.ebeaninternal.api.CQueryPlanKey;
 import com.avaje.ebeaninternal.api.LoadContext;
 import com.avaje.ebeaninternal.api.SpiEbeanServer;
@@ -74,7 +74,6 @@ import com.avaje.ebeanservice.docstore.api.mapping.DocMappingBuilder;
 import com.avaje.ebeanservice.docstore.api.mapping.DocPropertyMapping;
 import com.avaje.ebeanservice.docstore.api.mapping.DocPropertyType;
 import com.avaje.ebeanservice.docstore.api.mapping.DocumentMapping;
-import com.fasterxml.jackson.core.JsonParser;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -195,6 +194,7 @@ public class BeanDescriptor<T> implements MetaBeanInfo, BeanType<T> {
    * The type of bean this describes.
    */
   private final Class<T> beanType;
+  private final Class<?> rootBeanType;
 
   /**
    * This is not sent to a remote client.
@@ -387,6 +387,7 @@ public class BeanDescriptor<T> implements MetaBeanInfo, BeanType<T> {
     this.fullName = InternString.intern(deploy.getFullName());
 
     this.beanType = deploy.getBeanType();
+    this.rootBeanType = PersistenceContextUtil.root(beanType);
     this.prototypeEntityBean = createPrototypeEntityBean(beanType);
     
     this.namedQueries = deploy.getNamedQueries();
@@ -1667,6 +1668,20 @@ public class BeanDescriptor<T> implements MetaBeanInfo, BeanType<T> {
     return fullName;
   }
 
+  /**
+   * Put the bean into the persistence context.
+   */
+  public void contextPut(PersistenceContext pc, Object id, Object bean) {
+    pc.put(rootBeanType, id, bean);
+  }
+
+  /**
+   * Put the bean into the persistence context if it is absent.
+   */
+  public Object contextPutIfAbsent(PersistenceContext pc, Object id, EntityBean localBean) {
+    return pc.putIfAbsent(rootBeanType, id, localBean);
+  }
+
   /**
    * Helper method to return the unique property. If only one property makes up
    * the unique id then it's value is returned. If there is a concatenated

src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptorCacheHelp.java

@@ -245,7 +245,7 @@ public final class BeanDescriptorCacheHelp<T> {
       EntityBeanIntercept refEbi = ((EntityBean) refBean)._ebean_getIntercept();
 
       many.add(bc, (EntityBean) refBean);
-      persistenceContext.put(id, refBean);
+      targetDescriptor.contextPut(persistenceContext, id, refBean);
       refEbi.setPersistenceContext(persistenceContext);
     }
     return true;
@@ -343,8 +343,7 @@ public final class BeanDescriptorCacheHelp<T> {
     EntityBeanIntercept ebi = entityBean._ebean_getIntercept();
     ebi.setPersistenceContext(context);
     Object id = desc.getId(entityBean);
-    context.put(id, bean);
-
+    desc.contextPut(context, id, bean);
   }
   
   /**

src/main/java/com/avaje/ebeaninternal/server/deploy/BeanPropertyAssocOne.java

@@ -6,11 +6,8 @@ import com.avaje.ebean.SqlUpdate;
 import com.avaje.ebean.Transaction;
 import com.avaje.ebean.ValuePair;
 import com.avaje.ebean.bean.EntityBean;
-import com.avaje.ebean.bean.EntityBeanIntercept;
-import com.avaje.ebean.bean.PersistenceContext;
 import com.avaje.ebeaninternal.server.cache.CachedBeanData;
 import com.avaje.ebeaninternal.server.core.DefaultSqlUpdate;
-import com.avaje.ebeaninternal.server.deploy.id.IdBinder;
 import com.avaje.ebeaninternal.server.deploy.id.ImportedId;
 import com.avaje.ebeaninternal.server.deploy.meta.DeployBeanPropertyAssocOne;
 import com.avaje.ebeaninternal.server.el.ElPropertyChainBuilder;
@@ -34,28 +31,29 @@ import java.util.Map;
  */
 public class BeanPropertyAssocOne<T> extends BeanPropertyAssoc<T> {
 
-  private final boolean oneToOne;
+  protected final boolean oneToOne;
 
-  private final boolean oneToOneExported;
+  protected final boolean oneToOneExported;
 
-  private final boolean importedPrimaryKey;
+  protected final boolean importedPrimaryKey;
 
-  private final LocalHelp localHelp;
+  protected final AssocOneHelp localHelp;
 
-  private final BeanProperty[] embeddedProps;
+  protected final BeanProperty[] embeddedProps;
 
-  private final HashMap<String, BeanProperty> embeddedPropsMap;
+  protected final HashMap<String, BeanProperty> embeddedPropsMap;
 
   /**
    * The information for Imported foreign Keys.
    */
-  private ImportedId importedId;
+  protected ImportedId importedId;
 
-  private ExportedProperty[] exportedProperties;
+  protected ExportedProperty[] exportedProperties;
 
-  private String deleteByParentIdSql;
-  private String deleteByParentIdInSql;
-  BeanPropertyAssocMany<?> relationshipProperty;
+  protected String deleteByParentIdSql;
+  protected String deleteByParentIdInSql;
+
+  protected BeanPropertyAssocMany<?> relationshipProperty;
 
   /**
    * Create based on deploy information of an EmbeddedId.
@@ -629,269 +627,13 @@ public class BeanPropertyAssocOne<T> extends BeanPropertyAssoc<T> {
     }
   }
 
-  private LocalHelp createHelp(boolean embedded, boolean oneToOneExported) {
+  private AssocOneHelp createHelp(boolean embedded, boolean oneToOneExported) {
     if (embedded) {
-      return new Embedded();
+      return new AssocOneHelpEmbedded(this);
     } else if (oneToOneExported) {
-      return new ReferenceExported();
+      return new AssocOneHelpReferenceExported(this);
     } else {
-      return new Reference();
-    }
-  }
-
-  /**
-   * Local interface to handle Embedded, Reference and Reference Exported
-   * cases.
-   */
-  private abstract class LocalHelp {
-
-    abstract void loadIgnore(DbReadContext ctx);
-
-    abstract Object read(DbReadContext ctx) throws SQLException;
-
-    abstract Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException;
-
-    abstract void appendSelect(DbSqlContext ctx, boolean subQuery);
-
-    abstract void appendFrom(DbSqlContext ctx, SqlJoinType joinType);
-
-  }
-
-  private final class Embedded extends LocalHelp {
-
-    void loadIgnore(DbReadContext ctx) {
-      for (int i = 0; i < embeddedProps.length; i++) {
-        embeddedProps[i].loadIgnore(ctx);
-      }
-    }
-
-    @Override
-    Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException {
-      Object dbVal = read(ctx);
-      if (bean != null) {
-        // set back to the parent bean
-        setValue(bean, dbVal);
-        ctx.propagateState(dbVal);
-        return dbVal;
-
-      } else {
-        return null;
-      }
-    }
-
-    @Override
-    Object read(DbReadContext ctx) throws SQLException {
-
-      EntityBean embeddedBean = targetDescriptor.createEntityBean();
-
-      boolean notNull = false;
-      for (int i = 0; i < embeddedProps.length; i++) {
-        Object value = embeddedProps[i].readSet(ctx, embeddedBean);
-        if (value != null) {
-          notNull = true;
-        }
-      }
-      if (notNull) {
-        ctx.propagateState(embeddedBean);
-        return embeddedBean;
-      } else {
-        return null;
-      }
-    }
-
-    @Override
-    void appendFrom(DbSqlContext ctx, SqlJoinType joinType) {
-    }
-
-    @Override
-    void appendSelect(DbSqlContext ctx, boolean subQuery) {
-      for (int i = 0; i < embeddedProps.length; i++) {
-        embeddedProps[i].appendSelect(ctx, subQuery);
-      }
-    }
-  }
-
-  /**
-   * For imported reference - this is the common case.
-   */
-  private final class Reference extends LocalHelp {
-
-    Reference() {
-    }
-
-    @Override
-    void loadIgnore(DbReadContext ctx) {
-      targetIdBinder.loadIgnore(ctx);
-      if (targetInheritInfo != null) {
-        ctx.getDataReader().incrementPos(1);
-      }
-    }
-
-    @Override
-    Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException {
-      Object val = read(ctx);
-      if (bean != null) {
-        setValue(bean, val);
-        ctx.propagateState(val);
-      }
-      return val;
-    }
-
-    /**
-     * Read and set a Reference bean.
-     */
-    @Override
-    Object read(DbReadContext ctx) throws SQLException {
-
-      BeanDescriptor<?> rowDescriptor = null;
-      Class<?> rowType = targetType;
-      if (targetInheritInfo != null) {
-        // read discriminator to determine the type
-        InheritInfo rowInheritInfo = targetInheritInfo.readType(ctx);
-        if (rowInheritInfo != null) {
-          rowType = rowInheritInfo.getType();
-          rowDescriptor = rowInheritInfo.getBeanDescriptor();
-        }
-      }
-
-      // read the foreign key column(s)
-      Object id = targetIdBinder.read(ctx);
-      if (id == null) {
-        return null;
-      }
-
-      // check transaction context to see if it already exists
-      Object existing = ctx.getPersistenceContext().get(rowType, id);
-
-      if (existing != null) {
-        return existing;
-      }
-
-      Boolean readOnly = ctx.isReadOnly();
-      Object ref;
-      if (targetInheritInfo != null) {
-        // for inheritance hierarchy create the correct type for this row...
-        ref = rowDescriptor.createReference(readOnly, id);
-      } else {
-        ref = targetDescriptor.createReference(readOnly, id);
-      }
-
-      Object existingBean = ctx.getPersistenceContext().putIfAbsent(id, ref);
-      if (existingBean != null) {
-        // advanced case when we use multiple concurrent threads to
-        // build a single object graph, and another thread has since
-        // loaded a matching bean so we will use that instead.
-        ref = existingBean;
-
-      } else {
-        EntityBeanIntercept ebi = ((EntityBean) ref)._ebean_getIntercept();
-        if (Boolean.TRUE.equals(ctx.isReadOnly())) {
-          ebi.setReadOnly(true);
-        }
-        ctx.register(name, ebi);
-      }
-
-      return ref;
-    }
-
-    @Override
-    void appendFrom(DbSqlContext ctx, SqlJoinType joinType) {
-      if (targetInheritInfo != null) {
-        // add join to support the discriminator column
-        String relativePrefix = ctx.getRelativePrefix(name);
-        tableJoin.addJoin(joinType, relativePrefix, ctx);
-      }
-    }
-
-    /**
-     * Append columns for foreign key columns.
-     */
-    @Override
-    void appendSelect(DbSqlContext ctx, boolean subQuery) {
-
-      if (!subQuery && targetInheritInfo != null) {
-        // add discriminator column
-        String relativePrefix = ctx.getRelativePrefix(getName());
-        String tableAlias = ctx.getTableAlias(relativePrefix);
-        ctx.appendColumn(tableAlias, targetInheritInfo.getDiscriminatorColumn());
-      }
-      importedId.sqlAppend(ctx);
-    }
-  }
-
-  /**
-   * For OneToOne exported reference - not so common.
-   */
-  private final class ReferenceExported extends LocalHelp {
-
-    @Override
-    void loadIgnore(DbReadContext ctx) {
-      targetDescriptor.getIdBinder().loadIgnore(ctx);
-    }
-
-    /**
-     * Read and set a Reference bean.
-     */
-    @Override
-    Object readSet(DbReadContext ctx, EntityBean bean) throws SQLException {
-
-      Object dbVal = read(ctx);
-      if (bean != null) {
-        setValue(bean, dbVal);
-        ctx.propagateState(dbVal);
-      }
-      return dbVal;
-    }
-
-    @Override
-    Object read(DbReadContext ctx) throws SQLException {
-
-      // TODO: Support for Inheritance hierarchy on exported OneToOne ?
-      IdBinder idBinder = targetDescriptor.getIdBinder();
-      Object id = idBinder.read(ctx);
-      if (id == null) {
-        return null;
-      }
-
-      PersistenceContext persistCtx = ctx.getPersistenceContext();
-      Object existing = persistCtx.get(targetType, id);
-
-      if (existing != null) {
-        return existing;
-      }
-      Object ref = targetDescriptor.createReference(ctx.isReadOnly(), id);
-
-      EntityBeanIntercept ebi = ((EntityBean) ref)._ebean_getIntercept();
-      if (Boolean.TRUE.equals(ctx.isReadOnly())) {
-        ebi.setReadOnly(true);
-      }
-      persistCtx.put(id, ref);
-      ctx.register(name, ebi);
-      return ref;
-    }
-
-    /**
-     * Append columns for foreign key columns.
-     */
-    @Override
-    void appendSelect(DbSqlContext ctx, boolean subQuery) {
-
-      // set appropriate tableAlias for the exported id columns
-
-      String relativePrefix = ctx.getRelativePrefix(getName());
-      ctx.pushTableAlias(relativePrefix);
-
-      IdBinder idBinder = targetDescriptor.getIdBinder();
-      idBinder.appendSelect(ctx, subQuery);
-
-      ctx.popTableAlias();
-    }
-
-    @Override
-    void appendFrom(DbSqlContext ctx, SqlJoinType joinType) {
-
-      String relativePrefix = ctx.getRelativePrefix(getName());
-      tableJoin.addJoin(joinType, relativePrefix, ctx);
+      return new AssocOneHelpReference(this);
     }
   }
 

src/main/java/com/avaje/ebeaninternal/server/query/SqlTreeNodeBean.java

@@ -239,7 +239,7 @@ public class SqlTreeNodeBean implements SqlTreeNode {
         localBean = null;
       } else if (!temporalVersions) {
         // check the PersistenceContext to see if the bean already exists
-        contextBean = (EntityBean) persistenceContext.putIfAbsent(id, localBean);
+        contextBean = (EntityBean) localDesc.contextPutIfAbsent(persistenceContext, id, localBean);
         if (contextBean == null) {
           // bean just added to the persistenceContext
           contextBean = localBean;

src/main/java/com/avaje/ebeaninternal/server/text/json/ReadJson.java

@@ -120,7 +120,7 @@ public class ReadJson {
       return null;
     }
 
-    Object contextBean = persistenceContext.putIfAbsent(id, bean);
+    Object contextBean = beanDesc.contextPutIfAbsent(persistenceContext, id, bean);
     if (contextBean == null) {
       if (loadContext != null) {
         EntityBeanIntercept ebi = bean._ebean_getIntercept();

src/main/java/com/avaje/ebeaninternal/server/transaction/DefaultPersistenceContext.java

@@ -1,9 +1,9 @@
 package com.avaje.ebeaninternal.server.transaction;
 
 import com.avaje.ebean.bean.PersistenceContext;
+import com.avaje.ebean.bean.PersistenceContextUtil;
 import com.avaje.ebeaninternal.api.Monitor;
 
-import javax.persistence.Entity;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -31,7 +31,7 @@ public final class DefaultPersistenceContext implements PersistenceContext {
   /**
    * Map used hold caches. One cache per bean type.
    */
-  private final HashMap<String, ClassContext> typeCache = new HashMap<String, ClassContext>();
+  private final HashMap<Class<?>, ClassContext> typeCache = new HashMap<Class<?>, ClassContext>();
 
   private final Monitor monitor = new Monitor();
 
@@ -44,15 +44,15 @@ public final class DefaultPersistenceContext implements PersistenceContext {
   /**
    * Set an object into the PersistenceContext.
    */
-  public void put(Object id, Object bean) {
+  public void put(Class<?> rootType, Object id, Object bean) {
     synchronized (monitor) {
-      getClassContext(bean.getClass()).put(id, bean);
+      getClassContext(rootType).put(id, bean);
     }
   }
 
-  public Object putIfAbsent(Object id, Object bean) {
+  public Object putIfAbsent(Class<?> rootType, Object id, Object bean) {
     synchronized (monitor) {
-      return getClassContext(bean.getClass()).putIfAbsent(id, bean);
+      return getClassContext(rootType).putIfAbsent(id, bean);
     }
   }
 
@@ -76,7 +76,7 @@ public final class DefaultPersistenceContext implements PersistenceContext {
    */
   public int size(Class<?> beanType) {
     synchronized (monitor) {
-      ClassContext classMap = typeCache.get(beanType.getName());
+      ClassContext classMap = typeCache.get(beanType);
       return classMap == null ? 0 : classMap.size();
     }
   }
@@ -92,7 +92,7 @@ public final class DefaultPersistenceContext implements PersistenceContext {
 
   public void clear(Class<?> beanType) {
     synchronized (monitor) {
-      ClassContext classMap = typeCache.get(beanType.getName());
+      ClassContext classMap = typeCache.get(beanType);
       if (classMap != null) {
         classMap.clear();
       }
@@ -101,7 +101,7 @@ public final class DefaultPersistenceContext implements PersistenceContext {
 
   public void deleted(Class<?> beanType, Object id) {
     synchronized (monitor) {
-      ClassContext classMap = typeCache.get(beanType.getName());
+      ClassContext classMap = typeCache.get(beanType);
       if (classMap != null && id != null) {
         classMap.deleted(id);
       }
@@ -110,7 +110,7 @@ public final class DefaultPersistenceContext implements PersistenceContext {
 
   public void clear(Class<?> beanType, Object id) {
     synchronized (monitor) {
-      ClassContext classMap = typeCache.get(beanType.getName());
+      ClassContext classMap = typeCache.get(beanType);
       if (classMap != null && id != null) {
         classMap.remove(id);
       }
@@ -125,25 +125,15 @@ public final class DefaultPersistenceContext implements PersistenceContext {
 
   private ClassContext getClassContext(Class<?> beanType) {
 
-    String clsName = getBeanBaseType(beanType).getName();
-    ClassContext classMap = typeCache.get(clsName);
+    Class<?> rootType =  PersistenceContextUtil.root(beanType);
+    ClassContext classMap = typeCache.get(rootType);
     if (classMap == null) {
       classMap = new ClassContext();
-      typeCache.put(clsName, classMap);
+      typeCache.put(rootType, classMap);
     }
     return classMap;
   }
 
-  private Class<?> getBeanBaseType(Class<?> beanType) {
-    Class<?> parent = beanType.getSuperclass();
-
-    while (parent != null && parent.isAnnotationPresent(Entity.class)) {
-      beanType = parent;
-      parent = parent.getSuperclass();
-    }
-    return beanType;
-  }
-
   private static class ClassContext {
 
     private final Map<Object, Object> map = new HashMap<Object, Object>();

src/main/java/com/avaje/ebeaninternal/server/transaction/NoopPersistenceContext.java

@@ -11,12 +11,12 @@ import com.avaje.ebean.bean.PersistenceContext;
 public class NoopPersistenceContext implements PersistenceContext {
 
   @Override
-  public void put(Object id, Object bean) {
+  public void put(Class<?> rootType, Object id, Object bean) {
     // do nothing
   }
 
   @Override
-  public Object putIfAbsent(Object id, Object bean) {
+  public Object putIfAbsent(Class<?> rootType, Object id, Object bean) {
     // do nothing
     return null;
   }