refactor: expose Loaded type hints + add New hint

Related: #214, #691

packages/core/src/EntityManager.ts

@@ -5,7 +5,7 @@ import { Configuration, QueryHelper, RequestContext, Utils, ValidationError } fr
 import { EntityAssigner, EntityFactory, EntityLoader, EntityRepository, EntityValidator, IdentifiedReference, LoadStrategy, Reference, ReferenceType, wrap } from './entity';
 import { LockMode, UnitOfWork } from './unit-of-work';
 import { CountOptions, DeleteOptions, EntityManagerType, FindOneOptions, FindOneOrFailOptions, FindOptions, IDatabaseDriver, UpdateOptions } from './drivers';
-import { AnyEntity, Dictionary, EntityData, EntityMetadata, EntityName, FilterDef, FilterQuery, Loaded, Primary, Populate, PopulateMap, PopulateOptions } from './typings';
+import { AnyEntity, Dictionary, EntityData, EntityMetadata, EntityName, FilterDef, FilterQuery, Loaded, Primary, Populate, PopulateMap, PopulateOptions, New } from './typings';
 import { QueryOrderMap } from './enums';
 import { MetadataStorage } from './metadata';
 import { Transaction } from './connections';
@@ -77,12 +77,12 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
   /**
    * Finds all entities matching your `where` query.
    */
-  async find<T, P extends Populate<T> = any>(entityName: EntityName<T>, where: FilterQuery<T>, populate?: Populate<T>, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]>;
+  async find<T, P extends Populate<T> = any>(entityName: EntityName<T>, where: FilterQuery<T>, populate?: P, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]>;
 
   /**
    * Finds all entities matching your `where` query.
    */
