perf(core): improve result mapping and snapshotting (#4053)

Previously we had to snapshot the entity after it was loaded from
database and hydrated, because the data could be in a different format
(e.g. due to joined strategy). This is now improved and the snapshotting
is no longer needed.

This has a nice side effect, custom types "to db" methods are no longer
triggered after entity gets loaded, only the "from db" ones are.

packages/core/src/EntityManager.ts

@@ -464,7 +464,6 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
         ignoreLazyScalarProperties: true,
         lookup: false,
       });
-      em.unitOfWork.saveSnapshots();
 
       return cached.data;
     }
@@ -637,7 +636,9 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
       em.entityFactory.mergeData(meta, entity, pk!);
     }
 
-    em.unitOfWork.registerManaged(entity, data, { refresh: true });
+    // recompute the data as there might be some values missing (e.g. those with db column defaults)
+    const snapshot = this.comparator.prepareEntity(entity);
+    em.unitOfWork.registerManaged(entity, snapshot, { refresh: true });
 
     return entity;
   }
@@ -787,8 +788,10 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
       }
     }
 
-    for (const [entity, data] of entities) {
-      em.unitOfWork.registerManaged(entity, data, { refresh: true });
+    for (const [entity] of entities) {
+      // recompute the data as there might be some values missing (e.g. those with db column defaults)
+      const snapshot = this.comparator.prepareEntity(entity);
+      em.unitOfWork.registerManaged(entity, snapshot, { refresh: true });
     }
 
     return [...entities.keys()];
@@ -1034,7 +1037,6 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
     entity = Utils.isEntity<Entity>(data) ? data : em.entityFactory.create<Entity>(entityName, data as EntityData<Entity>, { merge: true, ...options });
     em.validator.validate(entity, data, childMeta ?? meta);
     em.unitOfWork.merge(entity);
-    em.unitOfWork.saveSnapshots();
 
     return entity!;
   }

packages/core/src/entity/EntityFactory.ts

@@ -152,21 +152,26 @@ export class EntityFactory {
     const schema = this.driver.getSchemaName(meta, options);
 
     if (Array.isArray(id)) {
+      // composite FK as PK needs to be wrapped for `getPrimaryKeyCondFromArray` to work correctly
+      if (!meta.compositePK && meta.getPrimaryProps()[0].reference !== ReferenceType.SCALAR) {
+        id = [id] as Primary<T>[];
+      }
+
       id = Utils.getPrimaryKeyCondFromArray(id, meta);
     }
 
     const pks = Utils.getOrderedPrimaryKeys<T>(id, meta, this.platform, options.convertCustomTypes);
 
-    if (Utils.isPrimaryKey(id)) {
-      id = { [meta.primaryKeys[0]]: id as Primary<T> };
-    }
-
-    const exists = this.unitOfWork.getById<T>(entityName, pks, schema);
+    const exists = this.unitOfWork.getById<T>(entityName, pks as Primary<T>, schema);
 
     if (exists) {
       return exists;
     }
 
+    if (Utils.isPrimaryKey(id)) {
+      id = { [meta.primaryKeys[0]]: id as Primary<T> };
+    }
+
     return this.create<T>(entityName, id as EntityData<T>, { ...options, initialized: false }) as T;
   }
 

packages/core/src/entity/EntityLoader.ts

@@ -265,9 +265,7 @@ export class EntityLoader {
       schema = children.find(e => e.__helper!.__schema)?.__helper!.__schema;
     }
 
-    const ids = Utils.unique(children.map(e => {
-      return Utils.getPrimaryKeyValues(Reference.unwrapReference(e), e.__meta!.primaryKeys, true);
-    }));
+    const ids = Utils.unique(children.map(e => e.__helper.getPrimaryKey()));
     const where = this.mergePrimaryCondition<T>(ids, fk, options, meta, this.metadata, this.driver.getPlatform());
     const fields = this.buildFields(options.fields, prop);
     const { refresh, filters, convertCustomTypes, lockMode, strategy, populateWhere, connectionType } = options;

packages/core/src/unit-of-work/UnitOfWork.ts

