[permissions] Override repository and manager methods #2 (#11929)

Closes https://github.com/twentyhq/core-team-issues/issues/747

packages/twenty-server/src/engine/twenty-orm/entity-manager/workspace-entity-manager.spec.ts

@@ -0,0 +1,255 @@
+import { ObjectRecordsPermissions } from 'twenty-shared/types';
+import { EntityManager } from 'typeorm';
+
+import { WorkspaceInternalContext } from 'src/engine/twenty-orm/interfaces/workspace-internal-context.interface';
+
+import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
+import { WorkspaceDataSource } from 'src/engine/twenty-orm/datasource/workspace.datasource';
+import { validateOperationIsPermittedOrThrow } from 'src/engine/twenty-orm/repository/permissions.utils';
+
+import { WorkspaceEntityManager } from './workspace-entity-manager';
+
+jest.mock('src/engine/twenty-orm/repository/permissions.utils', () => ({
+  validateOperationIsPermittedOrThrow: jest.fn(),
+}));
+
+jest.mock('../repository/workspace-select-query-builder', () => ({
+  WorkspaceSelectQueryBuilder: jest.fn().mockImplementation(() => ({
+    where: jest.fn().mockReturnThis(),
+    getMany: jest.fn().mockResolvedValue([]),
+    getOne: jest.fn().mockResolvedValue(null),
+    getManyAndCount: jest.fn().mockResolvedValue([[], 0]),
+    execute: jest
+      .fn()
+      .mockResolvedValue({ affected: 1, raw: [], generatedMaps: [] }),
+    setFindOptions: jest.fn().mockReturnThis(),
+  })),
+}));
+
+describe('WorkspaceEntityManager', () => {
+  let entityManager: WorkspaceEntityManager;
+  let mockInternalContext: WorkspaceInternalContext;
+  let mockDataSource: WorkspaceDataSource;
+  let mockPermissionOptions: {
+    shouldBypassPermissionChecks: boolean;
+    objectRecordsPermissions?: ObjectRecordsPermissions;
+  };
+
+  beforeEach(() => {
+    mockInternalContext = {
+      workspaceId: 'test-workspace-id',
+      objectMetadataMaps: {
+        idByNameSingular: {},
+      },
+      featureFlagsMap: {
+        [FeatureFlagKey.IsPermissionsV2Enabled]: true,
+      },
+    } as WorkspaceInternalContext;
+
+    mockDataSource = {
+      featureFlagMap: {
+        [FeatureFlagKey.IsPermissionsV2Enabled]: true,
+      },
+      permissionsPerRoleId: {},
+    } as WorkspaceDataSource;
+
+    mockPermissionOptions = {
+      shouldBypassPermissionChecks: false,
+      objectRecordsPermissions: {
+        'test-entity': {
+          canRead: true,
+          canUpdate: false,
+          canSoftDelete: false,
+          canDestroy: false,
+        },
+      },
+    };
+
+    // Mock TypeORM connection methods
+    const mockWorkspaceDataSource = {
+      getMetadata: jest.fn().mockReturnValue({
+        name: 'test-entity',
+        columns: [],
+        relations: [],
+        findInheritanceMetadata: jest.fn(),
+        findColumnWithPropertyPath: jest.fn(),
+      }),
+      createQueryBuilder: jest.fn().mockReturnValue({
+        delete: jest.fn().mockReturnThis(),
+        from: jest.fn().mockReturnThis(),
+        where: jest.fn().mockReturnThis(),
+        execute: jest
+          .fn()
+          .mockResolvedValue({ affected: 1, raw: [], generatedMaps: [] }),
+        getMany: jest.fn().mockResolvedValue([]),
+        getOne: jest.fn().mockResolvedValue(null),
+        getManyAndCount: jest.fn().mockResolvedValue([[], 0]),
+        update: jest.fn().mockReturnThis(),
+        softDelete: jest.fn().mockReturnThis(),
+        restore: jest.fn().mockReturnThis(),
+      }),
+      createQueryRunner: jest.fn().mockReturnValue({
+        connect: jest.fn(),
+        startTransaction: jest.fn(),
+        commitTransaction: jest.fn(),
+        rollbackTransaction: jest.fn(),
+        release: jest.fn(),
+        clearTable: jest.fn(),
+      }),
+    };
+
+    entityManager = new WorkspaceEntityManager(
+      mockInternalContext,
+      mockDataSource,
+    );
+
+    Object.defineProperty(entityManager, 'connection', {
+      get: () => mockWorkspaceDataSource,
+    });
+
+    jest.spyOn(entityManager as any, 'validatePermissions');
+    jest.spyOn(entityManager as any, 'createQueryBuilder');
+
+    jest
+      .spyOn(entityManager as any, 'extractTargetNameSingularFromEntityTarget')
+      .mockImplementation((entityName: string) => {
+        return entityName;
+      });
+
+    // Mock getFeatureFlagMap
+    jest.spyOn(entityManager as any, 'getFeatureFlagMap').mockReturnValue({
+      [FeatureFlagKey.IsPermissionsV2Enabled]: true,
+    });
+
+    // Mock typeORM's EntityManager methods
+    jest
+      .spyOn(EntityManager.prototype, 'save')
+      .mockImplementation(() => Promise.resolve({}));
+    jest
+      .spyOn(EntityManager.prototype, 'update')
+      .mockImplementation(() =>
+        Promise.resolve({ affected: 1, raw: [], generatedMaps: [] }),
+      );
+    jest
+      .spyOn(EntityManager.prototype, 'restore')
+      .mockImplementation(() =>
+        Promise.resolve({ affected: 1, raw: [], generatedMaps: [] }),
+      );
+    jest
+      .spyOn(EntityManager.prototype, 'clear')
+      .mockImplementation(() => Promise.resolve());
+    jest
+      .spyOn(EntityManager.prototype, 'preload')
+      .mockImplementation(() => Promise.resolve({}));
+
+    // Mock metadata methods
+    const mockMetadata = {
+      hasAllPrimaryKeys: jest.fn().mockReturnValue(true),
+      columns: [],
+      relations: [],
+      findInheritanceMetadata: jest.fn(),
+      findColumnWithPropertyPath: jest.fn(),
+    };
+
+    // Update mockWorkspaceDataSource to include metadata
+    mockWorkspaceDataSource.getMetadata = jest
+      .fn()
+      .mockReturnValue(mockMetadata);
+
+    // Reset the mock before each test
+    jest.clearAllMocks();
+  });
+
+  describe('Query Method', () => {
+    it('should call validatePermissions and validateOperationIsPermittedOrThrow for find', async () => {
+      await entityManager.find('test-entity', {}, mockPermissionOptions);
+
+      expect(entityManager.createQueryBuilder).toHaveBeenCalledWith(
+        'test-entity',
+        undefined,
+        undefined,
+        mockPermissionOptions,
+      );
+    });
+    it('should throw error for query', async () => {
+      expect(() => entityManager.query('SELECT * FROM test')).toThrow(
+        'Method not allowed.',
+      );
+    });
+  });
+
+  describe('Save Methods', () => {
+    it('should call validatePermissions and validateOperationIsPermittedOrThrow for save', async () => {
+      await entityManager.save(
+        'test-entity',
+        {},
+        { reload: false },
+        mockPermissionOptions,
+      );
+      expect(entityManager['validatePermissions']).toHaveBeenCalledWith(
+        'test-entity',
+        'update',
+        mockPermissionOptions,
+      );
+      expect(validateOperationIsPermittedOrThrow).toHaveBeenCalledWith({
+        entityName: 'test-entity',
+        operationType: 'update',
+        objectMetadataMaps: mockInternalContext.objectMetadataMaps,
+        objectRecordsPermissions:
+          mockPermissionOptions.objectRecordsPermissions,
+      });
+    });
+  });
+
+  describe('Update Methods', () => {
+    it('should call validatePermissions and validateOperationIsPermittedOrThrow for update', async () => {
+      await entityManager.update('test-entity', {}, {}, mockPermissionOptions);
+      expect(entityManager['validatePermissions']).toHaveBeenCalledWith(
+        'test-entity',
+        'update',
+        mockPermissionOptions,
+      );
+      expect(validateOperationIsPermittedOrThrow).toHaveBeenCalledWith({
+        entityName: 'test-entity',
+        operationType: 'update',
+        objectMetadataMaps: mockInternalContext.objectMetadataMaps,
+        objectRecordsPermissions:
+          mockPermissionOptions.objectRecordsPermissions,
+      });
+    });
+  });
+
+  describe('Other Methods', () => {
+    it('should call validatePermissions and validateOperationIsPermittedOrThrow for clear', async () => {
+      await entityManager.clear('test-entity', mockPermissionOptions);
+      expect(entityManager['validatePermissions']).toHaveBeenCalledWith(
+        'test-entity',
+        'delete',
+        mockPermissionOptions,
+      );
+      expect(validateOperationIsPermittedOrThrow).toHaveBeenCalledWith({
+        entityName: 'test-entity',
+        operationType: 'delete',
+        objectMetadataMaps: mockInternalContext.objectMetadataMaps,
+        objectRecordsPermissions:
+          mockPermissionOptions.objectRecordsPermissions,
+      });
+    });
+
+    it('should call validatePermissions and validateOperationIsPermittedOrThrow for preload', async () => {
+      await entityManager.preload('test-entity', {}, mockPermissionOptions);
+      expect(entityManager['validatePermissions']).toHaveBeenCalledWith(
+        'test-entity',
+        'select',
+        mockPermissionOptions,
+      );
+      expect(validateOperationIsPermittedOrThrow).toHaveBeenCalledWith({
+        entityName: 'test-entity',
+        operationType: 'select',
+        objectMetadataMaps: mockInternalContext.objectMetadataMaps,
+        objectRecordsPermissions:
+          mockPermissionOptions.objectRecordsPermissions,
+      });
+    });
+  });
+});