-  async find<T, P extends Populate<T> = any>(entityName: EntityName<T>, where: FilterQuery<T>, populate?: Populate<T> | FindOptions<T, P>, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]> {
+  async find<T, P extends Populate<T> = any>(entityName: EntityName<T>, where: FilterQuery<T>, populate?: P | FindOptions<T, P>, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]> {
     const options = Utils.isObject<FindOptions<T, P>>(populate) ? populate : { populate, orderBy, limit, offset } as FindOptions<T, P>;
     entityName = Utils.className(entityName);
     where = QueryHelper.processWhere(where, entityName, this.metadata);
@@ -382,8 +382,8 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
   /**
    * Creates new instance of given entity and populates it with given data
    */
-  create<T>(entityName: EntityName<T>, data: EntityData<T>): T {
-    return this.getEntityFactory().create(entityName, data, true, true);
+  create<T, P extends Populate<T> = string[]>(entityName: EntityName<T>, data: EntityData<T, P>): New<T, P> {
+    return this.getEntityFactory().create<T, P>(entityName, data, true, true);
   }
 
   /**

packages/core/src/entity/BaseEntity.ts

@@ -1,6 +1,6 @@
 import { wrap } from './wrap';
 import { IdentifiedReference, Reference } from './Reference';
-import { Dictionary, EntityData, IWrappedEntity } from '../typings';
+import { Dictionary, EntityData, IWrappedEntity, LoadedReference, Populate } from '../typings';
 import { AssignOptions, EntityAssigner } from './EntityAssigner';
 
 export abstract class BaseEntity<T, PK extends keyof T> implements IWrappedEntity<T, PK> {
@@ -13,8 +13,8 @@ export abstract class BaseEntity<T, PK extends keyof T> implements IWrappedEntit
     wrap(this, true).populated(populated);
   }
 
-  toReference(): IdentifiedReference<T, PK> {
-    return Reference.create<T, PK>(this as unknown as T);
+  toReference<PK2 extends PK = never, P extends Populate<T> = never>(): IdentifiedReference<T, PK2> & LoadedReference<T, P> {
+    return Reference.create<T, PK>(this as unknown as T) as IdentifiedReference<T, PK> & LoadedReference<T>;
   }
 
   toObject(ignoreFields: string[] = []): EntityData<T> {

packages/core/src/entity/EntityFactory.ts

@@ -1,5 +1,5 @@
 import { Utils } from '../utils';
-import { EntityData, EntityMetadata, EntityName, EntityProperty, Primary } from '../typings';
+import { EntityData, EntityMetadata, EntityName, EntityProperty, New, Populate, Primary } from '../typings';
 import { UnitOfWork } from '../unit-of-work';
 import { ReferenceType } from './enums';
 import { EntityManager, EventType, wrap } from '..';
@@ -16,9 +16,9 @@ export class EntityFactory {
   constructor(private readonly unitOfWork: UnitOfWork,
               private readonly em: EntityManager) { }
 
-  create<T>(entityName: EntityName<T>, data: EntityData<T>, initialized = true, newEntity = false): T {
+  create<T, P extends Populate<T> = keyof T>(entityName: EntityName<T>, data: EntityData<T>, initialized = true, newEntity = false): New<T, P> {
     if (Utils.isEntity<T>(data)) {
-      return data;
+      return data as New<T, P>;
     }
 
     entityName = Utils.className(entityName);
@@ -35,7 +35,7 @@ export class EntityFactory {
 
     this.runHooks(entity, meta);
 
-    return entity;
+    return entity as New<T, P>;
   }
 
   createReference<T>(entityName: EntityName<T>, id: Primary<T> | Primary<T>[] | Record<string, Primary<T>>): T {

packages/core/src/entity/EntityRepository.ts

@@ -1,5 +1,5 @@
 import { EntityManager } from '../EntityManager';
-import { EntityData, EntityName, AnyEntity, Primary, Populate, Loaded } from '../typings';
+import { EntityData, EntityName, AnyEntity, Primary, Populate, Loaded, New } from '../typings';
 import { QueryOrderMap } from '../enums';
 import { FilterQuery, FindOneOptions, FindOptions, FindOneOrFailOptions, IdentifiedReference, Reference } from '..';
 
@@ -36,21 +36,21 @@ export class EntityRepository<T extends AnyEntity<T>> {
   }
 
   async find<P extends Populate<T> = any>(where: FilterQuery<T>, options?: FindOptions<T, P>): Promise<Loaded<T, P>[]>;
-  async find<P extends Populate<T> = any>(where: FilterQuery<T>, populate?: string[] | boolean, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]>;
-  async find<P extends Populate<T> = any>(where: FilterQuery<T>, populate: string[] | boolean | FindOptions<T, P> = [], orderBy: QueryOrderMap = {}, limit?: number, offset?: number): Promise<Loaded<T, P>[]> {
+  async find<P extends Populate<T> = any>(where: FilterQuery<T>, populate?: P, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]>;
+  async find<P extends Populate<T> = any>(where: FilterQuery<T>, populate?: P | FindOptions<T, P>, orderBy: QueryOrderMap = {}, limit?: number, offset?: number): Promise<Loaded<T, P>[]> {
     return this.em.find<T, P>(this.entityName, where as FilterQuery<T>, populate as P, orderBy, limit, offset);
   }
 
   async findAndCount<P extends Populate<T> = any>(where: FilterQuery<T>, options?: FindOptions<T>): Promise<[Loaded<T, P>[], number]>;
-  async findAndCount<P extends Populate<T> = any>(where: FilterQuery<T>, populate?: string[] | boolean, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<[Loaded<T, P>[], number]>;
-  async findAndCount<P extends Populate<T> = any>(where: FilterQuery<T>, populate: string[] | boolean | FindOptions<T> = [], orderBy: QueryOrderMap = {}, limit?: number, offset?: number): Promise<[Loaded<T, P>[], number]> {
+  async findAndCount<P extends Populate<T> = any>(where: FilterQuery<T>, populate?: P, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<[Loaded<T, P>[], number]>;
+  async findAndCount<P extends Populate<T> = any>(where: FilterQuery<T>, populate?: P | FindOptions<T>, orderBy: QueryOrderMap = {}, limit?: number, offset?: number): Promise<[Loaded<T, P>[], number]> {
     return this.em.findAndCount<T, P>(this.entityName, where as FilterQuery<T>, populate as P, orderBy, limit, offset);
   }
 
-  async findAll<P extends Populate<T> = any>(options?: FindOptions<T>): Promise<Loaded<T, P>[]>;
-  async findAll<P extends Populate<T> = any>(populate?: string[] | boolean | true, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]>;
-  async findAll<P extends Populate<T> = any>(populate: string[] | boolean | true | FindOptions<T> = [], orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]> {
-    return this.em.find<T, P>(this.entityName, {}, populate as string[], orderBy, limit, offset);
+  async findAll<P extends Populate<T> = any>(options?: FindOptions<T, P>): Promise<Loaded<T, P>[]>;
+  async findAll<P extends Populate<T> = any>(populate?: P, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]>;
+  async findAll<P extends Populate<T> = any>(populate?: P | FindOptions<T, P>, orderBy?: QueryOrderMap, limit?: number, offset?: number): Promise<Loaded<T, P>[]> {
+    return this.em.find<T, P>(this.entityName, {}, populate as P, orderBy, limit, offset);
   }
 
   remove(entity: AnyEntity): EntityManager {
@@ -112,8 +112,8 @@ export class EntityRepository<T extends AnyEntity<T>> {
   /**
    * Creates new instance of given entity and populates it with given data
    */
-  create(data: EntityData<T>): T {
-    return this.em.create<T>(this.entityName, data);
+  create<P extends Populate<T> = string[]>(data: EntityData<T, P>): New<T, P> {
+    return this.em.create<T, P>(this.entityName, data);
   }
 
   /**

packages/core/src/entity/wrap.ts

@@ -6,24 +6,28 @@ import { BaseEntity } from './BaseEntity';
 /**
  * returns WrappedEntity instance associated with this entity. This includes all the internal properties like `__meta` or `__em`.
  */
-export function wrap<T>(entity: T, preferHelper: true): IWrappedEntityInternal<T, keyof T>;
+export function wrap<T, PK extends keyof T>(entity: T, preferHelper: true): IWrappedEntityInternal<T, PK>;
 
 /**
  * wraps entity type with WrappedEntity internal properties and helpers like init/isInitialized/populated/toJSON
  */
-export function wrap<T>(entity: T, preferHelper?: false): IWrappedEntity<T, keyof T>;
+export function wrap<T, PK extends keyof T>(entity: T, preferHelper?: false): IWrappedEntity<T, PK>;
 
 /**
  * wraps entity type with WrappedEntity internal properties and helpers like init/isInitialized/populated/toJSON
  * use `preferHelper = true` to have access to the internal `__` properties like `__meta` or `__em`
  */
-export function wrap<T>(entity: T, preferHelper = false): IWrappedEntity<T, keyof T> | IWrappedEntityInternal<T, keyof T> {
+export function wrap<T, PK extends keyof T>(entity: T, preferHelper = false): IWrappedEntity<T, PK> | IWrappedEntityInternal<T, PK> {
   if (entity instanceof BaseEntity && !preferHelper) {
-    return entity as IWrappedEntity<T, keyof T>;
+    return entity as unknown as IWrappedEntity<T, PK>;
   }
 
   if (entity instanceof ArrayCollection) {
-    return entity as unknown as WrappedEntity<T, keyof T>;
+    return entity as unknown as IWrappedEntity<T, PK>;
+  }
+
+  if (!entity) {
+    return entity as unknown as IWrappedEntity<T, PK>;
   }
 
   return (entity as Dictionary).__helper;

packages/core/src/index.ts

@@ -1,7 +1,7 @@
 /* istanbul ignore file */
 export {
   Constructor, Dictionary, PrimaryKeyType, Primary, IPrimaryKey, FilterQuery, IWrappedEntity, EntityName, EntityData,
-  AnyEntity, EntityProperty, EntityMetadata, QBFilterQuery, PopulateOptions, Populate,
+  AnyEntity, EntityProperty, EntityMetadata, QBFilterQuery, PopulateOptions, Populate, Loaded, New, LoadedReference, LoadedCollection,
 } from './typings';
 export * from './enums';
 export * from './exceptions';

packages/core/src/metadata/EntitySchema.ts

@@ -1,13 +1,12 @@
-import { AnyEntity, Constructor, Dictionary, EntityMetadata, EntityName, EntityProperty, NonFunctionPropertyNames } from '../typings';
+import { AnyEntity, Constructor, Dictionary, EntityMetadata, EntityName, EntityProperty, ExpandProperty, NonFunctionPropertyNames } from '../typings';
 import {
   EmbeddedOptions, EnumOptions, IndexOptions, ManyToManyOptions, ManyToOneOptions, OneToManyOptions, OneToOneOptions, PrimaryKeyOptions, PropertyOptions,
   SerializedPrimaryKeyOptions, UniqueOptions,
 } from '../decorators';
-import { BaseEntity, Cascade, Collection, EntityRepository, ReferenceType, LoadStrategy } from '../entity';
+import { BaseEntity, Cascade, EntityRepository, ReferenceType, LoadStrategy } from '../entity';
 import { Type } from '../types';
 import { Utils } from '../utils';
 
-type CollectionItem<T> = T extends Collection<infer K> ? K : T;
 type TypeType = string | NumberConstructor | StringConstructor | BooleanConstructor | DateConstructor | ArrayConstructor | Constructor<Type<any>>;
 type TypeDef<T> = { type: TypeType } | { customType: Type<any> } | { entity: string | (() => string | EntityName<T>) };
 type Property<T, O> =
@@ -22,7 +21,7 @@ type PropertyKey<T, U> = NonFunctionPropertyNames<Omit<T, keyof U>>;
 type Metadata<T, U> =
   & Omit<Partial<EntityMetadata<T>>, 'name' | 'properties'>
   & ({ name: string } | { class: Constructor<T>; name?: string })
-  & { properties?: { [K in PropertyKey<T, U> & string]-?: Property<CollectionItem<NonNullable<T[K]>>, T> } };
+  & { properties?: { [K in PropertyKey<T, U> & string]-?: Property<ExpandProperty<NonNullable<T[K]>>, T> } };
 
 export class EntitySchema<T extends AnyEntity<T> = AnyEntity, U extends AnyEntity<U> | undefined = undefined> {
 

packages/core/src/typings.ts

@@ -67,17 +67,17 @@ export type Query<T> = true extends IsEntity<T>
 export type FilterQuery<T> = Query<T> | { [PrimaryKeyType]?: any };
 export type QBFilterQuery<T = any> = FilterQuery<T> & Dictionary | FilterQuery<T>;
 
-export interface IWrappedEntity<T, PK extends keyof T> {
+export interface IWrappedEntity<T, PK extends keyof T, P = never> {
   isInitialized(): boolean;
   populated(populated?: boolean): void;
   init(populated?: boolean, lockMode?: LockMode): Promise<T>;
-  toReference(): IdentifiedReference<T, PK>;
+  toReference<PK2 extends PK = never, P2 extends P = never>(): IdentifiedReference<T, PK2> & LoadedReference<T, P2>;
   toObject(ignoreFields?: string[]): Dictionary;
   toJSON(...args: any[]): Dictionary;
   assign(data: any, options?: AssignOptions | boolean): T;
 }
 
-export interface IWrappedEntityInternal<T, PK extends keyof T> extends IWrappedEntity<T, PK> {
+export interface IWrappedEntityInternal<T, PK extends keyof T, P = keyof T> extends IWrappedEntity<T, PK, P> {
   __uuid: string;
   __meta: EntityMetadata<T>;
   __internal: { platform: Platform; metadata: MetadataStorage; validator: EntityValidator };
@@ -96,7 +96,7 @@ export type AnyEntity<T = any> = { [K in keyof T]?: T[K] } & { [PrimaryKeyType]?
 export type EntityClass<T extends AnyEntity<T>> = Function & { prototype: T };
 export type EntityClassGroup<T extends AnyEntity<T>> = { entity: EntityClass<T>; schema: EntityMetadata<T> | EntitySchema<T> };
 export type EntityName<T extends AnyEntity<T>> = string | EntityClass<T> | EntitySchema<T, any>;
-export type EntityData<T extends AnyEntity<T>> = { [K in keyof T]?: T[K] | Primary<T[K]> | EntityData<T[K]> | CollectionItem<T[K]>[] } & Dictionary;
+export type EntityData<T extends AnyEntity<T>, P extends Populate<T> = keyof T> = { [K in keyof T]?: T[K] | Primary<T[K]> | EntityData<T[K]> | CollectionItem<T[K]>[] } & Dictionary;
 
 export interface EntityProperty<T extends AnyEntity<T> = any> {
   name: string & keyof T;
@@ -240,15 +240,15 @@ export interface LoadedReference<T extends AnyEntity<T>, P = never> extends Refe
   get(): T & P;
 }
 
-export interface LoadedCollection<T extends AnyEntity<T>, U extends AnyEntity<U>, P = never> extends Collection<T, U> {
+export interface LoadedCollection<T extends AnyEntity<T>, P = never> extends Collection<T> {
   $: readonly (T & P)[];
   get(): readonly (T & P)[];
 }
 
 type MarkLoaded<T extends AnyEntity<T>, P, H = unknown> = P extends Reference<infer U>
   ? LoadedReference<U, Loaded<U, H>>
   : P extends Collection<infer U>
-    ? LoadedCollection<U, T, Loaded<U, H>>
+    ? LoadedCollection<U, Loaded<U, H>>
     : P;
 
 type LoadedIfInKeyHint<T, K extends keyof T, H> = K extends H ? MarkLoaded<T, T[K]> : T[K];
@@ -271,3 +271,5 @@ export type Loaded<T, P = unknown> = unknown extends P ? T : T & {
       ? LoadedIfInNestedHint<T, K, P>
       : LoadedIfInKeyHint<T, K, P>;
 };
+
+export type New<T, P = string[]> = Loaded<T, P>;

tests/EntityHelper.mongo.test.ts

@@ -165,6 +165,11 @@ describe('EntityAssignerMongo', () => {
     expect(author._id).toBeNull();
   });
 
+  test('wrap helper returns the argument when its falsy', async () => {
+    expect(wrap(null)).toBeNull();
+    expect(wrap(undefined)).toBeUndefined();
+  });
+
   test('setting m:1 reference is propagated to 1:m collection', async () => {
     const author = new Author('n', 'e');
     const book = new Book('t');

tests/EntityManager.mongo.test.ts

@@ -978,15 +978,15 @@ describe('EntityManagerMongo', () => {
     orm.em.clear();
 
     const repo = orm.em.getRepository(Book);
-    let books = await repo.findAll(['author', 'tags']);
+    const books = await repo.findAll(['author', 'tags']);
     expect(books.length).toBe(3);
     expect(books[0].tags.count()).toBe(2);
     await books[0].author.books.init();
     await orm.em.remove(books[0].author).flush();
     orm.em.clear();
 
-    books = await repo.findAll();
-    expect(books.length).toBe(0);
+    const books2 = await repo.findAll();
+    expect(books2.length).toBe(0);
   });
 
   test('cascade remove on m:1 reference', async () => {

tests/EntityManager.mongo2.test.ts

@@ -33,6 +33,12 @@ describe('EntityManagerMongo2', () => {
     const book5 = await orm.em.findOneOrFail(Book, bible, { populate: { publisher: true, tags: true, perex: true } });
     expect(book5.publisher!.$.name).toBe('Publisher 123');
     expect(book5.tags.$[0].name).toBe('t1');
+
+    const pub2 = orm.em.create(Publisher, { name: 'asd' });
+    const wrapped0 = wrap(pub2).toReference<'id' | '_id'>();
+    // @ts-expect-error
+    expect(wrapped0.books).toBeUndefined();
+    book5.publisher = wrapped0;
   });
 
   afterAll(async () => orm.close(true));