feat: extend twenty orm

packages/twenty-server/src/engine/twenty-orm/datasource/workspace.datasource.ts

@@ -0,0 +1,24 @@
+import {
+  DataSource,
+  EntityManager,
+  EntityTarget,
+  ObjectLiteral,
+  QueryRunner,
+} from 'typeorm';
+
+import { WorkspaceRepository } from 'src/engine/twenty-orm/repository/workspace.repository';
+import { WorkspaceEntityManager } from 'src/engine/twenty-orm/entity-manager/entity.manager';
+
+export class WorkspaceDataSource extends DataSource {
+  readonly manager: WorkspaceEntityManager;
+
+  override getRepository<Entity extends ObjectLiteral>(
+    target: EntityTarget<Entity>,
+  ): WorkspaceRepository<Entity> {
+    return this.manager.getRepository(target);
+  }
+
+  override createEntityManager(queryRunner?: QueryRunner): EntityManager {
+    return new WorkspaceEntityManager(this, queryRunner);
+  }
+}

packages/twenty-server/src/engine/twenty-orm/decorators/inject-workspace-datasource.decorator.ts

@@ -2,5 +2,6 @@ import { Inject } from '@nestjs/common';
 
 import { TWENTY_ORM_WORKSPACE_DATASOURCE } from 'src/engine/twenty-orm/twenty-orm.constants';
 
+// nit: The datasource can be null if it's used outside of an authenticated request context
 export const InjectWorkspaceDatasource = () =>
   Inject(TWENTY_ORM_WORKSPACE_DATASOURCE);

packages/twenty-server/src/engine/twenty-orm/decorators/inject-workspace-repository.decorator.ts

@@ -3,6 +3,7 @@ import { EntityClassOrSchema } from '@nestjs/typeorm/dist/interfaces/entity-clas
 
 import { getWorkspaceRepositoryToken } from 'src/engine/twenty-orm/utils/get-workspace-repository-token.util';
 
+// nit: The repository can be null if it's used outside of an authenticated request context
 export const InjectWorkspaceRepository = (
   entity: EntityClassOrSchema,
 ): ReturnType<typeof Inject> => Inject(getWorkspaceRepositoryToken(entity));

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

@@ -0,0 +1,24 @@
+import { EntityManager, EntityTarget, ObjectLiteral } from 'typeorm';
+
+import { WorkspaceRepository } from 'src/engine/twenty-orm/repository/workspace.repository';
+
+export class WorkspaceEntityManager extends EntityManager {
+  override getRepository<Entity extends ObjectLiteral>(
+    target: EntityTarget<Entity>,
+  ): WorkspaceRepository<Entity> {
+    // find already created repository instance and return it if found
+    const repoFromMap = this.repositories.get(target);
+
+    if (repoFromMap) return repoFromMap as WorkspaceRepository<Entity>;
+
+    const newRepository = new WorkspaceRepository<Entity>(
+      target,
+      this,
+      this.queryRunner,
+    );
+
+    this.repositories.set(target, newRepository);
+
+    return newRepository;
+  }
+}

packages/twenty-server/src/engine/twenty-orm/factories/entity-schema.factory.ts

@@ -5,6 +5,7 @@ import { EntitySchema } from 'typeorm';
 import { EntitySchemaColumnFactory } from 'src/engine/twenty-orm/factories/entity-schema-column.factory';
 import { EntitySchemaRelationFactory } from 'src/engine/twenty-orm/factories/entity-schema-relation.factory';
 import { metadataArgsStorage } from 'src/engine/twenty-orm/storage/metadata-args.storage';