packages/twenty-server/src/engine/twenty-orm/entity-manager/workspace-entity-manager.ts

@@ -1,15 +1,28 @@
+import { Entity } from '@microsoft/microsoft-graph-types';
 import { ObjectRecordsPermissions } from 'twenty-shared/types';
 import { isDefined } from 'twenty-shared/utils';
 import {
+  DeleteResult,
   EntityManager,
   EntityTarget,
   FindManyOptions,
+  FindOneOptions,
+  FindOptionsWhere,
   InsertResult,
+  ObjectId,
   ObjectLiteral,
   QueryRunner,
+  RemoveOptions,
   Repository,
+  SaveOptions,
   SelectQueryBuilder,
+  TypeORMError,
+  UpdateResult,
 } from 'typeorm';
+import { DeepPartial } from 'typeorm/common/DeepPartial';
+import { PickKeysByType } from 'typeorm/common/PickKeysByType';
+import { EntityNotFoundError } from 'typeorm/error/EntityNotFoundError';
+import { FindOptionsUtils } from 'typeorm/find-options/FindOptionsUtils';
 import { QueryDeepPartialEntity } from 'typeorm/query-builder/QueryPartialEntity';
 import { UpsertOptions } from 'typeorm/repository/UpsertOptions';
 
@@ -29,6 +42,11 @@ import {
 import { WorkspaceSelectQueryBuilder } from 'src/engine/twenty-orm/repository/workspace-select-query-builder';
 import { WorkspaceRepository } from 'src/engine/twenty-orm/repository/workspace.repository';
 
+type PermissionOptions = {
+  shouldBypassPermissionChecks?: boolean;
+  objectRecordsPermissions?: ObjectRecordsPermissions;
+};
+
 export class WorkspaceEntityManager extends EntityManager {
   private readonly internalContext: WorkspaceInternalContext;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -117,10 +135,11 @@ export class WorkspaceEntityManager extends EntityManager {
     alias?: string,
     queryRunner?: QueryRunner,
     options: {
-      shouldBypassPermissionChecks: boolean;
-      roleId?: string;
+      shouldBypassPermissionChecks?: boolean;
+      objectRecordsPermissions?: ObjectRecordsPermissions;
     } = {
       shouldBypassPermissionChecks: false,
+      objectRecordsPermissions: {},
     },
   ): SelectQueryBuilder<Entity> | WorkspaceSelectQueryBuilder<Entity> {
     let queryBuilder: SelectQueryBuilder<Entity>;
@@ -145,50 +164,23 @@ export class WorkspaceEntityManager extends EntityManager {
     if (!isPermissionsV2Enabled) {
       return queryBuilder;
     } else {
-      let objectPermissions = {};
-
-      if (options?.roleId) {
-        const dataSource = this.connection as WorkspaceDataSource;
-        const objectPermissionsByRoleId = dataSource.permissionsPerRoleId;
-
-        objectPermissions = objectPermissionsByRoleId?.[options.roleId] ?? {};
-      }
-
       return new WorkspaceSelectQueryBuilder(
         queryBuilder,
-        objectPermissions,
+        options?.objectRecordsPermissions ?? {},
         this.internalContext,
         options?.shouldBypassPermissionChecks ?? false,
       );
     }
   }
 
-  override find<Entity extends ObjectLiteral>(
-    target: EntityTarget<Entity>,
-    options?: FindManyOptions<Entity>,
-    permissionOptions?: {
-      shouldBypassPermissionChecks?: boolean;
-      objectRecordsPermissions?: ObjectRecordsPermissions;
-    },
-  ): Promise<Entity[]> {
-    this.validatePermissions(target, 'select', permissionOptions);
-
-    return super.find(target, options);
-  }
-
   override insert<Entity extends ObjectLiteral>(
     target: EntityTarget<Entity>,
-    entityOrEntities:
-      | QueryDeepPartialEntity<Entity>
-      | QueryDeepPartialEntity<Entity>[],
-    permissionOptions?: {
-      shouldBypassPermissionChecks?: boolean;
-      objectRecordsPermissions?: ObjectRecordsPermissions;
-    },
+    entity: QueryDeepPartialEntity<Entity> | QueryDeepPartialEntity<Entity>[],
+    permissionOptions?: PermissionOptions,
   ): Promise<InsertResult> {
     this.validatePermissions(target, 'insert', permissionOptions);
 
-    return super.insert(target, entityOrEntities);
+    return super.insert(target, entity);
   }
 
   override upsert<Entity extends ObjectLiteral>(
@@ -207,6 +199,141 @@ export class WorkspaceEntityManager extends EntityManager {
     return super.upsert(target, entityOrEntities, conflictPathsOrOptions);
   }
 
+  override update<Entity extends ObjectLiteral>(
+    target: EntityTarget<Entity>,
+    criteria:
+      | string
+      | string[]
+      | number
+      | number[]
+      | Date
+      | Date[]
+      | ObjectId
+      | ObjectId[]
+      | unknown,
+    partialEntity: QueryDeepPartialEntity<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<UpdateResult> {
+    this.validatePermissions(target, 'update', permissionOptions);
+
+    return super.update(target, criteria, partialEntity);
+  }
+
+  override save<Entity extends ObjectLiteral>(
+    entities: Entity[],
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity[]>;
+
+  override save<Entity extends ObjectLiteral>(
+    entity: Entity,
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity>;
+
+  override save<Entity, T extends DeepPartial<Entity>>(
+    targetOrEntity: EntityTarget<Entity>,
+    entities: T[],
+    options: SaveOptions & {
+      reload: false;
+    },
+    permissionOptions?: PermissionOptions,
+  ): Promise<T[]>;
+
+  override save<Entity, T extends DeepPartial<Entity>>(
+    targetOrEntity: EntityTarget<Entity>,
+    entities: T[],
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<(T & Entity)[]>;
+
+  override save<Entity, T extends DeepPartial<Entity>>(
+    targetOrEntity: EntityTarget<Entity>,
+    entity: T,
+    options: SaveOptions & {
+      reload: false;
+    },
+    permissionOptions?: PermissionOptions,
+  ): Promise<T>;
+
+  override save<Entity, T extends DeepPartial<Entity>>(
+    targetOrEntity: EntityTarget<Entity>,
+    entity: T,
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<T>;
+
+  override save<Entity extends ObjectLiteral, T extends DeepPartial<Entity>>(
+    targetOrEntity: EntityTarget<Entity> | Entity | Entity[],
+    entityOrMaybeOptions:
+      | T
+      | T[]
+      | SaveOptions
+      | (SaveOptions & { reload: false }),
+    maybeOptionsOrMaybePermissionOptions?:
+      | PermissionOptions
+      | SaveOptions
+      | (SaveOptions & { reload: false }),
+    permissionOptions?: PermissionOptions,
+  ): Promise<(T & Entity) | (T & Entity)[] | Entity | Entity[]> {
+    const permissionOptionsFromArgs =
+      maybeOptionsOrMaybePermissionOptions &&
+      ('shouldBypassPermissionChecks' in maybeOptionsOrMaybePermissionOptions ||
+        'objectRecordsPermissions' in maybeOptionsOrMaybePermissionOptions)
+        ? maybeOptionsOrMaybePermissionOptions
+        : permissionOptions;
+
+    this.validatePermissions(
+      targetOrEntity,
+      'update',
+      permissionOptionsFromArgs,
+    );
+
+    const target =
+      arguments.length > 1 &&
+      (typeof targetOrEntity === 'function' ||
+        typeof targetOrEntity === 'string')
+        ? (targetOrEntity as EntityTarget<Entity>)
+        : undefined;
+
+    const entityOrEntities = target
+      ? (entityOrMaybeOptions as T | T[])
+      : (targetOrEntity as Entity | Entity[]);
+
+    const options = target
+      ? (maybeOptionsOrMaybePermissionOptions as SaveOptions | undefined)
+      : (entityOrMaybeOptions as SaveOptions | undefined);
+
+    if (isDefined(target)) {
+      let entity: T | undefined;
+      let entities: T[] | undefined;
+
+      if (Array.isArray(entityOrEntities)) {
+        entities = entityOrEntities as T[];
+
+        return super.save(target, entities, options);
+      } else {
+        entity = entityOrEntities as T;
+
+        return super.save(target, entity, options);
+      }
+    } else {
+      return super.save(entityOrEntities as Entity | Entity[], options);
+    }
+  }
+
+  override increment<Entity extends ObjectLiteral>(
+    target: EntityTarget<Entity>,
+    criteria: unknown,
+    propertyPath: string,
+    value: number | string,
+    permissionOptions?: PermissionOptions,
+  ): Promise<UpdateResult> {
+    this.validatePermissions(target, 'update', permissionOptions);
+
+    return super.increment(target, criteria, propertyPath, value);
+  }
+
   private getRepositoryKey({
     target,
     dataSource,
@@ -214,7 +341,7 @@ export class WorkspaceEntityManager extends EntityManager {
     shouldBypassPermissionChecks,
   }: {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    target: EntityTarget<any>;
+    target: EntityTarget<unknown>;
     dataSource: WorkspaceDataSource;
     shouldBypassPermissionChecks: boolean;
     roleId?: string;
@@ -233,8 +360,8 @@ export class WorkspaceEntityManager extends EntityManager {
       : `${repositoryPrefix}${roleIdSuffix}${rolesPermissionsVersionSuffix}${featureFlagMapVersionSuffix}`;
   }
 
-  private validatePermissions<Entity extends ObjectLiteral>(
-    target: EntityTarget<Entity>,
+  validatePermissions<Entity extends ObjectLiteral>(
+    target: EntityTarget<Entity> | Entity,
     operationType: OperationType,
     permissionOptions?: {
       shouldBypassPermissionChecks?: boolean;
@@ -254,8 +381,13 @@ export class WorkspaceEntityManager extends EntityManager {
       return;
     }
 
+    const entityName =
+      typeof target === 'function' || typeof target === 'string'
+        ? this.extractTargetNameSingularFromEntityTarget(target)
+        : this.extractTargetNameSingularFromEntity(target);
+
     validateOperationIsPermittedOrThrow({
-      entityName: this.extractTargetNameSingularFromEntityTarget(target),
+      entityName,
       operationType,
       objectRecordsPermissions:
         permissionOptions?.objectRecordsPermissions ?? {},
@@ -264,9 +396,706 @@ export class WorkspaceEntityManager extends EntityManager {
   }
 
   private extractTargetNameSingularFromEntityTarget(
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    target: EntityTarget<any>,
+    target: EntityTarget<unknown>,
   ): string {
     return this.connection.getMetadata(target).name;
   }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  private extractTargetNameSingularFromEntity(entity: any): string {
+    return this.connection.getMetadata(entity.constructor).name;
+  }
+
+  // Forbidden methods
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  override query<T = any>(_query: string, _parameters?: any[]): Promise<T> {
+    throw new Error('Method not allowed.');
+  }
+
+  // Not in use methods - duplicated from TypeORM's EntityManager to use our createQueryBuilder
+
+  override find<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    options?: FindManyOptions<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity[]> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      FindOptionsUtils.extractFindManyOptionsAlias(options) || metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions(options || {})
+      .getMany();
+  }
+
+  override findBy<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    where: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity[]> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions({ where: where })
+      .getMany();
+  }
+
+  override findOne<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    options: FindOneOptions<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity | null> {
+    const metadata = this.connection.getMetadata(entityClass);
+    // prepare alias for built query
+    let alias = metadata.name;
+
+    if (options && options.join) {
+      alias = options.join.alias;
+    }
+    if (!options.where) {
+      throw new Error(
+        `You must provide selection conditions in order to find a single row.`,
+      );
+    }
+
+    // create query builder and apply find options
+    return this.createQueryBuilder(
+      entityClass,
+      alias,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions({
+        ...options,
+        take: 1,
+      })
+      .getOne();
+  }
+
+  override findOneBy<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    where: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity | null> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    // create query builder and apply find options
+    return this.createQueryBuilder(
+      entityClass,
+      metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions({
+        where,
+        take: 1,
+      })
+      .getOne();
+  }
+
+  override findAndCount<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    options?: FindManyOptions<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<[Entity[], number]> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      FindOptionsUtils.extractFindManyOptionsAlias(options) || metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions(options || {})
+      .getManyAndCount();
+  }
+
+  override findAndCountBy<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    where: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<[Entity[], number]> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions({ where })
+      .getManyAndCount();
+  }
+
+  override findOneOrFail<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    options: FindOneOptions<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity> {
+    return this.findOne(entityClass, options, permissionOptions).then(
+      (value) => {
+        if (value === null) {
+          return Promise.reject(new EntityNotFoundError(entityClass, options));
+        }
+
+        return Promise.resolve(value);
+      },
+    );
+  }
+
+  override findOneByOrFail<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    where: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity> {
+    return this.findOneBy(entityClass, where, permissionOptions).then(
+      (value) => {
+        if (value === null) {
+          return Promise.reject(new EntityNotFoundError(entityClass, where));
+        }
+
+        return Promise.resolve(value);
+      },
+    );
+  }
+
+  override delete<Entity extends ObjectLiteral>(
+    targetOrEntity: EntityTarget<Entity>,
+    criteria: unknown,
+    permissionOptions?: PermissionOptions,
+  ): Promise<DeleteResult> {
+    this.validatePermissions(targetOrEntity, 'delete', permissionOptions);
+
+    return super.delete(targetOrEntity, criteria);
+  }
+
+  override remove<Entity>(
+    entity: Entity,
+    options?: RemoveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity>;
+
+  override remove<Entity>(
+    targetOrEntity: EntityTarget<Entity>,
+    entity: Entity,
+    options?: RemoveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity>;
+
+  override remove<Entity>(
+    entity: Entity[],
+    options?: RemoveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity>;
+
+  override remove<Entity>(
+    targetOrEntity: EntityTarget<Entity>,
+    entity: Entity[],
+    options?: RemoveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity[]>;
+
+  override remove<Entity extends ObjectLiteral>(
+    targetOrEntity: EntityTarget<Entity> | Entity[] | Entity,
+    entityOrMaybeOptions: Entity | Entity[] | RemoveOptions,
+    maybeOptionsOrMaybePermissionOptions?: RemoveOptions | PermissionOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity | Entity[]> {
+    const permissionOptionsFromArgs =
+      maybeOptionsOrMaybePermissionOptions &&
+      ('shouldBypassPermissionChecks' in maybeOptionsOrMaybePermissionOptions ||
+        'objectRecordsPermissions' in maybeOptionsOrMaybePermissionOptions)
+        ? (maybeOptionsOrMaybePermissionOptions as PermissionOptions)
+        : permissionOptions;
+
+    this.validatePermissions(
+      targetOrEntity,
+      'delete',
+      permissionOptionsFromArgs,
+    );
+
+    const target =
+      typeof targetOrEntity === 'function' || typeof targetOrEntity === 'string'
+        ? (targetOrEntity as EntityTarget<Entity>)
+        : undefined;
+
+    const entityOrEntities = target
+      ? (entityOrMaybeOptions as Entity | Entity[])
+      : (targetOrEntity as Entity | Entity[]);
+
+    const options = target
+      ? (maybeOptionsOrMaybePermissionOptions as RemoveOptions | undefined)
+      : (entityOrMaybeOptions as RemoveOptions);
+
+    if (isDefined(target)) {
+      if (Array.isArray(entityOrEntities)) {
+        return super.remove(target, entityOrEntities as Entity[], options);
+      } else {
+        return super.remove(target, entityOrEntities as Entity, options);
+      }
+    } else {
+      return super.remove(entityOrEntities as Entity | Entity[], options);
+    }
+  }
+
+  override softDelete<Entity extends ObjectLiteral>(
+    targetOrEntity: EntityTarget<Entity>,
+    criteria: unknown,
+    permissionOptions?: PermissionOptions,
+  ): Promise<UpdateResult> {
+    // if user passed empty criteria or empty list of criterias, then throw an error
+    if (
+      criteria === undefined ||
+      criteria === null ||
+      criteria === '' ||
+      (Array.isArray(criteria) && criteria.length === 0)
+    ) {
+      return Promise.reject(
+        new TypeORMError(
+          `Empty criteria(s) are not allowed for the softDelete method.`,
+        ),
+      );
+    }
+    if (
+      typeof criteria === 'string' ||
+      typeof criteria === 'number' ||
+      criteria instanceof Date ||
+      Array.isArray(criteria)
+    ) {
+      return this.createQueryBuilder(
+        undefined,
+        undefined,
+        this.queryRunner,
+        permissionOptions,
+      )
+        .softDelete()
+        .from(targetOrEntity)
+        .whereInIds(criteria)
+        .execute();
+    } else {
+      return this.createQueryBuilder(
+        undefined,
+        undefined,
+        this.queryRunner,
+        permissionOptions,
+      )
+        .softDelete()
+        .from(targetOrEntity)
+        .where(criteria)
+        .execute();
+    }
+  }
+
+  override softRemove<Entity extends ObjectLiteral>(
+    entities: Entity[],
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity[]>;
+
+  override softRemove<Entity extends ObjectLiteral>(
+    entities: Entity,
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity>;
+
+  override softRemove<
+    Entity extends ObjectLiteral,
+    T extends DeepPartial<Entity>,
+  >(
+    targetOrEntity: EntityTarget<Entity>,
+    entities: T[],
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<T[]>;
+
+  override softRemove<
+    Entity extends ObjectLiteral,
+    T extends DeepPartial<Entity>,
+  >(
+    targetOrEntity: EntityTarget<Entity>,
+    entities: T,
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<T>;
+
+  override async softRemove<
+    Entity extends ObjectLiteral,
+    T extends DeepPartial<Entity>,
+  >(
+    targetOrEntityOrEntities: Entity | Entity[] | EntityTarget<Entity>,
+    entitiesOrMaybeOptions: T | T[] | SaveOptions,
+    maybeOptionsOrMaybePermissionOptions?: SaveOptions | PermissionOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity | Entity[] | T | T[]> {
+    const permissionOptionsFromArgs =
+      maybeOptionsOrMaybePermissionOptions &&
+      ('shouldBypassPermissionChecks' in maybeOptionsOrMaybePermissionOptions ||
+        'objectRecordsPermissions' in maybeOptionsOrMaybePermissionOptions)
+        ? (maybeOptionsOrMaybePermissionOptions as PermissionOptions)
+        : permissionOptions;
+
+    this.validatePermissions(
+      targetOrEntityOrEntities,
+      'soft-delete',
+      permissionOptionsFromArgs,
+    );
+
+    const target =
+      typeof targetOrEntityOrEntities === 'function' ||
+      typeof targetOrEntityOrEntities === 'string'
+        ? (targetOrEntityOrEntities as EntityTarget<Entity>)
+        : undefined;
+
+    const entityOrEntities = target
+      ? (entitiesOrMaybeOptions as T | T[])
+      : (targetOrEntityOrEntities as Entity | Entity[]);
+
+    const options = target
+      ? (maybeOptionsOrMaybePermissionOptions as SaveOptions | undefined)
+      : (entitiesOrMaybeOptions as SaveOptions);
+
+    if (isDefined(target)) {
+      if (Array.isArray(entityOrEntities)) {
+        return super.softRemove(target, entityOrEntities as T[], options);
+      } else {
+        return super.softRemove(target, entityOrEntities as T, options);
+      }
+    } else {
+      if (Array.isArray(entityOrEntities)) {
+        return super.softRemove(entityOrEntities as Entity | Entity[], options);
+      } else {
+        return super.softRemove(entityOrEntities as Entity, options);
+      }
+    }
+  }
+
+  override recover<Entity>(
+    entities: Entity[],
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity[]>;
+
+  override recover<Entity>(
+    entity: Entity,
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity>;
+
+  override recover<Entity, T extends DeepPartial<Entity>>(
+    targetOrEntity: EntityTarget<Entity>,
+    entities: T[],
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<T[]>;
+
+  override recover<Entity, T extends DeepPartial<Entity>>(
+    targetOrEntity: EntityTarget<Entity>,
+    entity: T,
+    options?: SaveOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<T>;
+
+  override recover<Entity extends ObjectLiteral, T extends DeepPartial<Entity>>(
+    targetOrEntityOrEntities: EntityTarget<Entity> | Entity | Entity[],
+    entityOrEntitiesOrMaybeOptions: T | T[] | SaveOptions,
+    maybeOptionsOrMaybePermissionOptions?: SaveOptions | PermissionOptions,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity | Entity[] | T | T[]> {
+    const permissionOptionsFromArgs =
+      maybeOptionsOrMaybePermissionOptions &&
+      ('shouldBypassPermissionChecks' in maybeOptionsOrMaybePermissionOptions ||
+        'objectRecordsPermissions' in maybeOptionsOrMaybePermissionOptions)
+        ? (maybeOptionsOrMaybePermissionOptions as PermissionOptions)
+        : permissionOptions;
+
+    this.validatePermissions(
+      targetOrEntityOrEntities,
+      'restore',
+      permissionOptionsFromArgs,
+    );
+
+    const target =
+      typeof targetOrEntityOrEntities === 'function' ||
+      typeof targetOrEntityOrEntities === 'string'
+        ? (targetOrEntityOrEntities as EntityTarget<Entity>)
+        : undefined;
+
+    const options = target
+      ? (maybeOptionsOrMaybePermissionOptions as SaveOptions | undefined)
+      : (entityOrEntitiesOrMaybeOptions as SaveOptions);
+
+    if (target) {
+      if (Array.isArray(entityOrEntitiesOrMaybeOptions)) {
+        return super.recover(
+          target,
+          entityOrEntitiesOrMaybeOptions as T[],
+          options,
+        );
+      } else {
+        return super.recover(
+          target,
+          entityOrEntitiesOrMaybeOptions as T,
+          options,
+        );
+      }
+    } else {
+      if (Array.isArray(entityOrEntitiesOrMaybeOptions)) {
+        return super.recover(
+          entityOrEntitiesOrMaybeOptions as Entity | Entity[],
+          options,
+        );
+      } else {
+        return super.recover(entityOrEntitiesOrMaybeOptions as Entity, options);
+      }
+    }
+  }
+
+  override restore<Entity extends ObjectLiteral>(
+    targetOrEntity: EntityTarget<Entity>,
+    criteria: unknown,
+    permissionOptions?: PermissionOptions,
+  ): Promise<UpdateResult> {
+    // if user passed empty criteria or empty list of criterias, then throw an error
+    if (
+      criteria === undefined ||
+      criteria === null ||
+      criteria === '' ||
+      (Array.isArray(criteria) && criteria.length === 0)
+    ) {
+      return Promise.reject(
+        new TypeORMError(
+          `Empty criteria(s) are not allowed for the restore method.`,
+        ),
+      );
+    }
+    if (
+      typeof criteria === 'string' ||
+      typeof criteria === 'number' ||
+      criteria instanceof Date ||
+      Array.isArray(criteria)
+    ) {
+      return this.createQueryBuilder(
+        undefined,
+        undefined,
+        this.queryRunner,
+        permissionOptions,
+      )
+        .restore()
+        .from(targetOrEntity)
+        .whereInIds(criteria)
+        .execute();
+    } else {
+      return this.createQueryBuilder(
+        undefined,
+        undefined,
+        this.queryRunner,
+        permissionOptions,
+      )
+        .restore()
+        .from(targetOrEntity)
+        .where(criteria)
+        .execute();
+    }
+  }
+
+  override exists<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    options?: FindManyOptions<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<boolean> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      FindOptionsUtils.extractFindManyOptionsAlias(options) || metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions(options || {})
+      .getExists();
+  }
+
+  override existsBy<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    where: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<boolean> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions({ where })
+      .getExists();
+  }
+
+  override count<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    options?: FindManyOptions<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<number> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      FindOptionsUtils.extractFindManyOptionsAlias(options) || metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions(options || {})
+      .getCount();
+  }
+
+  override countBy<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    where: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<number> {
+    const metadata = this.connection.getMetadata(entityClass);
+
+    return this.createQueryBuilder(
+      entityClass,
+      metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions({ where })
+      .getCount();
+  }
+
+  async callAggregateFunCustom(
+    entityClass: EntityTarget<Entity>,
+    fnName: string,
+    columnName: string,
+    where = {},
+    permissionOptions?: PermissionOptions,
+  ) {
+    const metadata = this.connection.getMetadata(entityClass);
+    const column = metadata.columns.find(
+      (item) => item.propertyPath === columnName,
+    );
+
+    if (!column) {
+      throw new TypeORMError(
+        `Column "${columnName}" was not found in table "${metadata.name}"`,
+      );
+    }
+    const result = await this.createQueryBuilder(
+      entityClass,
+      metadata.name,
+      this.queryRunner,
+      permissionOptions,
+    )
+      .setFindOptions({ where })
+      .select(
+        `${fnName}(${this.connection.driver.escape(column.databaseName)})`,
+        fnName,
+      )
+      .getRawOne();
+
+    return result[fnName] === null ? null : parseFloat(result[fnName]);
+  }
+
+  override sum<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    columnName: PickKeysByType<Entity, number>,
+    where?: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<number | null> {
+    return this.callAggregateFunCustom(
+      entityClass,
+      'SUM',
+      columnName,
+      where,
+      permissionOptions,
+    );
+  }
+
+  override average<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    columnName: PickKeysByType<Entity, number>,
+    where?: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<number | null> {
+    return this.callAggregateFunCustom(
+      entityClass,
+      'AVG',
+      columnName,
+      where,
+      permissionOptions,
+    );
+  }
+
+  override minimum<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    columnName: PickKeysByType<Entity, number>,
+    where?: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<number | null> {
+    return this.callAggregateFunCustom(
+      entityClass,
+      'MIN',
+      columnName,
+      where,
+      permissionOptions,
+    );
+  }
+
+  override maximum<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    columnName: PickKeysByType<Entity, number>,
+    where?: FindOptionsWhere<Entity> | FindOptionsWhere<Entity>[],
+    permissionOptions?: PermissionOptions,
+  ): Promise<number | null> {
+    return this.callAggregateFunCustom(
+      entityClass,
+      'MAX',
+      columnName,
+      where,
+      permissionOptions,
+    );
+  }
+
+  override clear<Entity>(
+    entityClass: EntityTarget<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<void> {
+    this.validatePermissions(entityClass, 'delete', permissionOptions);
+
+    return super.clear(entityClass);
+  }
+
+  override preload<Entity extends ObjectLiteral>(
+    entityClass: EntityTarget<Entity>,
+    entityLike: DeepPartial<Entity>,
+    permissionOptions?: PermissionOptions,
+  ): Promise<Entity | undefined> {
+    this.validatePermissions(entityClass, 'select', permissionOptions);
+
+    return super.preload(entityClass, entityLike);
+  }
+
+  override decrement<Entity extends ObjectLiteral>(
+    target: EntityTarget<Entity>,
+    criteria: unknown,
+    propertyPath: string,
+    value: number | string,
+    permissionOptions?: PermissionOptions,
+  ): Promise<UpdateResult> {
+    this.validatePermissions(target, 'update', permissionOptions);
+
+    return super.decrement(target, criteria, propertyPath, value);
+  }
 }