Fix server integration tests 2 (#10818)

## Context
- Removing search* integration tests instead of fixing them because they
will be replaced by global search very soon
- Fixed billing + add missing seeds to make them work
- Fixed integration tests not using consistently the correct "test" db
- Fixed ci not running the with-db-reset configuration due to nx
configuration being used twice for different level of the command
- Enriched .env.test
- Fixed parts where exceptions were not thrown properly and not caught
by exception handler to convert to 400 when needed
- Refactored feature flag service that had 2 different implementations
in lab and admin panel + added tests
- Fixed race condition when migrations are created at the same timestamp
and doing the same type of operation, in this case object deletion could
break because table could be deleted earlier than its relations
- Fixed many integration tests that were not up to date since the CI has
been broken for a while

---------

Co-authored-by: Charles Bochet <charlesBochet@users.noreply.github.com>

packages/twenty-server/src/engine/core-modules/admin-panel/__tests__/admin-panel.service.spec.ts

@@ -9,29 +9,12 @@ import {
 import { LoginTokenService } from 'src/engine/core-modules/auth/token/services/login-token.service';
 import { DomainManagerService } from 'src/engine/core-modules/domain-manager/services/domain-manager.service';
 import { EnvironmentService } from 'src/engine/core-modules/environment/environment.service';
-import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
-import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
 import { User } from 'src/engine/core-modules/user/user.entity';
-import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 
 const UserFindOneMock = jest.fn();
-const WorkspaceFindOneMock = jest.fn();
-const FeatureFlagUpdateMock = jest.fn();
-const FeatureFlagSaveMock = jest.fn();
 const LoginTokenServiceGenerateLoginTokenMock = jest.fn();
 const EnvironmentServiceGetAllMock = jest.fn();
 
-jest.mock(
-  'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum',
-  () => {
-    return {
-      FeatureFlagKey: {
-        IsFlagEnabled: 'IS_FLAG_ENABLED',
-      },
-    };
-  },
-);
-
 jest.mock(
   '../../environment/constants/environment-variables-group-metadata',
   () => ({
@@ -62,25 +45,12 @@ describe('AdminPanelService', () => {
     const module: TestingModule = await Test.createTestingModule({
       providers: [
         AdminPanelService,
-        {
-          provide: getRepositoryToken(Workspace, 'core'),
-          useValue: {
-            findOne: WorkspaceFindOneMock,
-          },
-        },
         {
           provide: getRepositoryToken(User, 'core'),
           useValue: {
             findOne: UserFindOneMock,
           },
         },
-        {
-          provide: getRepositoryToken(FeatureFlag, 'core'),
-          useValue: {
-            update: FeatureFlagUpdateMock,
-            save: FeatureFlagSaveMock,
-          },
-        },
         {
           provide: LoginTokenService,
           useValue: {
@@ -112,80 +82,6 @@ describe('AdminPanelService', () => {
     expect(service).toBeDefined();
   });
 
-  it('should update an existing feature flag if it exists', async () => {
-    const workspaceId = 'workspace-id';
-    const featureFlag = 'IsFlagEnabled' as FeatureFlagKey;
-    const value = true;
-    const existingFlag = {
-      id: 'flag-id',
-      key: 'IS_FLAG_ENABLED',
-      value: false,
-    };
-
-    WorkspaceFindOneMock.mockReturnValueOnce({
-      id: workspaceId,
-      featureFlags: [existingFlag],
-    });
-
-    await service.updateWorkspaceFeatureFlags(workspaceId, featureFlag, value);
-
-    expect(FeatureFlagUpdateMock).toHaveBeenCalledWith(existingFlag.id, {
-      value,
-    });
-    expect(FeatureFlagSaveMock).not.toHaveBeenCalled();
-  });
-
-  it('should create a new feature flag if it does not exist', async () => {
-    const workspaceId = 'workspace-id';
-    const featureFlag = 'IsFlagEnabled' as FeatureFlagKey;
-    const value = true;
-
-    WorkspaceFindOneMock.mockReturnValueOnce({
-      id: workspaceId,
-      featureFlags: [],
-    });
-
-    await service.updateWorkspaceFeatureFlags(workspaceId, featureFlag, value);
-
-    expect(FeatureFlagSaveMock).toHaveBeenCalledWith({
-      key: 'IS_FLAG_ENABLED',
-      value,
-      workspaceId,
-    });
-    expect(FeatureFlagUpdateMock).not.toHaveBeenCalled();
-  });
-
-  it('should throw an exception if the workspace is not found', async () => {
-    const workspaceId = 'non-existent-workspace';
-    const featureFlag = 'IsFlagEnabled' as FeatureFlagKey;
-    const value = true;
-
-    WorkspaceFindOneMock.mockReturnValueOnce(null);
-
-    await expect(
-      service.updateWorkspaceFeatureFlags(workspaceId, featureFlag, value),
-    ).rejects.toThrowError(
-      new AuthException('Workspace not found', AuthExceptionCode.INVALID_INPUT),
-    );
-  });
-
-  it('should throw an exception if the flag is not found', async () => {
-    const workspaceId = 'non-existent-workspace';
-    const featureFlag = 'IsUnknownFlagEnabled' as FeatureFlagKey;
-    const value = true;
-
-    WorkspaceFindOneMock.mockReturnValueOnce(null);
-
-    await expect(
-      service.updateWorkspaceFeatureFlags(workspaceId, featureFlag, value),
-    ).rejects.toThrowError(
-      new AuthException(
-        'Invalid feature flag key',
-        AuthExceptionCode.INVALID_INPUT,
-      ),
-    );
-  });
-
   it('should impersonate a user and return workspace and loginToken on success', async () => {
     const mockUser = {
       id: 'user-id',

packages/twenty-server/src/engine/core-modules/admin-panel/admin-panel.module.ts

@@ -7,20 +7,20 @@ import { AdminPanelResolver } from 'src/engine/core-modules/admin-panel/admin-pa
 import { AdminPanelService } from 'src/engine/core-modules/admin-panel/admin-panel.service';
 import { AuthModule } from 'src/engine/core-modules/auth/auth.module';
 import { DomainManagerModule } from 'src/engine/core-modules/domain-manager/domain-manager.module';
-import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
+import { FeatureFlagModule } from 'src/engine/core-modules/feature-flag/feature-flag.module';
 import { HealthModule } from 'src/engine/core-modules/health/health.module';
 import { RedisClientModule } from 'src/engine/core-modules/redis-client/redis-client.module';
 import { User } from 'src/engine/core-modules/user/user.entity';
-import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 
 @Module({
   imports: [
-    TypeOrmModule.forFeature([User, Workspace, FeatureFlag], 'core'),
+    TypeOrmModule.forFeature([User], 'core'),
     AuthModule,
     DomainManagerModule,
     HealthModule,
     RedisClientModule,
     TerminusModule,
+    FeatureFlagModule,
   ],
   providers: [AdminPanelResolver, AdminPanelService, AdminPanelHealthService],
   exports: [AdminPanelService],

packages/twenty-server/src/engine/core-modules/admin-panel/admin-panel.resolver.ts

@@ -12,6 +12,9 @@ import { UserLookup } from 'src/engine/core-modules/admin-panel/dtos/user-lookup
 import { UserLookupInput } from 'src/engine/core-modules/admin-panel/dtos/user-lookup.input';
 import { QueueMetricsTimeRange } from 'src/engine/core-modules/admin-panel/enums/queue-metrics-time-range.enum';
 import { AuthGraphqlApiExceptionFilter } from 'src/engine/core-modules/auth/filters/auth-graphql-api-exception.filter';
+import { FeatureFlagException } from 'src/engine/core-modules/feature-flag/feature-flag.exception';
+import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
+import { UserInputError } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
 import { HealthIndicatorId } from 'src/engine/core-modules/health/enums/health-indicator-id.enum';
 import { WorkerHealthIndicator } from 'src/engine/core-modules/health/indicators/worker.health';
 import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
@@ -30,6 +33,7 @@ export class AdminPanelResolver {
     private adminService: AdminPanelService,
     private adminPanelHealthService: AdminPanelHealthService,
     private workerHealthIndicator: WorkerHealthIndicator,
+    private featureFlagService: FeatureFlagService,
   ) {}
 
   @UseGuards(WorkspaceAuthGuard, UserAuthGuard, ImpersonateGuard)
@@ -53,13 +57,21 @@ export class AdminPanelResolver {
   async updateWorkspaceFeatureFlag(
     @Args() updateFlagInput: UpdateWorkspaceFeatureFlagInput,
   ): Promise<boolean> {
-    await this.adminService.updateWorkspaceFeatureFlags(
-      updateFlagInput.workspaceId,
-      updateFlagInput.featureFlag,
-      updateFlagInput.value,
-    );
+    try {
+      await this.featureFlagService.upsertWorkspaceFeatureFlag({
+        workspaceId: updateFlagInput.workspaceId,
+        featureFlag: updateFlagInput.featureFlag,
+        value: updateFlagInput.value,
+      });
 
-    return true;
+      return true;
+    } catch (error) {
+      if (error instanceof FeatureFlagException) {
+        throw new UserInputError(error.message);
+      }
+
+      throw error;
+    }
   }
 
   @UseGuards(WorkspaceAuthGuard, UserAuthGuard, AdminPanelGuard)

packages/twenty-server/src/engine/core-modules/admin-panel/admin-panel.service.ts

@@ -18,15 +18,8 @@ import { EnvironmentVariablesGroup } from 'src/engine/core-modules/environment/e
 import { EnvironmentService } from 'src/engine/core-modules/environment/environment.service';
 import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
 import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
-import {
-  FeatureFlagException,
-  FeatureFlagExceptionCode,
-} from 'src/engine/core-modules/feature-flag/feature-flag.exception';
-import { featureFlagValidator } from 'src/engine/core-modules/feature-flag/validates/feature-flag.validate';
 import { User } from 'src/engine/core-modules/user/user.entity';
 import { userValidator } from 'src/engine/core-modules/user/user.validate';
-import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
-import { workspaceValidator } from 'src/engine/core-modules/workspace/workspace.validate';
 
 @Injectable()
 export class AdminPanelService {
@@ -36,10 +29,6 @@ export class AdminPanelService {
     private readonly domainManagerService: DomainManagerService,
     @InjectRepository(User, 'core')
     private readonly userRepository: Repository<User>,
-    @InjectRepository(Workspace, 'core')
-    private readonly workspaceRepository: Repository<Workspace>,
-    @InjectRepository(FeatureFlag, 'core')
-    private readonly featureFlagRepository: Repository<FeatureFlag>,
   ) {}
 
   async impersonate(userId: string, workspaceId: string) {
@@ -131,44 +120,6 @@ export class AdminPanelService {
     };
   }
 
-  async updateWorkspaceFeatureFlags(
-    workspaceId: string,
-    featureFlag: FeatureFlagKey,
-    value: boolean,
-  ) {
-    featureFlagValidator.assertIsFeatureFlagKey(
-      featureFlag,
-      new FeatureFlagException(
-        'Invalid feature flag key',
-        FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
-      ),
-    );
-
-    const workspace = await this.workspaceRepository.findOne({
-      where: { id: workspaceId },
-      relations: ['featureFlags'],
-    });
-
-    workspaceValidator.assertIsDefinedOrThrow(
-      workspace,
-      new AuthException('Workspace not found', AuthExceptionCode.INVALID_INPUT),
-    );
-
-    const existingFlag = workspace.featureFlags?.find(
-      (flag) => flag.key === FeatureFlagKey[featureFlag],
-    );
-
-    if (existingFlag) {
-      await this.featureFlagRepository.update(existingFlag.id, { value });
-    } else {
-      await this.featureFlagRepository.save({
-        key: FeatureFlagKey[featureFlag],
-        value,
-        workspaceId: workspace.id,
-      });
-    }
-  }
-
   getEnvironmentVariablesGrouped(): EnvironmentVariablesOutput {
     const rawEnvVars = this.environmentService.getAll();
     const groupedData = new Map<

packages/twenty-server/src/engine/core-modules/domain-manager/services/domain-manager.service.ts

@@ -1,17 +1,17 @@
 import { Injectable, Logger } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';
 
-import { Repository } from 'typeorm';
 import { isDefined } from 'twenty-shared';
+import { Repository } from 'typeorm';
 
+import { SEED_APPLE_WORKSPACE_ID } from 'src/database/typeorm-seeds/core/workspaces';
+import { WorkspaceSubdomainCustomDomainAndIsCustomDomainEnabledType } from 'src/engine/core-modules/domain-manager/domain-manager.type';
 import { CustomDomainValidRecords } from 'src/engine/core-modules/domain-manager/dtos/custom-domain-valid-records';
 import { generateRandomSubdomain } from 'src/engine/core-modules/domain-manager/utils/generate-random-subdomain';
 import { getSubdomainFromEmail } from 'src/engine/core-modules/domain-manager/utils/get-subdomain-from-email';
 import { getSubdomainNameFromDisplayName } from 'src/engine/core-modules/domain-manager/utils/get-subdomain-name-from-display-name';
 import { EnvironmentService } from 'src/engine/core-modules/environment/environment.service';
 import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
-import { WorkspaceSubdomainCustomDomainAndIsCustomDomainEnabledType } from 'src/engine/core-modules/domain-manager/domain-manager.type';
-import { SEED_APPLE_WORKSPACE_ID } from 'src/database/typeorm-seeds/core/workspaces';
 import { workspaceValidator } from 'src/engine/core-modules/workspace/workspace.validate';
 
 @Injectable()

packages/twenty-server/src/engine/core-modules/feature-flag/services/__tests__/feature-flag.service.spec.ts

@@ -0,0 +1,268 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { getRepositoryToken } from '@nestjs/typeorm';
+
+import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
+import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
+import {
+  FeatureFlagException,
+  FeatureFlagExceptionCode,
+} from 'src/engine/core-modules/feature-flag/feature-flag.exception';
+import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
+import { featureFlagValidator } from 'src/engine/core-modules/feature-flag/validates/feature-flag.validate';
+import { publicFeatureFlagValidator } from 'src/engine/core-modules/feature-flag/validates/is-public-feature-flag.validate';
+
+jest.mock(
+  'src/engine/core-modules/feature-flag/validates/is-public-feature-flag.validate',
+);
+jest.mock(
+  'src/engine/core-modules/feature-flag/validates/feature-flag.validate',
+);
+
+describe('FeatureFlagService', () => {
+  let service: FeatureFlagService;
+
+  const mockFeatureFlagRepository = {
+    findOneBy: jest.fn(),
+    find: jest.fn(),
+    upsert: jest.fn(),
+  };
+
+  const workspaceId = 'workspace-id';
+  const featureFlag = FeatureFlagKey.IsWorkflowEnabled;
+
+  beforeEach(async () => {
+    jest.clearAllMocks();
+
+    (
+      publicFeatureFlagValidator.assertIsPublicFeatureFlag as jest.Mock
+    ).mockReset();
+    (featureFlagValidator.assertIsFeatureFlagKey as jest.Mock).mockReset();
+
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [
+        FeatureFlagService,
+        {
+          provide: getRepositoryToken(FeatureFlag, 'core'),
+          useValue: mockFeatureFlagRepository,
+        },
+      ],
+    }).compile();
+
+    service = module.get<FeatureFlagService>(FeatureFlagService);
+  });
+
+  it('should be defined', () => {
+    expect(service).toBeDefined();
+  });
+
+  describe('isFeatureEnabled', () => {
+    it('should return true when feature flag is enabled', async () => {
+      // Prepare
+      mockFeatureFlagRepository.findOneBy.mockResolvedValue({
+        key: featureFlag,
+        value: true,
+        workspaceId,
+      });
+
+      // Act
+      const result = await service.isFeatureEnabled(featureFlag, workspaceId);
+
+      // Assert
+      expect(result).toBe(true);
+      expect(mockFeatureFlagRepository.findOneBy).toHaveBeenCalledWith({
+        workspaceId,
+        key: featureFlag,
+        value: true,
+      });
+    });
+
+    it('should return false when feature flag is not found', async () => {
+      // Prepare
+      mockFeatureFlagRepository.findOneBy.mockResolvedValue(null);
+
+      // Act
+      const result = await service.isFeatureEnabled(featureFlag, workspaceId);
+
+      // Assert
+      expect(result).toBe(false);
+    });
+
+    it('should return false when feature flag value is false', async () => {
+      // Prepare
+      mockFeatureFlagRepository.findOneBy.mockResolvedValue({
+        key: featureFlag,
+        value: false,
+        workspaceId,
+      });
+
+      // Act
+      const result = await service.isFeatureEnabled(featureFlag, workspaceId);
+
+      // Assert
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('getWorkspaceFeatureFlags', () => {
+    it('should return all feature flags for a workspace', async () => {
+      // Prepare
+      const mockFeatureFlags = [
+        { key: FeatureFlagKey.IsWorkflowEnabled, value: true, workspaceId },
+        { key: FeatureFlagKey.IsCopilotEnabled, value: false, workspaceId },
+      ];
+
+      mockFeatureFlagRepository.find.mockResolvedValue(mockFeatureFlags);
+
+      // Act
+      const result = await service.getWorkspaceFeatureFlags(workspaceId);
+
+      // Assert
+      expect(result).toEqual(mockFeatureFlags);
+      expect(mockFeatureFlagRepository.find).toHaveBeenCalledWith({
+        where: { workspaceId },
+      });
+    });
+  });
+
+  describe('getWorkspaceFeatureFlagsMap', () => {
+    it('should return a map of feature flags for a workspace', async () => {
+      // Prepare
+      const mockFeatureFlags = [
+        { key: FeatureFlagKey.IsWorkflowEnabled, value: true, workspaceId },
+        { key: FeatureFlagKey.IsCopilotEnabled, value: false, workspaceId },
+      ];
+
+      mockFeatureFlagRepository.find.mockResolvedValue(mockFeatureFlags);
+
+      // Act
+      const result = await service.getWorkspaceFeatureFlagsMap(workspaceId);
+
+      // Assert
+      expect(result).toEqual({
+        [FeatureFlagKey.IsWorkflowEnabled]: true,
+        [FeatureFlagKey.IsCopilotEnabled]: false,
+      });
+    });
+  });
+
+  describe('enableFeatureFlags', () => {
+    it('should enable multiple feature flags for a workspace', async () => {
+      // Prepare
+      const keys = [
+        FeatureFlagKey.IsWorkflowEnabled,
+        FeatureFlagKey.IsCopilotEnabled,
+      ];
+
+      mockFeatureFlagRepository.upsert.mockResolvedValue({});
+
+      // Act
+      await service.enableFeatureFlags(keys, workspaceId);
+
+      // Assert
+      expect(mockFeatureFlagRepository.upsert).toHaveBeenCalledWith(
+        keys.map((key) => ({ workspaceId, key, value: true })),
+        {
+          conflictPaths: ['workspaceId', 'key'],
+          skipUpdateIfNoValuesChanged: true,
+        },
+      );
+    });
+  });
+
+  describe('upsertWorkspaceFeatureFlag', () => {
+    it('should upsert a feature flag for a workspace', async () => {
+      // Prepare
+      const value = true;
+      const mockFeatureFlag = {
+        key: featureFlag,
+        value,
+        workspaceId,
+      };
+
+      mockFeatureFlagRepository.upsert.mockResolvedValue({
+        generatedMaps: [mockFeatureFlag],
+      });
+
+      (
+        featureFlagValidator.assertIsFeatureFlagKey as jest.Mock
+      ).mockImplementation(() => true);
+
+      // Act
+      const result = await service.upsertWorkspaceFeatureFlag({
+        workspaceId,
+        featureFlag,
+        value,
+      });
+
+      // Assert
+      expect(result).toEqual(mockFeatureFlag);
+      expect(mockFeatureFlagRepository.upsert).toHaveBeenCalledWith(
+        {
+          key: FeatureFlagKey[featureFlag],
+          value,
+          workspaceId,
+        },
+        {
+          conflictPaths: ['workspaceId', 'key'],
+          skipUpdateIfNoValuesChanged: true,
+        },
+      );
+    });
+
+    it('should throw an exception when feature flag key is invalid', async () => {
+      // Prepare
+      const invalidFeatureFlag = 'INVALID_KEY' as FeatureFlagKey;
+      const value = true;
+
+      (
+        featureFlagValidator.assertIsFeatureFlagKey as jest.Mock
+      ).mockImplementation(() => {
+        throw new FeatureFlagException(
+          'Invalid feature flag key',
+          FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
+        );
+      });
+
+      // Act & Assert
+      await expect(
+        service.upsertWorkspaceFeatureFlag({
+          workspaceId,
+          featureFlag: invalidFeatureFlag,
+          value,
+        }),
+      ).rejects.toThrow(
+        new FeatureFlagException(
+          'Invalid feature flag key',
+          FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
+        ),
+      );
+    });
+
+    it('should throw an exception when non-public feature flag is used with shouldBePublic=true', async () => {
+      // Prepare
+      (
+        publicFeatureFlagValidator.assertIsPublicFeatureFlag as jest.Mock
+      ).mockImplementation(() => {
+        throw new FeatureFlagException(
+          'Invalid feature flag key, flag is not public',
+          FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
+        );
+      });
+
+      // Act & Assert
+      await expect(
+        service.upsertWorkspaceFeatureFlag({
+          workspaceId,
+          featureFlag,
+          value: true,
+          shouldBePublic: true,
+        }),
+      ).rejects.toThrow(
+        new FeatureFlagException(
+          'Invalid feature flag key, flag is not public',
+          FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
+        ),
+      );
+    });
+  });
+});

packages/twenty-server/src/engine/core-modules/feature-flag/services/feature-flag.service.ts

@@ -7,6 +7,12 @@ import { FeatureFlagMap } from 'src/engine/core-modules/feature-flag/interfaces/
 
 import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
 import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
+import {
+  FeatureFlagException,
+  FeatureFlagExceptionCode,
+} from 'src/engine/core-modules/feature-flag/feature-flag.exception';
+import { featureFlagValidator } from 'src/engine/core-modules/feature-flag/validates/feature-flag.validate';
+import { publicFeatureFlagValidator } from 'src/engine/core-modules/feature-flag/validates/is-public-feature-flag.validate';
 
 @Injectable()
 export class FeatureFlagService {
@@ -64,4 +70,48 @@ export class FeatureFlagService {
       },
     );
   }
+
+  public async upsertWorkspaceFeatureFlag({
+    workspaceId,
+    featureFlag,
+    value,
+    shouldBePublic = false,
+  }: {
+    workspaceId: string;
+    featureFlag: FeatureFlagKey;
+    value: boolean;
+    shouldBePublic?: boolean;
+  }): Promise<FeatureFlag> {
+    if (shouldBePublic) {
+      publicFeatureFlagValidator.assertIsPublicFeatureFlag(
+        featureFlag,
+        new FeatureFlagException(
+          'Invalid feature flag key, flag is not public',
+          FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
+        ),
+      );
+    }
+
+    featureFlagValidator.assertIsFeatureFlagKey(
+      featureFlag,
+      new FeatureFlagException(
+        'Invalid feature flag key',
+        FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
+      ),
+    );
+
+    const upsertResult = await this.featureFlagRepository.upsert(
+      {
+        key: FeatureFlagKey[featureFlag],
+        value,
+        workspaceId: workspaceId,
+      },
+      {
+        conflictPaths: ['workspaceId', 'key'],
+        skipUpdateIfNoValuesChanged: true,
+      },
+    );
+
+    return upsertResult.generatedMaps[0] as FeatureFlag;
+  }
 }

packages/twenty-server/src/engine/core-modules/lab/lab.module.ts

@@ -1,22 +1,13 @@
 import { Module } from '@nestjs/common';
-import { TypeOrmModule } from '@nestjs/typeorm';
 
-import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
 import { FeatureFlagModule } from 'src/engine/core-modules/feature-flag/feature-flag.module';
-import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module';
 
 import { LabResolver } from './lab.resolver';
 
-import { LabService } from './services/lab.service';
-
 @Module({
-  imports: [
-    TypeOrmModule.forFeature([FeatureFlag, Workspace], 'core'),
-    FeatureFlagModule,
-    PermissionsModule,
-  ],
-  providers: [LabService, LabResolver],
-  exports: [LabService],
+  imports: [FeatureFlagModule, PermissionsModule],
+  providers: [LabResolver],
+  exports: [],
 })
 export class LabModule {}

packages/twenty-server/src/engine/core-modules/lab/lab.resolver.ts

@@ -3,8 +3,10 @@ import { Args, Mutation, Resolver } from '@nestjs/graphql';
 
 import { AuthGraphqlApiExceptionFilter } from 'src/engine/core-modules/auth/filters/auth-graphql-api-exception.filter';
 import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
+import { FeatureFlagException } from 'src/engine/core-modules/feature-flag/feature-flag.exception';
+import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
+import { UserInputError } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
 import { UpdateLabPublicFeatureFlagInput } from 'src/engine/core-modules/lab/dtos/update-lab-public-feature-flag.input';
-import { LabService } from 'src/engine/core-modules/lab/services/lab.service';
 import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
 import { SettingsPermissionsGuard } from 'src/engine/guards/settings-permissions.guard';
@@ -16,7 +18,7 @@ import { PermissionsGraphqlApiExceptionFilter } from 'src/engine/metadata-module
 @UseFilters(AuthGraphqlApiExceptionFilter, PermissionsGraphqlApiExceptionFilter)
 @UseGuards(SettingsPermissionsGuard(SettingsPermissions.WORKSPACE))
 export class LabResolver {
-  constructor(private labService: LabService) {}
+  constructor(private featureFlagService: FeatureFlagService) {}
 
   @UseGuards(WorkspaceAuthGuard)
   @Mutation(() => FeatureFlag)
@@ -24,6 +26,18 @@ export class LabResolver {
     @Args('input') input: UpdateLabPublicFeatureFlagInput,
     @AuthWorkspace() workspace: Workspace,
   ): Promise<FeatureFlag> {
-    return this.labService.updateLabPublicFeatureFlag(workspace.id, input);
+    try {
+      return await this.featureFlagService.upsertWorkspaceFeatureFlag({
+        workspaceId: workspace.id,
+        featureFlag: input.publicFeatureFlag,
+        value: input.value,
+        shouldBePublic: true,
+      });
+    } catch (error) {
+      if (error instanceof FeatureFlagException) {
+        throw new UserInputError(error.message);
+      }
+      throw error;
+    }
   }
 }

packages/twenty-server/src/engine/core-modules/lab/services/lab.service.ts

@@ -1,81 +0,0 @@
-import { Injectable } from '@nestjs/common';
-import { InjectRepository } from '@nestjs/typeorm';
-
-import { Repository } from 'typeorm';
-
-import {
-  AuthException,
-  AuthExceptionCode,
-} from 'src/engine/core-modules/auth/auth.exception';
-import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
-import { FeatureFlag } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
-import {
-  FeatureFlagException,
-  FeatureFlagExceptionCode,
-} from 'src/engine/core-modules/feature-flag/feature-flag.exception';
-import { featureFlagValidator } from 'src/engine/core-modules/feature-flag/validates/feature-flag.validate';
-import { publicFeatureFlagValidator } from 'src/engine/core-modules/feature-flag/validates/is-public-feature-flag.validate';
-import { UpdateLabPublicFeatureFlagInput } from 'src/engine/core-modules/lab/dtos/update-lab-public-feature-flag.input';
-import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
-import { workspaceValidator } from 'src/engine/core-modules/workspace/workspace.validate';
-
-@Injectable()
-export class LabService {
-  constructor(
-    @InjectRepository(FeatureFlag, 'core')
-    private readonly featureFlagRepository: Repository<FeatureFlag>,
-    @InjectRepository(Workspace, 'core')
-    private readonly workspaceRepository: Repository<Workspace>,
-  ) {}
-
-  async updateLabPublicFeatureFlag(
-    workspaceId: string,
-    payload: UpdateLabPublicFeatureFlagInput,
-  ): Promise<FeatureFlag> {
-    featureFlagValidator.assertIsFeatureFlagKey(
-      payload.publicFeatureFlag,
-      new FeatureFlagException(
-        'Invalid feature flag key',
-        FeatureFlagExceptionCode.INVALID_FEATURE_FLAG_KEY,
-      ),
-    );
-
-    publicFeatureFlagValidator.assertIsPublicFeatureFlag(
-      FeatureFlagKey[payload.publicFeatureFlag],
-      new FeatureFlagException(
-        'Feature flag is not public',
-        FeatureFlagExceptionCode.FEATURE_FLAG_IS_NOT_PUBLIC,
-      ),
-    );
-
-    const workspace = await this.workspaceRepository.findOne({
-      where: { id: workspaceId },
-    });
-
-    workspaceValidator.assertIsDefinedOrThrow(
-      workspace,
-      new AuthException('Workspace not found', AuthExceptionCode.INVALID_INPUT),
-    );
-
-    const existingFlag = await this.featureFlagRepository.findOne({
-      where: {
-        workspaceId,
-        key: FeatureFlagKey[payload.publicFeatureFlag],
-      },
-    });
-
-    if (existingFlag) {
-      await this.featureFlagRepository.update(existingFlag.id, {
-        value: payload.value,
-      });
-
-      return { ...existingFlag, value: payload.value };
-    }
-
-    return this.featureFlagRepository.save({
-      key: FeatureFlagKey[payload.publicFeatureFlag],
-      value: payload.value,
-      workspaceId,
-    });
-  }
-}