import { type Entity } from '@microsoft/microsoft-graph-types'; import { isDefined } from 'class-validator'; import { type ObjectsPermissionsByRoleId } from 'twenty-shared/types'; import { DataSource, type DataSourceOptions, type EntityTarget, type ObjectLiteral, type QueryRunner, type ReplicationMode, type SelectQueryBuilder, } from 'typeorm'; import { EntityManagerFactory } from 'typeorm/entity-manager/EntityManagerFactory'; import { type FeatureFlagMap } from 'src/engine/core-modules/feature-flag/interfaces/feature-flag-map.interface'; import { type WorkspaceDataSourceInterface } from 'src/engine/twenty-orm/interfaces/workspace-datasource.interface'; import { type WorkspaceInternalContext } from 'src/engine/twenty-orm/interfaces/workspace-internal-context.interface'; import { type AuthContext } from 'src/engine/core-modules/auth/types/auth-context.type'; import { PermissionsException, PermissionsExceptionCode, } from 'src/engine/metadata-modules/permissions/permissions.exception'; import { WorkspaceEntityManager } from 'src/engine/twenty-orm/entity-manager/workspace-entity-manager'; import { type WorkspaceQueryRunner } from 'src/engine/twenty-orm/query-runner/workspace-query-runner'; import { type WorkspaceRepository } from 'src/engine/twenty-orm/repository/workspace.repository'; import { type RolePermissionConfig } from 'src/engine/twenty-orm/types/role-permission-config'; type CreateQueryBuilderOptions = { calledByWorkspaceEntityManager?: boolean; }; export class WorkspaceDataSource extends DataSource implements WorkspaceDataSourceInterface { readonly isGlobalFlow = false; readonly internalContext: WorkspaceInternalContext; readonly manager: WorkspaceEntityManager; featureFlagMapVersion: string; featureFlagMap: FeatureFlagMap; rolesPermissionsVersion: string; permissionsPerRoleId: ObjectsPermissionsByRoleId; dataSourceWithOverridenCreateQueryBuilder: WorkspaceDataSource; isPoolSharingEnabled: boolean; constructor( internalContext: WorkspaceInternalContext, options: DataSourceOptions, featureFlagMapVersion: string, featureFlagMap: FeatureFlagMap, rolesPermissionsVersion: string, permissionsPerRoleId: ObjectsPermissionsByRoleId, isPoolSharingEnabled: boolean, ) { super(options); this.internalContext = internalContext; this.featureFlagMap = featureFlagMap; this.featureFlagMapVersion = featureFlagMapVersion; // Recreate manager after internalContext has been initialized this.manager = this.createEntityManager(); this.rolesPermissionsVersion = rolesPermissionsVersion; this.permissionsPerRoleId = permissionsPerRoleId; this.isPoolSharingEnabled = isPoolSharingEnabled; } override getRepository( target: EntityTarget, permissionOptions?: RolePermissionConfig, authContext?: AuthContext, ): WorkspaceRepository { return this.manager.getRepository(target, permissionOptions, authContext); } override createEntityManager( queryRunner?: QueryRunner, ): WorkspaceEntityManager { return new WorkspaceEntityManager(this.internalContext, this, queryRunner); } override createQueryRunner( mode = 'master' as ReplicationMode, ): WorkspaceQueryRunner { const queryRunner = this.driver.createQueryRunner(mode); const manager = this.createEntityManager(queryRunner); Object.assign(queryRunner, { manager: manager }); // eslint-disable-next-line @typescript-eslint/no-explicit-any return queryRunner as any as WorkspaceQueryRunner; } // Do not use, only for specific permission-related purpose createQueryRunnerForEntityPersistExecutor( mode = 'master' as ReplicationMode, ) { if (this.dataSourceWithOverridenCreateQueryBuilder) { const queryRunner = this.driver.createQueryRunner(mode); const manager = new EntityManagerFactory().create( this.dataSourceWithOverridenCreateQueryBuilder, queryRunner, ); Object.assign(queryRunner, { manager: manager }); return queryRunner; } const dataSourceWithOverridenCreateQueryBuilder = Object.assign( Object.create(Object.getPrototypeOf(this)), this, { createQueryBuilder: ( entityOrRunner: EntityTarget | QueryRunner, alias?: string, queryRunner?: QueryRunner, ) => { if (isDefined(alias) && typeof alias === 'string') { const entity = entityOrRunner as EntityTarget; return this.createQueryBuilder(entity, alias, queryRunner, { calledByWorkspaceEntityManager: true, }); } else { const runner = entityOrRunner as QueryRunner; return this.createQueryBuilder(runner, { calledByWorkspaceEntityManager: true, }); } }, }, ); const queryRunner = this.driver.createQueryRunner(mode); const manager = new EntityManagerFactory().create( dataSourceWithOverridenCreateQueryBuilder, queryRunner, ); Object.assign(queryRunner, { manager: manager }); return queryRunner; } override createQueryBuilder( entityClass: EntityTarget, alias: string, queryRunner?: QueryRunner, options?: CreateQueryBuilderOptions, ): SelectQueryBuilder; override createQueryBuilder( queryRunner?: QueryRunner, options?: CreateQueryBuilderOptions, // eslint-disable-next-line @typescript-eslint/no-explicit-any ): SelectQueryBuilder; // Only callable from workspaceEntityManager to guarantee a permission check was run override createQueryBuilder( // eslint-disable-next-line @typescript-eslint/no-explicit-any queryRunnerOrEntityClass?: QueryRunner | EntityTarget, aliasOrOptions?: string | CreateQueryBuilderOptions, queryRunner?: QueryRunner, options?: CreateQueryBuilderOptions, // eslint-disable-next-line @typescript-eslint/no-explicit-any ): SelectQueryBuilder { let calledByWorkspaceEntityManager; const isCalledWithEntityTarget = isDefined(aliasOrOptions) && typeof aliasOrOptions === 'string'; if (isCalledWithEntityTarget) { calledByWorkspaceEntityManager = options?.calledByWorkspaceEntityManager; } else { calledByWorkspaceEntityManager = ( aliasOrOptions as CreateQueryBuilderOptions )?.calledByWorkspaceEntityManager; } if (!(calledByWorkspaceEntityManager === true)) { throw new PermissionsException( 'Method not allowed because permissions are not implemented at datasource level.', PermissionsExceptionCode.METHOD_NOT_ALLOWED, ); } if (isCalledWithEntityTarget) { // eslint-disable-next-line @typescript-eslint/no-explicit-any const entityClass = queryRunnerOrEntityClass as EntityTarget; return super.createQueryBuilder( entityClass, aliasOrOptions as string, queryRunner, ); } else { const queryRunner = queryRunnerOrEntityClass as QueryRunner; return super.createQueryBuilder(queryRunner); } } // eslint-disable-next-line @typescript-eslint/no-explicit-any override query( query: string, // eslint-disable-next-line @typescript-eslint/no-explicit-any parameters?: any[], queryRunner?: QueryRunner, options?: { shouldBypassPermissionChecks?: boolean; }, ): Promise { if (!options?.shouldBypassPermissionChecks) { throw new PermissionsException( 'Method not allowed because permissions are not implemented at datasource level.', PermissionsExceptionCode.METHOD_NOT_ALLOWED, ); } return super.query(query, parameters, queryRunner); } setRolesPermissionsVersion(rolesPermissionsVersion: string) { this.rolesPermissionsVersion = rolesPermissionsVersion; } setRolesPermissions(permissionsPerRoleId: ObjectsPermissionsByRoleId) { this.permissionsPerRoleId = permissionsPerRoleId; } setFeatureFlagMap(featureFlagMap: FeatureFlagMap) { this.featureFlagMap = featureFlagMap; } setFeatureFlagMapVersion(featureFlagMapVersion: string) { this.featureFlagMapVersion = featureFlagMapVersion; } override async destroy(): Promise { if (this.isPoolSharingEnabled) { // eslint-disable-next-line no-console console.log( `PromiseMemoizer Event: A WorkspaceDataSource for workspace ${this.internalContext.workspaceId} is being cleared. Actual pool closure managed by PgPoolSharedService. Not calling dataSource.destroy().`, ); // We should NOT call dataSource.destroy() here, because that would end // the shared pool, potentially affecting other active users of that pool. // The PgPoolSharedService is responsible for the lifecycle of shared pools. return Promise.resolve(); } else { // eslint-disable-next-line no-console console.log( `PromiseMemoizer Event: A WorkspaceDataSource for workspace ${this.internalContext.workspaceId} is being cleared. Calling safelyDestroyDataSource.`, ); return super.destroy(); } } }