feat(core): allow setting default schema on EntityManager (#4717)

Hi,

I want to start with thanking you all for this great ORM that is so
actively maintained and offer so many great features. One thing that
really catches my eye is the per request (shell)EntityManager that I
think will make this ORM a winner when it comes to multi tenancy. Almost
all of the other big ORM libraries are based on a global concept.

My aim with this PR is to make a shoot at implementing schema based
multi tenancy support with help of the per request EntityManager. With
this PR it's possible to set a schema when forking the EntityManager.
That schema will be applied as a default schema if no other schema is
present. It makes use of the schema options introduced in:
#2296

docs/docs/multiple-schemas.md

@@ -32,6 +32,41 @@ const qb = em.createQueryBuilder(User);
 await qb.insert({ email: 'foo@bar.com' }).withSchema('client-123');
 ```
 
+## Default schema on `EntityManager`
+
+Instead of defining schema per entity or operation it's possible to `.fork()` EntityManger and define a default schema that will be used with wildcard schemas.
+
+```ts
+const fork = em.fork({ schema: 'client-123' });
+await fork.findOne(User, { ... });
+
+// Will yield the same result as 
+const user = await em.findOne(User, { ... }, { schema: 'client-123' });
+```
+
+When creating an entity the fork will set default schema
+
+```ts
+const fork = em.fork({ schema: 'client-123' });
+const user = new User();
+user.email = 'foo@bar.com';
+await fork.persistAndFlush(user);
+
+// Will yield the same result as
+const qb = em.createQueryBuilder(User);
+await qb.insert({ email: 'foo@bar.com' }).withSchema('client-123');
+```
+
+You can also set or clear schema
+
+```ts
+em.schema = 'client-123';
+const fork = em.fork({ schema: 'client-1234' });
+fork.schema = null;
+```
+
+`EntityManager.schema` Respects the context, so global EM will give you the contextual schema if executed inside [request context handler](https://mikro-orm.io/docs/identity-map#-requestcontext-helper)
+
 ## Wildcard Schema
 
 Since v5, MikroORM also supports defining entities that can exist in multiple schemas. To do that, we just specify wildcard schema:
@@ -57,7 +92,7 @@ export class Book {
 
 Entities like this will be by default ignored when using `SchemaGenerator`, as we need to specify which schema to use. For that we need to use the `schema` option of the `createSchema/updateSchema/dropSchema` methods or the `--schema` CLI parameter.
 
-On runtime, the wildcard schema will be replaced with either `FindOptions.schema`, or with the `schema` option from the ORM config.
+On runtime, the wildcard schema will be replaced with either `FindOptions.schema`, `EntityManager.schema` or with the `schema` option from the ORM config.
 
 ### Note about migrations
 

docs/versioned_docs/version-5.8/multiple-schemas.md

@@ -32,6 +32,41 @@ const qb = em.createQueryBuilder(User);
 await qb.insert({ email: 'foo@bar.com' }).withSchema('client-123');
 ```
 
+## Default schema on `EntityManager`
+
+Instead of defining schema per entity or operation it's possible to `.fork()` EntityManger and define a default schema that will be used with wildcard schemas.
+
+```ts
+const fork = em.fork({ schema: 'client-123' });
+await fork.findOne(User, { ... });
+
+// Will yield the same result as 
+const user = await em.findOne(User, { ... }, { schema: 'client-123' });
+```
+
+When creating an entity the fork will set default schema
+
+```ts
+const fork = em.fork({ schema: 'client-123' });
+const user = new User();
+user.email = 'foo@bar.com';
+await fork.persistAndFlush(user);
+
+// Will yield the same result as
+const qb = em.createQueryBuilder(User);
+await qb.insert({ email: 'foo@bar.com' }).withSchema('client-123');
+```
+
+You can also set or clear schema
+
+```ts
+em.schema = 'client-123';
+const fork = em.fork({ schema: 'client-1234' });
+fork.schema = null;
+```
+
+`EntityManager.schema` Respects the context, so global EM will give you the contextual schema if executed inside [request context handler](https://mikro-orm.io/docs/identity-map#-requestcontext-helper)
+
 ## Wildcard Schema
 
 Since v5, MikroORM also supports defining entities that can exist in multiple schemas. To do that, we just specify wildcard schema:
@@ -57,7 +92,7 @@ export class Book {
 
 Entities like this will be by default ignored when using `SchemaGenerator`, as we need to specify which schema to use. For that we need to use the `schema` option of the `createSchema/updateSchema/dropSchema` methods or the `--schema` CLI parameter.
 
-On runtime, the wildcard schema will be replaced with either `FindOptions.schema`, or with the `schema` option from the ORM config.
+On runtime, the wildcard schema will be replaced with either `FindOptions.schema`, `EntityManager.schema` or with the `schema` option from the ORM config.
 
 ### Note about migrations
 

packages/core/src/EntityManager.ts

@@ -90,6 +90,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
   private transactionContext?: Transaction;
   private disableTransactions = this.config.get('disableTransactions');
   private flushMode?: FlushMode;
+  private _schema?: string;
 
   /**
    * @internal
@@ -174,6 +175,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
     }
 
     const em = this.getContext();
+    options.schema ??= em._schema;
     await em.tryFlush(entityName, options);
     entityName = Utils.className(entityName);
     where = await em.processWhere(entityName, where, options, 'read') as FilterQuery<Entity>;
@@ -498,6 +500,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
     }
 
     const em = this.getContext();
+    options.schema ??= em._schema;
     await em.tryFlush(entityName, options);
     entityName = Utils.className(entityName);
     const meta = em.metadata.get<Entity>(entityName);
@@ -612,6 +615,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    */
   async upsert<Entity extends object>(entityNameOrEntity: EntityName<Entity> | Entity, data?: EntityData<Entity> | Entity, options: UpsertOptions<Entity> = {}): Promise<Entity> {
     const em = this.getContext(false);
+    options.schema ??= em._schema;
 
     let entityName: EntityName<Entity>;
     let where: FilterQuery<Entity>;
@@ -748,6 +752,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    */
   async upsertMany<Entity extends object>(entityNameOrEntity: EntityName<Entity> | Entity[], data?: (EntityData<Entity> | Entity)[], options: UpsertManyOptions<Entity> = {}): Promise<Entity[]> {
     const em = this.getContext(false);
+    options.schema ??= em._schema;
 
     let entityName: string;
     let propIndex: number;
@@ -1079,6 +1084,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    */
   async insert<Entity extends object>(entityNameOrEntity: EntityName<Entity> | Entity, data?: EntityData<Entity> | Entity, options: NativeInsertUpdateOptions<Entity> = {}): Promise<Primary<Entity>> {
     const em = this.getContext(false);
+    options.schema ??= em._schema;
 
     let entityName;
 
@@ -1113,6 +1119,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    */
   async insertMany<Entity extends object>(entityNameOrEntities: EntityName<Entity> | Entity[], data?: EntityData<Entity>[] | Entity[], options: NativeInsertUpdateOptions<Entity> = {}): Promise<Primary<Entity>[]> {
     const em = this.getContext(false);
+    options.schema ??= em._schema;
 
     let entityName;
 
@@ -1153,6 +1160,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    */
   async nativeUpdate<Entity extends object>(entityName: EntityName<Entity>, where: FilterQuery<Entity>, data: EntityData<Entity>, options: UpdateOptions<Entity> = {}): Promise<number> {
     const em = this.getContext(false);
+    options.schema ??= em._schema;
 
     entityName = Utils.className(entityName);
     data = QueryHelper.processObjectParams(data);
@@ -1169,6 +1177,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    */
   async nativeDelete<Entity extends object>(entityName: EntityName<Entity>, where: FilterQuery<Entity>, options: DeleteOptions<Entity> = {}): Promise<number> {
     const em = this.getContext(false);
+    options.schema ??= em._schema;
 
     entityName = Utils.className(entityName);
     where = await em.processWhere(entityName, where, options, 'delete');
@@ -1220,6 +1229,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
       return em.merge((entityName as Dictionary).constructor.name, entityName as unknown as EntityData<Entity>, data as MergeOptions);
     }
 
+    options.schema ??= em._schema;
     entityName = Utils.className(entityName as string);
     em.validator.validatePrimaryKey(data as EntityData<Entity>, em.metadata.get(entityName));
     let entity = em.unitOfWork.tryGetById<Entity>(entityName, data as FilterQuery<Entity>, options.schema, false);
@@ -1248,6 +1258,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    */
   create<Entity extends object>(entityName: EntityName<Entity>, data: RequiredEntityData<Entity>, options: CreateOptions = {}): Entity {
     const em = this.getContext();
+    options.schema ??= em._schema;
     const entity = em.entityFactory.create(entityName, data, {
       ...options,
       newEntity: !options.managed,
@@ -1293,6 +1304,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
    * Gets a reference to the entity identified by the given type and identifier without actually loading it, if the entity is not yet loaded
    */
   getReference<Entity extends object>(entityName: EntityName<Entity>, id: Primary<Entity>, options: GetReferenceOptions = {}): Entity | Reference<Entity> {
+    options.schema ??= this.schema;
     options.convertCustomTypes ??= false;
     const meta = this.metadata.get(Utils.className(entityName));
 
@@ -1320,8 +1332,13 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
     Entity extends object,
     Hint extends string = never,
   >(entityName: EntityName<Entity>, where: FilterQuery<Entity> = {} as FilterQuery<Entity>, options: CountOptions<Entity, Hint> = {}): Promise<number> {
-    options = { ...options };
     const em = this.getContext(false);
+
+    // Shallow copy options since the object will be modified when deleting orderBy
+    options = {
+      schema: em._schema,
+      ...options,
+    };
     entityName = Utils.className(entityName);
     where = await em.processWhere(entityName, where, options as FindOptions<Entity, Hint>, 'read') as FilterQuery<Entity>;
     options.populate = em.preparePopulate(entityName, options) as unknown as Populate<Entity>;
@@ -1506,6 +1523,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
     }
 
     const em = this.getContext();
+    options.schema ??= em._schema;
     const entityName = (entities[0] as Dictionary).constructor.name;
     const preparedPopulate = em.preparePopulate<Entity>(entityName, { populate: populate as true });
     await em.entityLoader.populate(entityName, entities, preparedPopulate, options);
@@ -1539,6 +1557,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
 
     fork.filters = { ...em.filters };
     fork.filterParams = Utils.copy(em.filterParams);
+    fork._schema = options.schema ?? em.schema;
 
     if (!options.clear) {
       for (const entity of em.unitOfWork.getIdentityMap()) {
@@ -1834,6 +1853,22 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
     await this.getContext().resultCache.remove(cacheKey);
   }
 
+  /**
+   * Returns the default schema of this EntityManager. Respects the context, so global EM will give you the contextual schema
+   * if executed inside request context handler.
+   */
+  get schema(): string | undefined {
+    return this.getContext(false)._schema;
+  }
+
+  /**
+   * Sets the default schema of this EntityManager. Respects the context, so global EM will set the contextual schema
+   * if executed inside request context handler.
+   */
+  set schema(schema: string | null | undefined) {
+    this.getContext(false)._schema = schema ?? undefined;
+  }
+
   /**
    * Returns the ID of this EntityManager. Respects the context, so global EM will give you the contextual ID
    * if executed inside request context handler.
@@ -1878,4 +1913,6 @@ export interface ForkOptions {
   flushMode?: FlushMode;
   /** disable transactions for this fork */
   disableTransactions?: boolean;
+  /** default schema to use for this fork */
+  schema?: string;
 }

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

@@ -559,6 +559,8 @@ export class UnitOfWork {
       return;
     }
 
+    // Set entityManager default schema
+    wrapped.__schema ??= this.em.schema;
     this.initIdentifier(entity);
 
     for (const prop of helper(entity).__meta.relations) {

packages/core/src/utils/RequestContext.ts

@@ -22,17 +22,17 @@ export class RequestContext {
   /**
    * Creates new RequestContext instance and runs the code inside its domain.
    */
-  static create(em: EntityManager | EntityManager[], next: (...args: any[]) => void): void {
-    const ctx = this.createContext(em);
+  static create(em: EntityManager | EntityManager[], next: (...args: any[]) => void, options: CreateContextOptions = {}): void {
+    const ctx = this.createContext(em, options);
     this.storage.run(ctx, next);
   }
 
   /**
    * Creates new RequestContext instance and runs the code inside its domain.
    * Async variant, when the `next` handler needs to be awaited (like in Koa).
    */
-  static async createAsync<T>(em: EntityManager | EntityManager[], next: (...args: any[]) => Promise<T>): Promise<T> {
-    const ctx = this.createContext(em);
+  static async createAsync<T>(em: EntityManager | EntityManager[], next: (...args: any[]) => Promise<T>, options: CreateContextOptions = {}): Promise<T> {
+    const ctx = this.createContext(em, options);
     return this.storage.run(ctx, next);
   }
 
@@ -51,16 +51,20 @@ export class RequestContext {
     return context ? context.map.get(name) : undefined;
   }
 
-  private static createContext(em: EntityManager | EntityManager[]): RequestContext {
+  private static createContext(em: EntityManager | EntityManager[], options: CreateContextOptions = {}): RequestContext {
     const forks = new Map<string, EntityManager>();
 
     if (Array.isArray(em)) {
-      em.forEach(em => forks.set(em.name, em.fork({ useContext: true })));
+      em.forEach(em => forks.set(em.name, em.fork({ useContext: true, schema: options.schema })));
     } else {
-      forks.set(em.name, em.fork({ useContext: true }));
+      forks.set(em.name, em.fork({ useContext: true, schema: options.schema }));
     }
 
     return new RequestContext(forks);
   }
 
 }
+
+export interface CreateContextOptions {
+  schema?: string;
+}

tests/EntityManager.mongo.test.ts

@@ -457,10 +457,10 @@ describe('EntityManagerMongo', () => {
 
     const spy = jest.spyOn(EntityManager.prototype, 'getContext');
     const fork2 = orm.em.fork({ disableContextResolution: true });
-    expect(spy).toBeCalledTimes(2);
+    expect(spy).toBeCalledTimes(3);
 
     const fork3 = orm.em.fork({ disableContextResolution: false });
-    expect(spy).toBeCalledTimes(5);
+    expect(spy).toBeCalledTimes(7);
   });
 
   test('findOne with empty where will throw', async () => {