@@ -1,5 +1,14 @@
 import { AsyncLocalStorage } from 'async_hooks';
-import type { AnyEntity, Dictionary, EntityData, EntityMetadata, EntityProperty, FilterQuery, IPrimaryKeyValue, Primary } from '../typings';
+import type {
+  AnyEntity,
+  Dictionary,
+  EntityData,
+  EntityMetadata,
+  EntityProperty,
+  FilterQuery,
+  IPrimaryKeyValue,
+  Primary,
+} from '../typings';
 import { Collection, EntityIdentifier, helper, Reference } from '../entity';
 import { ChangeSet, ChangeSetType } from './ChangeSet';
 import { ChangeSetComputer } from './ChangeSetComputer';
@@ -37,7 +46,6 @@ export class UnitOfWork {
   private readonly queuedActions = new Set<string>();
   private readonly loadedEntities = new Set<AnyEntity>();
   private readonly flushQueue: (() => Promise<void>)[] = [];
-  private readonly snapshotQueue: (() => void)[] = [];
   private working = false;
   private insideHooks = false;
 
@@ -63,24 +71,12 @@ export class UnitOfWork {
     // as there can be some entity with already changed state that is not yet flushed
     if (wrapped.__initialized && (!visited || !wrapped.__originalEntityData)) {
       wrapped.__originalEntityData = this.comparator.prepareEntity(entity);
-      this.queuedActions.delete(wrapped.__meta.className);
       wrapped.__touched = false;
     }
 
     this.cascade(entity, Cascade.MERGE, visited ?? new Set<AnyEntity>());
   }
 
-  /**
-   * @internal
-   */
-  saveSnapshots() {
-    for (const callback of this.snapshotQueue) {
-      callback();
-    }
-
-    this.snapshotQueue.length = 0;
-  }
-
   /**
    * @internal
    */
@@ -91,31 +87,25 @@ export class UnitOfWork {
       return entity;
     }
 
-    if (options?.loaded && helper(entity).__initialized && !helper(entity).__onLoadFired) {
+    const wrapped = helper(entity);
+
+    if (options?.loaded && wrapped.__initialized && !wrapped.__onLoadFired) {
       this.loadedEntities.add(entity as AnyEntity);
     }
 
-    const wrapped = helper(entity);
     wrapped.__em ??= this.em;
     wrapped.__managed = true;
 
     if (data && (options?.refresh || !wrapped.__originalEntityData)) {
-      Object.keys(data).forEach(key => helper(entity).__loadedProperties.add(key));
-      this.queuedActions.delete(wrapped.__meta.className);
+      Object.keys(data).forEach(key => wrapped.__loadedProperties.add(key));
 
-      // assign the data early, delay recompute via snapshot queue unless we refresh the state
-      if (options?.refresh) {
-        wrapped.__originalEntityData = this.comparator.prepareEntity(entity);
-      } else {
-        wrapped.__originalEntityData = data;
-      }
-
-      wrapped.__touched = false;
-      this.snapshotQueue.push(() => {
-        // we can't use the `data` directly here as it can contain fetch joined data, that can't be used for diffing the state
-        wrapped.__originalEntityData = this.comparator.prepareEntity(entity);
-        wrapped.__touched = false;
+      wrapped.__meta.relations.forEach(prop => {
+        if (Utils.isPlainObject(data[prop.name as string])) {
+          data[prop.name as string] = Utils.getPrimaryKeyValues(data[prop.name as string], wrapped.__meta.primaryKeys, true);
+        }
       });
+      wrapped.__originalEntityData = data;
+      wrapped.__touched = false;
     }
 
     return entity;
@@ -125,8 +115,6 @@ export class UnitOfWork {
    * @internal
    */
   async dispatchOnLoadEvent(): Promise<void> {
-    this.saveSnapshots();
-
     for (const entity of this.loadedEntities) {
       if (this.eventManager.hasListeners(EventType.onLoad, entity.__meta!)) {
         await this.eventManager.dispatchEvent(EventType.onLoad, { entity, em: this.em });
@@ -248,7 +236,6 @@ export class UnitOfWork {
     this.initIdentifier(entity);
     this.changeSets.set(entity, cs);
     this.persistStack.delete(entity);
-    this.queuedActions.delete(cs.name);
     helper(entity).__originalEntityData = this.comparator.prepareEntity(entity);
     helper(entity).__touched = false;
   }
@@ -863,8 +850,6 @@ export class UnitOfWork {
     for (const changeSet of this.changeSets.values()) {
       this.takeCollectionSnapshots(changeSet.entity, visited);
     }
-
-    this.saveSnapshots();
   }
 
   private async commitCreateChangeSets<T extends object>(changeSets: ChangeSet<T>[], ctx?: Transaction): Promise<void> {
@@ -887,8 +872,6 @@ export class UnitOfWork {
       this.registerManaged<T>(changeSet.entity, changeSet.payload, { refresh: true });
       await this.runHooks(EventType.afterCreate, changeSet);
     }
-
-    this.saveSnapshots();
   }
 
   private findExtraUpdates<T extends object>(changeSet: ChangeSet<T>, props: EntityProperty<T>[]): void {

packages/core/src/utils/EntityComparator.ts

@@ -272,6 +272,10 @@ export class EntityComparator {
         idx += pk.fieldNames.length;
       }
 
+      if (parts.length < 2) {
+        return parts[0];
+      }
+
       return '[' + parts.join(', ') + ']';
     };
 
@@ -292,6 +296,16 @@ export class EntityComparator {
 
       if (prop.type === 'boolean') {
         lines.push(`  if (typeof ${propName(prop.fieldNames[0])} !== 'undefined') { ret${this.wrap(prop.name)} = ${propName(prop.fieldNames[0])} == null ? ${propName(prop.fieldNames[0])} : !!${propName(prop.fieldNames[0])}; ${propName(prop.fieldNames[0], 'mapped')} = true; }`);
+      } else if (prop.type.toLowerCase() === 'date') {
+        if (prop.embedded && meta.properties[prop.embedded[0]].object) {
+          const entityKey = 'ret.' + prop.fieldNames[0];
+          const entityKeyOptional = entityKey.replace(/\./g, '?.');
+          lines.push(`  if (typeof ${entityKeyOptional} !== 'undefined') { ${entityKey} = ${entityKey} == null ? ${entityKey} : new Date(${entityKey}); }`);
+        } else {
+          lines.push(`  if (typeof ${propName(prop.fieldNames[0])} !== 'undefined') { ret${this.wrap(prop.name)} = new Date(${propName(prop.fieldNames[0])}); ${propName(prop.fieldNames[0], 'mapped')} = true; }`);
+        }
+      } else if (prop.reference === ReferenceType.EMBEDDED && prop.object && !this.platform.convertsJsonAutomatically()) {
+        lines.push(`  if (typeof ${propName(prop.fieldNames[0])} !== 'undefined') { ret${this.wrap(prop.name)} = ${propName(prop.fieldNames[0])} == null ? ${propName(prop.fieldNames[0])} : JSON.parse(${propName(prop.fieldNames[0])}); ${propName(prop.fieldNames[0], 'mapped')} = true; }`);
       } else {
         lines.push(`  if (typeof ${propName(prop.fieldNames[0])} !== 'undefined') { ret${this.wrap(prop.name)} = ${propName(prop.fieldNames[0])}; ${propName(prop.fieldNames[0], 'mapped')} = true; }`);
       }
@@ -314,7 +328,7 @@ export class EntityComparator {
     }
 
     let tail = '';
-    let ret = parts
+    const ret = parts
       .map(k => {
         if (k.match(/^\[idx_\d+]$/)) {
           tail += k;
@@ -328,16 +342,6 @@ export class EntityComparator {
       })
       .filter(k => k)
       .join(' && ');
-    const isRef = [ReferenceType.ONE_TO_ONE, ReferenceType.MANY_TO_ONE].includes(prop.reference) && !prop.mapToPk;
-    const isSetter = isRef && !!(prop.inversedBy || prop.mappedBy);
-
-    if (prop.primary || isSetter) {
-      ret += ` && entity${entityKey} != null`;
-    }
-
-    if (isRef) {
-      ret += ` && (entity${entityKey} == null || entity${entityKey}.__helper?.hasPrimaryKey())`;
-    }
 
     return ret;
   }
@@ -464,9 +468,20 @@ export class EntityComparator {
           ret += `    ret${dataKey} = entity${entityKey};\n`;
         }
       } else {
-        const meta2 = this.metadata.find(prop.type);
-        context.set(`getPrimaryKeyValues_${convertorKey}`, (val: any) => val && Utils.getPrimaryKeyValues(val, meta2!.primaryKeys, true, true));
-        ret += `    ret${dataKey} = getPrimaryKeyValues_${convertorKey}(entity${entityKey});\n`;
+        const toArray = (val: unknown): unknown => {
+          if (Utils.isPlainObject(val)) {
+            return Object.values(val).map(v => toArray(v));
+          }
+
+          return val;
+        };
+
+        context.set('toArray', toArray);
+        ret += `    if (entity${entityKey} === null) {\n`;
+        ret += `      ret${dataKey} = null;\n`;
+        ret += `    } else if (typeof entity${entityKey} !== 'undefined') {\n`;
+        ret += `      ret${dataKey} = toArray(entity${entityKey}.__helper.getPrimaryKey(true));\n`;
+        ret += `    }\n`;
       }
 
       return ret + '  }\n';

packages/knex/src/AbstractSqlDriver.ts

@@ -340,7 +340,7 @@ export abstract class AbstractSqlDriver<Connection extends AbstractSqlConnection
         props.forEach(prop => {
           if (prop.fieldNames.length > 1) {
             params.push(...(row[prop.name] as unknown[] ?? prop.fieldNames.map(() => null)));
-            keys.push(...prop.fieldNames.map(_ => '?'));
+            keys.push(...(row[prop.name] as unknown[] ?? prop.fieldNames).map(() => '?'));
           } else if (prop.customType && 'convertToDatabaseValueSQL' in prop.customType && !this.platform.isRaw(row[prop.name])) {
             keys.push(prop.customType.convertToDatabaseValueSQL!('?', this.platform));
             params.push(row[prop.name]);

packages/knex/src/query/QueryBuilder.ts

@@ -911,7 +911,6 @@ export class QueryBuilder<T extends object = AnyEntity> {
 
     if (data) {
       if (Utils.isEntity(data)) {
-        // data = helper(data).toJSON();
         data = this.em?.getComparator().prepareEntity(data as T) ?? serialize(data as T);
       }
 

packages/knex/src/query/QueryBuilderHelper.ts

@@ -778,7 +778,7 @@ export class QueryBuilderHelper {
       }
 
       if (prop.joinColumns && Array.isArray(data[k])) {
-        const copy = data[k];
+        const copy = Utils.flatten(data[k]);
         delete data[k];
         prop.joinColumns.forEach((joinColumn, idx) => data[joinColumn] = copy[idx]);
 

tests/EntityFactory.test.ts

@@ -220,7 +220,7 @@ describe('EntityFactory', () => {
     expect(mock.mock.calls).toHaveLength(3);
     expect(mock.mock.calls[0][0]).toMatch(/db\.getCollection\('book-tag'\)\.insertMany\(\[ { name: 't1' } ], {}\);/);
     expect(mock.mock.calls[1][0]).toMatch(/db\.getCollection\('author'\)\.insertMany\(\[ { createdAt: ISODate\('.*'\), updatedAt: ISODate\('.*'\), foo: 'bar', name: 'Jon', email: 'jon@snow\.com', termsAccepted: false } ], {}\);/);
-    expect(mock.mock.calls[2][0]).toMatch(/db\.getCollection\('books-table'\)\.insertMany\(\[ { createdAt: ISODate\('.*'\), title: 'B1', publisher: ObjectId\('5b0d19b28b21c648c2c8a600'\), author: ObjectId\('.*'\), tags: \[ ObjectId\('.*'\), ObjectId\('5b0d19b28b21c648c2c8a601'\) ] } ], {}\);/);
+    expect(mock.mock.calls[2][0]).toMatch(/db\.getCollection\('books-table'\)\.insertMany\(\[ { createdAt: ISODate\('.*'\), title: 'B1', author: ObjectId\('.*'\), publisher: ObjectId\('5b0d19b28b21c648c2c8a600'\), tags: \[ ObjectId\('.*'\), ObjectId\('5b0d19b28b21c648c2c8a601'\) ] } ], {}\);/);
 
     orm.em.clear();
     mock.mock.calls.length = 0;
@@ -249,7 +249,7 @@ describe('EntityFactory', () => {
     expect(mock.mock.calls.length).toBe(3);
     expect(mock.mock.calls[0][0]).toMatch(/db\.getCollection\('book-tag'\)\.insertMany\(\[ { name: 't1' } ], {}\);/);
     expect(mock.mock.calls[1][0]).toMatch(/db\.getCollection\('author'\)\.insertMany\(\[ { createdAt: ISODate\('.*'\), updatedAt: ISODate\('.*'\), foo: 'bar', name: 'Jon', email: 'jon2@snow\.com', termsAccepted: false } ], {}\);/);
-    expect(mock.mock.calls[2][0]).toMatch(/db\.getCollection\('books-table'\)\.insertMany\(\[ { createdAt: ISODate\('.*'\), title: 'B1', publisher: ObjectId\('5b0d19b28b21c648c2c8a600'\), author: ObjectId\('.*'\), tags: \[ ObjectId\('.*'\), ObjectId\('5b0d19b28b21c648c2c8a601'\) ] } ], {}\);/);
+    expect(mock.mock.calls[2][0]).toMatch(/db\.getCollection\('books-table'\)\.insertMany\(\[ { createdAt: ISODate\('.*'\), title: 'B1', author: ObjectId\('.*'\), publisher: ObjectId\('5b0d19b28b21c648c2c8a600'\), tags: \[ ObjectId\('.*'\), ObjectId\('5b0d19b28b21c648c2c8a601'\) ] } ], {}\);/);
   });
 
   test('em.create() should not mutate the input object (GH issue 1294)', async () => {

tests/EntityManager.postgre.test.ts

@@ -2019,7 +2019,7 @@ describe('EntityManagerPostgre', () => {
 
     expect(mock.mock.calls[0][0]).toMatch('begin');
     expect(mock.mock.calls[1][0]).toMatch('insert into "author2" ("created_at", "updated_at", "name", "email", "terms_accepted") values ($1, $2, $3, $4, $5) returning "id", "created_at", "updated_at", "age", "terms_accepted"');
-    expect(mock.mock.calls[2][0]).toMatch('insert into "book2" ("uuid_pk", "created_at", "title", "price", "publisher_id", "author_id") values ($1, $2, $3, $4, $5, $6) returning "uuid_pk", "created_at", "title"');
+    expect(mock.mock.calls[2][0]).toMatch('insert into "book2" ("uuid_pk", "created_at", "title", "price", "author_id", "publisher_id") values ($1, $2, $3, $4, $5, $6) returning "uuid_pk", "created_at", "title"');
     expect(mock.mock.calls[3][0]).toMatch('commit');
     expect(mock.mock.calls[4][0]).toMatch('begin');
     expect(mock.mock.calls[5][0]).toMatch('update "book2" set "publisher_id" = $1 where "uuid_pk" = $2');

tests/features/embeddables/GH2391.test.ts

@@ -67,7 +67,7 @@ describe('onCreate and onUpdate in embeddables (GH 2283 and 2391)', () => {
   });
 
   test(`GH issue 2283, 2391`, async () => {
-    let line = orm.em.create(MyEntity, {});
+    let line = orm.em.create(MyEntity, {}, { persist: false });
     await orm.em.fork().persistAndFlush(line);
 
     expect(!!line.audit1.created).toBeTruthy();

tests/features/embeddables/__snapshots__/embeddable-custom-types.postgres.test.ts.snap

@@ -20,7 +20,7 @@ drop table if exists "parent" cascade;
 exports[`embedded entities with custom types snapshot generator 1`] = `
 "function(entity) {
   const ret = {};
-  if (typeof entity.id !== 'undefined' && entity.id != null) {
+  if (typeof entity.id !== 'undefined') {
     ret.id = entity.id;
   }