+import { ObjectLiteralStorage } from 'src/engine/twenty-orm/storage/object-literal.storage';
 
 @Injectable()
 export class EntitySchemaFactory {
@@ -33,11 +34,15 @@ export class EntitySchemaFactory {
       relationMetadataArgsCollection,
     );
 
-    return new EntitySchema({
+    const entitySchema = new EntitySchema({
       name: objectMetadataArgs.nameSingular,
       tableName: objectMetadataArgs.nameSingular,
       columns,
       relations,
     });
+
+    ObjectLiteralStorage.setObjectLiteral(entitySchema, target);
+
+    return entitySchema;
   }
 }

packages/twenty-server/src/engine/twenty-orm/factories/index.ts

@@ -1,11 +1,13 @@
 import { EntitySchemaColumnFactory } from 'src/engine/twenty-orm/factories/entity-schema-column.factory';
 import { EntitySchemaRelationFactory } from 'src/engine/twenty-orm/factories/entity-schema-relation.factory';
 import { EntitySchemaFactory } from 'src/engine/twenty-orm/factories/entity-schema.factory';
+import { ScopedWorkspaceDatasourceFactory } from 'src/engine/twenty-orm/factories/scoped-workspace-datasource.factory';
 import { WorkspaceDatasourceFactory } from 'src/engine/twenty-orm/factories/workspace-datasource.factory';
 
 export const entitySchemaFactories = [
   EntitySchemaColumnFactory,
   EntitySchemaRelationFactory,
   EntitySchemaFactory,
   WorkspaceDatasourceFactory,
+  ScopedWorkspaceDatasourceFactory,
 ];

packages/twenty-server/src/engine/twenty-orm/factories/scoped-workspace-datasource.factory.ts

@@ -0,0 +1,25 @@
+import { Inject, Injectable, Scope } from '@nestjs/common';
+import { REQUEST } from '@nestjs/core';
+
+import { EntitySchema } from 'typeorm';
+
+import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
+import { WorkspaceDatasourceFactory } from 'src/engine/twenty-orm/factories/workspace-datasource.factory';
+
+@Injectable({ scope: Scope.REQUEST })
+export class ScopedWorkspaceDatasourceFactory {
+  constructor(
+    @Inject(REQUEST) private readonly request: Request,
+    private readonly workspaceDataSourceFactory: WorkspaceDatasourceFactory,
+  ) {}
+
+  public async create(entities: EntitySchema[]) {
+    const workspace: Workspace | undefined = this.request['req']?.['workspace'];
+
+    if (!workspace) {
+      return null;
+    }
+
+    return this.workspaceDataSourceFactory.create(entities, workspace.id);
+  }
+}

packages/twenty-server/src/engine/twenty-orm/factories/workspace-datasource.factory.ts

@@ -1,42 +1,33 @@
-import { Inject, Injectable, Scope } from '@nestjs/common';
-import { REQUEST } from '@nestjs/core';
+import { Injectable } from '@nestjs/common';
 
-import { DataSource, EntitySchema } from 'typeorm';
+import { EntitySchema } from 'typeorm';
 
 import { DataSourceService } from 'src/engine/metadata-modules/data-source/data-source.service';
 import { EnvironmentService } from 'src/engine/integrations/environment/environment.service';
-import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 import { DataSourceStorage } from 'src/engine/twenty-orm/storage/data-source.storage';
+import { WorkspaceDataSource } from 'src/engine/twenty-orm/datasource/workspace.datasource';
 
-@Injectable({ scope: Scope.REQUEST })
+@Injectable()
 export class WorkspaceDatasourceFactory {
   constructor(
-    @Inject(REQUEST) private readonly request: Request,
     private readonly dataSourceService: DataSourceService,
     private readonly environmentService: EnvironmentService,
   ) {}
 
-  public async createWorkspaceDatasource(entities: EntitySchema[]) {
-    const workspace: Workspace = this.request['req']['workspace'];
-
-    if (!workspace) {
-      return null;
-    }
-
-    const storedWorkspaceDataSource = DataSourceStorage.getDataSource(
-      workspace.id,
-    );
+  public async create(entities: EntitySchema[], workspaceId: string) {
+    const storedWorkspaceDataSource =
+      DataSourceStorage.getDataSource(workspaceId);
 
     if (storedWorkspaceDataSource) {
       return storedWorkspaceDataSource;
     }
 
     const dataSourceMetadata =
       await this.dataSourceService.getLastDataSourceMetadataFromWorkspaceIdOrFail(
-        workspace.id,
+        workspaceId,
       );
 
-    const workspaceDataSource = new DataSource({
+    const workspaceDataSource = new WorkspaceDataSource({
       url:
         dataSourceMetadata.url ??
         this.environmentService.get('PG_DATABASE_URL'),
@@ -51,7 +42,7 @@ export class WorkspaceDatasourceFactory {
 
     await workspaceDataSource.initialize();
 
-    DataSourceStorage.setDataSource(workspace.id, workspaceDataSource);
+    DataSourceStorage.setDataSource(workspaceId, workspaceDataSource);
 
     return workspaceDataSource;
   }