fix(core): do not override loaded entity state when loading it again

Previously, when you fetched one entity, updated it and then fetched it again from other place before persisting, the entity state would be overridden, resulting in missing update query and lost state).

lib/EntityManager.ts

@@ -78,7 +78,7 @@ export class EntityManager {
   async findOne<T extends IEntityType<T>>(entityName: EntityName<T>, where: FilterQuery<T> | IPrimaryKey, populate?: string[] | FindOneOptions, orderBy?: Record<string, QueryOrder>): Promise<T | null> {
     entityName = Utils.className(entityName);
     this.validator.validateEmptyWhere(where);
-
+    where = SmartQueryHelper.processWhere(where as FilterQuery<T>, entityName);
     let entity = this.getUnitOfWork().tryGetById<T>(entityName, where);
     const options = Utils.isObject<FindOneOptions>(populate) ? populate : { populate, orderBy };
 
@@ -87,7 +87,6 @@ export class EntityManager {
       return entity;
     }
 
-    where = SmartQueryHelper.processWhere(where as FilterQuery<T>, entityName);
     this.validator.validateParams(where);
     const data = await this.driver.findOne(entityName, where, options.populate, options.orderBy);
 
@@ -169,7 +168,13 @@ export class EntityManager {
 
     entityName = Utils.className(entityName);
     this.validator.validatePrimaryKey(data!, this.metadata[entityName]);
-    const entity = Utils.isEntity<T>(data) ? data : this.getEntityFactory().create<T>(entityName, data!, true);
+    let entity = this.getUnitOfWork().tryGetById<T>(entityName, data!);
+
+    if (entity && entity.isInitialized()) {
+      return entity;
+    }
+
+    entity = Utils.isEntity<T>(data) ? data : this.getEntityFactory().create<T>(entityName, data!, true);
 
     // add to IM immediately - needed for self-references that can be part of `data` (and do not trigger cascade merge)
     this.getUnitOfWork().merge(entity, [entity]);
@@ -191,7 +196,7 @@ export class EntityManager {
    */
   getReference<T extends IEntityType<T>>(entityName: EntityName<T>, id: IPrimaryKey): T {
     const entity = this.getEntityFactory().createReference<T>(entityName, id);
-    this.getUnitOfWork().merge(entity);
+    this.getUnitOfWork().merge(entity, [], false);
 
     return entity;
   }

lib/unit-of-work/ChangeSetComputer.ts

@@ -75,11 +75,11 @@ export class ChangeSetComputer {
 
   private processOneToOne<T extends IEntityType<T>>(prop: EntityProperty<T>, changeSet: ChangeSet<T>): void {
     // check diff, if we had a value on 1:1 before and now it changed (nulled or replaced), we need to trigger orphan removal)
-    const data = this.originalEntityData[changeSet.entity.__uuid];
+    const data = this.originalEntityData[changeSet.entity.__uuid] as EntityData<T>;
 
     if (prop.orphanRemoval && data && data[prop.name] && prop.name in changeSet.payload) {
       const em = changeSet.entity.__em;
-      const orphan = em.getReference(data[prop.name].constructor.name, data[prop.name].__primaryKey);
+      const orphan = em.getReference(prop.type, data[prop.name]);
       em.getUnitOfWork().scheduleOrphanRemoval(orphan);
     }
   }

lib/unit-of-work/UnitOfWork.ts

@@ -14,7 +14,7 @@ export class UnitOfWork {
   private readonly identityMap = {} as Record<string, IEntity>;
 
   /** holds copy of identity map so we can compute changes when persisting managed entities */
-  private readonly originalEntityData = {} as Record<string, IEntity>;
+  private readonly originalEntityData = {} as Record<string, EntityData<IEntity>>;
 
   /** map of wrapped primary keys so we can compute change set without eager commit */
   private readonly identifierMap = {} as Record<string, EntityIdentifier>;
@@ -31,13 +31,17 @@ export class UnitOfWork {
 
   constructor(private readonly em: EntityManager) { }
 
-  merge<T extends IEntityType<T>>(entity: T, visited: IEntity[] = []): void {
+  merge<T extends IEntityType<T>>(entity: T, visited: IEntity[] = [], mergeData = true): void {
     if (!entity.__primaryKey) {
       return;
     }
 
     this.identityMap[`${entity.constructor.name}-${entity.__serializedPrimaryKey}`] = entity;
-    this.originalEntityData[entity.__uuid] = Utils.copy(entity);
+
+    if (!this.originalEntityData[entity.__uuid] || mergeData) {
+      this.originalEntityData[entity.__uuid] = Utils.prepareEntity(entity);
+    }
+
     this.cascade(entity, Cascade.MERGE, visited);
   }
 
@@ -47,11 +51,13 @@ export class UnitOfWork {
   }
 
   tryGetById<T extends IEntityType<T>>(entityName: string, where: FilterQuery<T> | IPrimaryKey): T | null {
-    if (!Utils.isPrimaryKey(where)) {
+    const pk = Utils.extractPK(where, this.metadata[entityName]);
+
+    if (!pk) {
       return null;
     }
 
-    return this.getById<T>(entityName, where);
+    return this.getById<T>(entityName, pk);
   }
 
   getIdentityMap(): Record<string, IEntity> {
@@ -168,7 +174,7 @@ export class UnitOfWork {
     if (changeSet) {
       this.changeSets.push(changeSet);
       this.cleanUpStack(this.persistStack, entity);
-      this.originalEntityData[entity.__uuid] = Utils.copy(entity);
+      this.originalEntityData[entity.__uuid] = Utils.prepareEntity(entity);
     }
   }
 
@@ -213,8 +219,10 @@ export class UnitOfWork {
     await this.runHooks(`before${type}`, changeSet.entity, changeSet.payload);
     await this.changeSetPersister.persistToDatabase(changeSet);
 
-    if (changeSet.type !== ChangeSetType.DELETE) {
-      this.em.merge(changeSet.entity);
+    switch (changeSet.type) {
+      case ChangeSetType.CREATE: this.em.merge(changeSet.entity); break;
+      case ChangeSetType.UPDATE: this.merge(changeSet.entity); break;
+      case ChangeSetType.DELETE: this.unsetIdentity(changeSet.entity); break;
     }
 
     await this.runHooks(`after${type}`, changeSet.entity);
@@ -271,7 +279,7 @@ export class UnitOfWork {
 
     switch (type) {
       case Cascade.PERSIST: this.persist(entity, visited); break;
-      case Cascade.MERGE: this.merge(entity, visited); break;
+      case Cascade.MERGE: this.merge(entity, visited, false); break;
       case Cascade.REMOVE: this.remove(entity, visited); break;
     }
 

tests/EntityManager.mongo.test.ts

@@ -1148,10 +1148,12 @@ describe('EntityManagerMongo', () => {
     const baz2 = FooBaz.create('fz2');
     bar.baz = baz1;
     await orm.em.persist(bar);
+    expect(orm.em.getUnitOfWork()['originalEntityData'][bar.__uuid].baz).toEqual(baz1._id);
 
     // replacing reference with value will trigger orphan removal
     bar.baz = baz2;
     await orm.em.persist(bar);
+    expect(orm.em.getUnitOfWork()['originalEntityData'][bar.__uuid].baz).toEqual(baz2._id);
     await expect(orm.em.findOne(FooBaz, baz1)).resolves.toBeNull();
     await expect(orm.em.findOne(FooBaz, baz2)).resolves.not.toBeNull();
 
@@ -1173,6 +1175,56 @@ describe('EntityManagerMongo', () => {
     await expect(orm.em.commit()).rejects.toThrowError('Transactions are not supported by current driver');
   });
 
+  test('loading connected entity will not update identity map for associations', async () => {
+    const author = new Author('Jon Snow', 'snow@wall.st');
+    author.favouriteBook = new Book('b1', author);
+    await orm.em.persist(author);
+    orm.em.clear();
+
+    const a = (await orm.em.findOne(Author, author, ['favouriteBook']))!;
+    expect(a).not.toBe(author);
+    a.name = 'test 1';
+    a.favouriteBook.title = 'test 2';
+    const a1 = (await orm.em.findOne(Author, { favouriteBook: a.favouriteBook }))!;
+    const b1 = (await orm.em.findOne(Book, { author }))!;
+    expect(a.name).toBe('test 1');
+    expect(a.favouriteBook.title).toBe('test 2');
+    expect(a1.name).toBe('test 1');
+    expect(b1.title).toBe('test 2');
+    await orm.em.persist(a);
+    orm.em.clear();
+
+    const a2 = (await orm.em.findOne(Author, author))!;
+    const b2 = (await orm.em.findOne(Book, { author }))!;
+    expect(a2.name).toBe('test 1');
+    expect(b2.title).toBe('test 2');
+  });
+
+  test('getReference will not update identity map copy', async () => {
+    const author = new Author('Jon Snow', 'snow@wall.st');
+    author.favouriteBook = new Book('b1', author);
+    await orm.em.persist(author);
+    orm.em.clear();
+
+    const a = (await orm.em.findOne(Author, author, ['favouriteBook']))!;
+    expect(a).not.toBe(author);
+    a.name = 'test 1';
+    a.favouriteBook.title = 'test 2';
+    const a1 = orm.em.getReference(Author, a.id)!;
+    const b1 = orm.em.getReference(Book, a.favouriteBook.id)!;
+    expect(a.name).toBe('test 1');
+    expect(a.favouriteBook.title).toBe('test 2');
+    expect(a1.name).toBe('test 1');
+    expect(b1.title).toBe('test 2');
+    await orm.em.persist(a);
+    orm.em.clear();
+
+    const a2 = (await orm.em.findOne(Author, author))!;
+    const b2 = (await orm.em.findOne(Book, { author }))!;
+    expect(a2.name).toBe('test 1');
+    expect(b2.title).toBe('test 2');
+  });
+
   afterAll(async () => orm.close(true));
 
 });

tests/bootstrap.ts

@@ -8,6 +8,8 @@ import { BaseEntity2 } from './entities-sql/BaseEntity2';
 import { FooBar2 } from './entities-sql/FooBar2';
 import { BaseEntity22 } from './entities-sql/BaseEntity22';
 import { PostgreSqlConnection } from '../lib/connections/PostgreSqlConnection';
+import { FooBaz } from './entities/FooBaz';
+import { FooBar } from './entities/FooBar';
 
 const { BaseEntity4, Author3, Book3, BookTag3, Publisher3, Test3 } = require('./entities-js');
 
@@ -96,6 +98,8 @@ export async function wipeDatabase(em: EntityManager) {
   await em.getRepository(BookTag).remove({});
   await em.getRepository(Publisher).remove({});
   await em.getRepository(Test).remove({});
+  await em.getRepository(FooBar).remove({});
+  await em.getRepository(FooBaz).remove({});
   em.clear();
 }