fix: add concurrency guard to Lambda function build/execute lifecycle

https://sonarly.com/issue/19253?type=bug Concurrent workflow executions of the same logic function race on Lambda create/delete, causing one execution to time out when its function is deleted mid-poll by another execution's `build()` call. Fix: Three changes to `lambda.driver.ts` to address the concurrent execution race condition: 1. **Import `ResourceConflictException`** from `@aws-sdk/client-lambda` (line 19) — needed to catch 409 Conflict responses from AWS Lambda API. 2. **Handle `ResourceConflictException` on all three `CreateFunctionCommand` call sites:** - `ensureYarnInstallLambdaExists()` (line 396-402) - `ensureBuilderLambdaExists()` (line 491-496) - `build()` (line 968-974) When a concurrent execution already created the same function, the 409 is safely caught and execution continues to `waitFunctionActive`, which will wait for the existing function to become ready. 3. **Re-check `isAlreadyBuilt` after layer resolution in `build()`** (lines 926-938) — this is the key fix for the root cause. When `isSdkLayerStale=true`, the first check at line 905 is bypassed, but `ensureSdkLayer()` takes significant time (downloading archive, creating zip, publishing layer). During this time, a concurrent execution may have already rebuilt the function with the correct layers. The re-check prevents this execution from calling `delete()` on a function that another execution is currently waiting on via `waitFunctionActive`, which was the direct cause of the TimeoutError.
Deprecate dataSource table with dual-write to workspace.databaseSchema (#19059 )
2026-03-28 17:48:35 +00:00 · 2026-03-28 11:29:19 +01:00
20 changed files with 199 additions and 78 deletions
@@ -780,8 +780,7 @@ type Workspace {
  viewGroups: [ViewGroup!]
  viewSorts: [ViewSort!]
  metadataVersion: Float!
-  databaseUrl: String!
-  databaseSchema: String!
+  databaseSchema: String
  subdomain: String!
  customDomain: String
  isGoogleAuthEnabled: Boolean!
@@ -1749,6 +1748,7 @@ enum FeatureFlagKey {
  IS_CONNECTED_ACCOUNT_MIGRATED
  IS_GRAPHQL_QUERY_TIMING_ENABLED
  IS_RECORD_TABLE_WIDGET_ENABLED
+  IS_DATASOURCE_MIGRATED
 }

 type ClientConfig {
@@ -565,8 +565,7 @@ export interface Workspace {
    viewGroups?: ViewGroup[]
    viewSorts?: ViewSort[]
    metadataVersion: Scalars['Float']
-    databaseUrl: Scalars['String']
-    databaseSchema: Scalars['String']
+    databaseSchema?: Scalars['String']
    subdomain: Scalars['String']
    customDomain?: Scalars['String']
    isGoogleAuthEnabled: Scalars['Boolean']
@@ -1428,7 +1427,7 @@ export interface PublicFeatureFlag {
    __typename: 'PublicFeatureFlag'
 }

-export type FeatureFlagKey = 'IS_UNIQUE_INDEXES_ENABLED' | 'IS_JSON_FILTER_ENABLED' | 'IS_AI_ENABLED' | 'IS_APPLICATION_ENABLED' | 'IS_MARKETPLACE_ENABLED' | 'IS_RECORD_PAGE_LAYOUT_EDITING_ENABLED' | 'IS_PUBLIC_DOMAIN_ENABLED' | 'IS_EMAILING_DOMAIN_ENABLED' | 'IS_DASHBOARD_V2_ENABLED' | 'IS_ATTACHMENT_MIGRATED' | 'IS_NOTE_TARGET_MIGRATED' | 'IS_TASK_TARGET_MIGRATED' | 'IS_ROW_LEVEL_PERMISSION_PREDICATES_ENABLED' | 'IS_JUNCTION_RELATIONS_ENABLED' | 'IS_COMMAND_MENU_ITEM_ENABLED' | 'IS_NAVIGATION_MENU_ITEM_ENABLED' | 'IS_DATE_TIME_WHOLE_DAY_FILTER_ENABLED' | 'IS_NAVIGATION_MENU_ITEM_EDITING_ENABLED' | 'IS_DRAFT_EMAIL_ENABLED' | 'IS_USAGE_ANALYTICS_ENABLED' | 'IS_RICH_TEXT_V1_MIGRATED' | 'IS_DIRECT_GRAPHQL_EXECUTION_ENABLED' | 'IS_RECORD_PAGE_LAYOUT_GLOBAL_EDITION_ENABLED' | 'IS_CONNECTED_ACCOUNT_MIGRATED' | 'IS_GRAPHQL_QUERY_TIMING_ENABLED' | 'IS_RECORD_TABLE_WIDGET_ENABLED'
+export type FeatureFlagKey = 'IS_UNIQUE_INDEXES_ENABLED' | 'IS_JSON_FILTER_ENABLED' | 'IS_AI_ENABLED' | 'IS_APPLICATION_ENABLED' | 'IS_MARKETPLACE_ENABLED' | 'IS_RECORD_PAGE_LAYOUT_EDITING_ENABLED' | 'IS_PUBLIC_DOMAIN_ENABLED' | 'IS_EMAILING_DOMAIN_ENABLED' | 'IS_DASHBOARD_V2_ENABLED' | 'IS_ATTACHMENT_MIGRATED' | 'IS_NOTE_TARGET_MIGRATED' | 'IS_TASK_TARGET_MIGRATED' | 'IS_ROW_LEVEL_PERMISSION_PREDICATES_ENABLED' | 'IS_JUNCTION_RELATIONS_ENABLED' | 'IS_COMMAND_MENU_ITEM_ENABLED' | 'IS_NAVIGATION_MENU_ITEM_ENABLED' | 'IS_DATE_TIME_WHOLE_DAY_FILTER_ENABLED' | 'IS_NAVIGATION_MENU_ITEM_EDITING_ENABLED' | 'IS_DRAFT_EMAIL_ENABLED' | 'IS_USAGE_ANALYTICS_ENABLED' | 'IS_RICH_TEXT_V1_MIGRATED' | 'IS_DIRECT_GRAPHQL_EXECUTION_ENABLED' | 'IS_RECORD_PAGE_LAYOUT_GLOBAL_EDITION_ENABLED' | 'IS_CONNECTED_ACCOUNT_MIGRATED' | 'IS_GRAPHQL_QUERY_TIMING_ENABLED' | 'IS_RECORD_TABLE_WIDGET_ENABLED' | 'IS_DATASOURCE_MIGRATED'

 export interface ClientConfig {
    appVersion?: Scalars['String']
@@ -3656,7 +3655,6 @@ export interface WorkspaceGenqlSelection{
    viewGroups?: ViewGroupGenqlSelection
    viewSorts?: ViewSortGenqlSelection
    metadataVersion?: boolean | number
-    databaseUrl?: boolean | number
    databaseSchema?: boolean | number
    subdomain?: boolean | number
    customDomain?: boolean | number
@@ -9159,7 +9157,8 @@ export const enumFeatureFlagKey = {
   IS_RECORD_PAGE_LAYOUT_GLOBAL_EDITION_ENABLED: 'IS_RECORD_PAGE_LAYOUT_GLOBAL_EDITION_ENABLED' as const,
   IS_CONNECTED_ACCOUNT_MIGRATED: 'IS_CONNECTED_ACCOUNT_MIGRATED' as const,
   IS_GRAPHQL_QUERY_TIMING_ENABLED: 'IS_GRAPHQL_QUERY_TIMING_ENABLED' as const,
-   IS_RECORD_TABLE_WIDGET_ENABLED: 'IS_RECORD_TABLE_WIDGET_ENABLED' as const
+   IS_RECORD_TABLE_WIDGET_ENABLED: 'IS_RECORD_TABLE_WIDGET_ENABLED' as const,
+   IS_DATASOURCE_MIGRATED: 'IS_DATASOURCE_MIGRATED' as const
 }

 export const enumConfigSource = {
@@ -1637,9 +1637,6 @@ export default {
            "metadataVersion": [
                11
            ],
-            "databaseUrl": [
-                1
-            ],
            "databaseSchema": [
                1
            ],
@@ -1707,6 +1707,7 @@ export enum FeatureFlagKey {
  IS_COMMAND_MENU_ITEM_ENABLED = 'IS_COMMAND_MENU_ITEM_ENABLED',
  IS_CONNECTED_ACCOUNT_MIGRATED = 'IS_CONNECTED_ACCOUNT_MIGRATED',
  IS_DASHBOARD_V2_ENABLED = 'IS_DASHBOARD_V2_ENABLED',
+  IS_DATASOURCE_MIGRATED = 'IS_DATASOURCE_MIGRATED',
  IS_DATE_TIME_WHOLE_DAY_FILTER_ENABLED = 'IS_DATE_TIME_WHOLE_DAY_FILTER_ENABLED',
  IS_DIRECT_GRAPHQL_EXECUTION_ENABLED = 'IS_DIRECT_GRAPHQL_EXECUTION_ENABLED',
  IS_DRAFT_EMAIL_ENABLED = 'IS_DRAFT_EMAIL_ENABLED',
@@ -6051,8 +6052,7 @@ export type Workspace = {
  createdAt: Scalars['DateTime'];
  currentBillingSubscription?: Maybe<BillingSubscription>;
  customDomain?: Maybe<Scalars['String']>;
-  databaseSchema: Scalars['String'];
-  databaseUrl: Scalars['String'];
+  databaseSchema?: Maybe<Scalars['String']>;
  defaultRole?: Maybe<Role>;
  deletedAt?: Maybe<Scalars['DateTime']>;
  displayName?: Maybe<Scalars['String']>;
@@ -158,8 +158,7 @@ export const mockCurrentWorkspace = {
    },
  ],
  workspaceMembersCount: 1,
-  databaseSchema: '',
-  databaseUrl: '',
+  databaseSchema: null,
  isTwoFactorAuthenticationEnforced: false,
  eventLogRetentionDays: 90,
  __typename: 'Workspace',
@@ -6,7 +6,7 @@ import { In, MoreThanOrEqual, type Repository } from 'typeorm';

 import { MigrationCommandRunner } from 'src/database/commands/command-runners/migration.command-runner';
 import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
-import { DataSourceService } from 'src/engine/metadata-modules/data-source/data-source.service';
+import { type DataSourceService } from 'src/engine/metadata-modules/data-source/data-source.service';
 import { GlobalWorkspaceDataSource } from 'src/engine/twenty-orm/global-workspace-datasource/global-workspace-datasource';
 import { type GlobalWorkspaceOrmManager } from 'src/engine/twenty-orm/global-workspace-datasource/global-workspace-orm.manager';
 import { buildSystemAuthContext } from 'src/engine/twenty-orm/utils/build-system-auth-context.util';
@@ -0,0 +1,39 @@
+import { type MigrationInterface, type QueryRunner } from 'typeorm';
+
+export class DropWorkspaceDatabaseUrlColumn1774688563000
+  implements MigrationInterface
+{
+  name = 'DropWorkspaceDatabaseUrlColumn1774688563000';
+
+  public async up(queryRunner: QueryRunner): Promise<void> {
+    await queryRunner.query(
+      `ALTER TABLE "core"."workspace" DROP COLUMN IF EXISTS "databaseUrl"`,
+    );
+
+    await queryRunner.query(
+      `UPDATE "core"."workspace" SET "databaseSchema" = NULL WHERE "databaseSchema" = ''`,
+    );
+    await queryRunner.query(
+      `ALTER TABLE "core"."workspace" ALTER COLUMN "databaseSchema" DROP NOT NULL`,
+    );
+    await queryRunner.query(
+      `ALTER TABLE "core"."workspace" ALTER COLUMN "databaseSchema" DROP DEFAULT`,
+    );
+  }
+
+  public async down(queryRunner: QueryRunner): Promise<void> {
+    await queryRunner.query(
+      `UPDATE "core"."workspace" SET "databaseSchema" = '' WHERE "databaseSchema" IS NULL`,
+    );
+    await queryRunner.query(
+      `ALTER TABLE "core"."workspace" ALTER COLUMN "databaseSchema" SET NOT NULL`,
+    );
+    await queryRunner.query(
+      `ALTER TABLE "core"."workspace" ALTER COLUMN "databaseSchema" SET DEFAULT ''`,
+    );
+
+    await queryRunner.query(
+      `ALTER TABLE "core"."workspace" ADD "databaseUrl" character varying NOT NULL DEFAULT ''`,
+    );
+  }
+}
@@ -1,14 +1,17 @@
 import { Injectable } from '@nestjs/common';

 import { makeExecutableSchema } from '@graphql-tools/schema';
+import { isNonEmptyString } from '@sniptt/guards';
 import { GraphQLSchema, printSchema } from 'graphql';
 import { gql } from 'graphql-tag';
+import { FeatureFlagKey } from 'twenty-shared/types';
 import { isDefined } from 'twenty-shared/utils';

 import { ScalarsExplorerService } from 'src/engine/api/graphql/services/scalars-explorer.service';
 import { workspaceResolverBuilderMethodNames } from 'src/engine/api/graphql/workspace-resolver-builder/factories/factories';
 import { WorkspaceResolverFactory } from 'src/engine/api/graphql/workspace-resolver-builder/workspace-resolver.factory';
 import { WorkspaceGraphQLSchemaGenerator } from 'src/engine/api/graphql/workspace-schema-builder/workspace-graphql-schema.factory';
+import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
 import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { DataSourceService } from 'src/engine/metadata-modules/data-source/data-source.service';
 import {
@@ -28,24 +31,33 @@ import { TWENTY_STANDARD_APPLICATION } from 'src/engine/workspace-manager/twenty
@Injectable()
 export class WorkspaceSchemaFactory {
  constructor(
-    private readonly dataSourceService: DataSourceService,
    private readonly scalarsExplorerService: ScalarsExplorerService,
    private readonly workspaceGraphQLSchemaGenerator: WorkspaceGraphQLSchemaGenerator,
    private readonly workspaceResolverFactory: WorkspaceResolverFactory,
    private readonly workspaceCacheStorageService: WorkspaceCacheStorageService,
    private readonly workspaceManyOrAllFlatEntityMapsCacheService: WorkspaceManyOrAllFlatEntityMapsCacheService,
+    private readonly featureFlagService: FeatureFlagService,
+    private readonly dataSourceService: DataSourceService,
  ) {}

  async createGraphQLSchema(
    workspace: WorkspaceEntity,
    applicationId?: string,
  ): Promise<GraphQLSchema> {
-    const dataSourcesMetadata =
-      await this.dataSourceService.getDataSourcesMetadataFromWorkspaceId(
-        workspace.id,
-      );
+    const isDataSourceMigrated = await this.featureFlagService.isFeatureEnabled(
+      FeatureFlagKey.IS_DATASOURCE_MIGRATED,
+      workspace.id,
+    );

-    if (!dataSourcesMetadata || dataSourcesMetadata.length === 0) {
+    const hasSchema = isDataSourceMigrated
+      ? isNonEmptyString(workspace.databaseSchema)
+      : (
+          await this.dataSourceService.getDataSourcesMetadataFromWorkspaceId(
+            workspace.id,
+          )
+        ).length > 0;
+
+    if (!hasSchema) {
      return new GraphQLSchema({});
    }

@@ -16,6 +16,7 @@ import {
  type ListLayerVersionsCommandInput,
  LogType,
  PublishLayerVersionCommand,
+  ResourceConflictException,
  ResourceNotFoundException,
  waitUntilFunctionActiveV2,
 } from '@aws-sdk/client-lambda';
@@ -392,7 +393,13 @@ export class LambdaDriver implements LogicFunctionDriver {
        MemorySize: YARN_INSTALL_LAMBDA_MEMORY_MB,
      };

-      await lambdaClient.send(new CreateFunctionCommand(params));
+      try {
+        await lambdaClient.send(new CreateFunctionCommand(params));
+      } catch (error) {
+        if (!(error instanceof ResourceConflictException)) {
+          throw error;
+        }
+      }
    } finally {
      await temporaryDirManager.clean();
    }
@@ -481,7 +488,13 @@ export class LambdaDriver implements LogicFunctionDriver {
        MemorySize: BUILDER_LAMBDA_MEMORY_MB,
      };

-      await lambdaClient.send(new CreateFunctionCommand(params));
+      try {
+        await lambdaClient.send(new CreateFunctionCommand(params));
+      } catch (error) {
+        if (!(error instanceof ResourceConflictException)) {
+          throw error;
+        }
+      }
    } finally {
      await temporaryDirManager.clean();
    }
@@ -900,8 +913,6 @@ export class LambdaDriver implements LogicFunctionDriver {
      return;
    }

-    await this.delete(flatLogicFunction);
-
    const depsLayerArn = await this.getLayerArn({
      flatApplication,
      applicationUniversalIdentifier,
@@ -912,6 +923,22 @@ export class LambdaDriver implements LogicFunctionDriver {
      applicationUniversalIdentifier,
    });

+    // Re-check after layer resolution: a concurrent execution may have
+    // already rebuilt the function with the correct layers while we were
+    // waiting for ensureSdkLayer. Skipping the delete+create avoids
+    // destroying a function another execution is currently polling.
+    if (
+      await this.isAlreadyBuilt({
+        flatLogicFunction,
+        flatApplication,
+        applicationUniversalIdentifier,
+      })
+    ) {
+      return;
+    }
+
+    await this.delete(flatLogicFunction);
+
    const temporaryDirManager = new TemporaryDirManager();

    const { sourceTemporaryDir, lambdaZipPath } =
@@ -938,7 +965,13 @@ export class LambdaDriver implements LogicFunctionDriver {

      const command = new CreateFunctionCommand(params);

-      await (await this.getLambdaClient()).send(command);
+      try {
+        await (await this.getLambdaClient()).send(command);
+      } catch (error) {
+        if (!(error instanceof ResourceConflictException)) {
+          throw error;
+        }
+      }
    } finally {
      await temporaryDirManager.clean();
    }
@@ -234,13 +234,9 @@ export class WorkspaceEntity {
  @Column({ default: 1 })
  metadataVersion: number;

-  @Field()
-  @Column({ default: '' })
-  databaseUrl: string;
-
-  @Field()
-  @Column({ default: '' })
-  databaseSchema: string;
+  @Field(() => String, { nullable: true })
+  @Column({ type: 'varchar', nullable: true, default: null })
+  databaseSchema: string | null;

  @Field()
  @Column({ unique: true })
@@ -14,6 +14,9 @@ import { WorkspaceRelatedEntity } from 'src/engine/workspace-manager/types/works

 export type DataSourceType = DataSourceOptions['type'];

+// @deprecated - This entity is being deprecated in favor of storing
+// databaseSchema directly on WorkspaceEntity.
+// During the transition, writes go to both tables (dual-write).
@Entity('dataSource')
@Index('IDX_DATA_SOURCE_WORKSPACE_ID_CREATED_AT', ['workspaceId', 'createdAt'])
 export class DataSourceEntity extends WorkspaceRelatedEntity {
@@ -1,11 +1,13 @@
 import { Module } from '@nestjs/common';
 import { TypeOrmModule } from '@nestjs/typeorm';

+import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
+
 import { DataSourceEntity } from './data-source.entity';
 import { DataSourceService } from './data-source.service';

@Module({
-  imports: [TypeOrmModule.forFeature([DataSourceEntity])],
+  imports: [TypeOrmModule.forFeature([DataSourceEntity, WorkspaceEntity])],
  providers: [DataSourceService],
  exports: [DataSourceService],
 })
@@ -3,6 +3,7 @@ import { InjectRepository } from '@nestjs/typeorm';

 import { type FindManyOptions, Repository } from 'typeorm';

+import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import {
  DataSourceException,
  DataSourceExceptionCode,
@@ -10,22 +11,32 @@ import {

 import { DataSourceEntity } from './data-source.entity';

+// @deprecated - This service is being deprecated. During the transition,
+// writes go to both the dataSource table and workspace table (dual-write).
+// Reads should progressively migrate to use workspace.databaseSchema
+// or the deterministic getWorkspaceSchemaName(workspaceId) utility.
@Injectable()
 export class DataSourceService {
  constructor(
    @InjectRepository(DataSourceEntity)
    private readonly dataSourceMetadataRepository: Repository<DataSourceEntity>,
+    @InjectRepository(WorkspaceEntity)
+    private readonly workspaceRepository: Repository<WorkspaceEntity>,
  ) {}

  async createDataSourceMetadata(
    workspaceId: string,
    workspaceSchema: string,
  ): Promise<DataSourceEntity> {
-    // TODO: Double check if this is the correct way to do this
    const dataSource = await this.dataSourceMetadataRepository.findOne({
      where: { workspaceId },
    });

+    // Dual-write: always keep workspace.databaseSchema in sync
+    await this.workspaceRepository.update(workspaceId, {
+      databaseSchema: workspaceSchema,
+    });
+
    if (dataSource) {
      return dataSource;
    }
@@ -36,12 +47,16 @@ export class DataSourceService {
    });
  }

+  // @deprecated - Use workspace.activationStatus or workspace.databaseSchema
+  // to check if a workspace has been initialized instead.
  async getManyDataSourceMetadata(
    options: FindManyOptions<DataSourceEntity> = {},
  ): Promise<DataSourceEntity[]> {
    return this.dataSourceMetadataRepository.find(options);
  }

+  // @deprecated - Use workspace.databaseSchema or
+  // getWorkspaceSchemaName(workspaceId) instead.
  async getDataSourcesMetadataFromWorkspaceId(
    workspaceId: string,
  ): Promise<DataSourceEntity[]> {
@@ -51,6 +66,8 @@ export class DataSourceService {
    });
  }

+  // @deprecated - Use workspace.databaseSchema or
+  // getWorkspaceSchemaName(workspaceId) instead.
  async getLastDataSourceMetadataFromWorkspaceId(
    workspaceId: string,
  ): Promise<DataSourceEntity | null> {
@@ -60,6 +77,8 @@ export class DataSourceService {
    });
  }

+  // @deprecated - Use workspace.databaseSchema or
+  // getWorkspaceSchemaName(workspaceId) instead.
  async getLastDataSourceMetadataFromWorkspaceIdOrFail(
    workspaceId: string,
  ): Promise<DataSourceEntity> {
@@ -78,5 +97,10 @@ export class DataSourceService {

  async delete(workspaceId: string): Promise<void> {
    await this.dataSourceMetadataRepository.delete({ workspaceId });
+
+    // Dual-write: clear workspace.databaseSchema on delete
+    await this.workspaceRepository.update(workspaceId, {
+      databaseSchema: null,
+    });
  }
 }
@@ -17,7 +17,6 @@ export const buildDefaultIndexesForCustomObject = ({
  defaultFlatFieldForCustomObjectMaps: DefaultFlatFieldForCustomObjectMaps;
 }) => {
  const tsFlatVectorIndexUniversalIdentifier = v4();
-  const positionIndexUniversalIdentifier = v4();
  const createdAt = new Date();
  const tsVectorFlatIndex = generateFlatIndexMetadataWithNameOrThrow({
    objectFlatFieldMetadatas,
@@ -49,39 +48,9 @@ export const buildDefaultIndexesForCustomObject = ({
    flatObjectMetadata,
  });

-  const positionFlatIndex = generateFlatIndexMetadataWithNameOrThrow({
-    objectFlatFieldMetadatas,
-    flatIndex: {
-      createdAt: createdAt.toISOString(),
-      universalFlatIndexFieldMetadatas: [
-        {
-          createdAt: createdAt.toISOString(),
-          fieldMetadataUniversalIdentifier:
-            defaultFlatFieldForCustomObjectMaps.fields.position
-              .universalIdentifier,
-          indexMetadataUniversalIdentifier: positionIndexUniversalIdentifier,
-          order: 0,
-          updatedAt: createdAt.toISOString(),
-        },
-      ],
-
-      indexType: IndexType.BTREE,
-      indexWhereClause: null,
-      isCustom: false,
-      isUnique: false,
-      objectMetadataUniversalIdentifier: flatObjectMetadata.universalIdentifier,
-      universalIdentifier: positionIndexUniversalIdentifier,
-      updatedAt: createdAt.toISOString(),
-      applicationUniversalIdentifier:
-        flatObjectMetadata.applicationUniversalIdentifier,
-    },
-    flatObjectMetadata,
-  });
-
  return {
    indexes: {
      tsVectorFlatIndex,
-      positionFlatIndex,
    },
  } as const satisfies { indexes: Record<string, UniversalFlatIndexMetadata> };
 };
@@ -1,6 +1,7 @@
 import { Module } from '@nestjs/common';

 import { TokenModule } from 'src/engine/core-modules/auth/token/token.module';
+import { FeatureFlagModule } from 'src/engine/core-modules/feature-flag/feature-flag.module';
 import { JwtModule } from 'src/engine/core-modules/jwt/jwt.module';
 import { DataSourceModule } from 'src/engine/metadata-modules/data-source/data-source.module';
 import { WorkspaceManyOrAllFlatEntityMapsCacheModule } from 'src/engine/metadata-modules/flat-entity/services/workspace-many-or-all-flat-entity-maps-cache.module';
@@ -10,6 +11,7 @@ import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/
@Module({
  imports: [
    DataSourceModule,
+    FeatureFlagModule,
    WorkspaceCacheStorageModule,
    WorkspaceManyOrAllFlatEntityMapsCacheModule,
    TokenModule,
@@ -1,6 +1,8 @@
 import { Injectable } from '@nestjs/common';

+import { isNonEmptyString } from '@sniptt/guards';
 import { type Request, type Response } from 'express';
+import { FeatureFlagKey } from 'twenty-shared/types';
 import { type APP_LOCALES, SOURCE_LOCALE } from 'twenty-shared/translations';
 import { isDefined } from 'twenty-shared/utils';

@@ -9,6 +11,7 @@ import { AuthGraphqlApiExceptionFilter } from 'src/engine/core-modules/auth/filt
 import { AccessTokenService } from 'src/engine/core-modules/auth/token/services/access-token.service';
 import { getAuthExceptionRestStatus } from 'src/engine/core-modules/auth/utils/get-auth-exception-rest-status.util';
 import { ExceptionHandlerService } from 'src/engine/core-modules/exception-handler/exception-handler.service';
+import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
 import { ErrorCode } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
 import { JwtWrapperService } from 'src/engine/core-modules/jwt/services/jwt-wrapper.service';
 import { DataSourceService } from 'src/engine/metadata-modules/data-source/data-source.service';
@@ -29,6 +32,7 @@ export class MiddlewareService {
    private readonly workspaceStorageCacheService: WorkspaceCacheStorageService,
    private readonly flatEntityMapsCacheService: WorkspaceManyOrAllFlatEntityMapsCacheService,
    private readonly dataSourceService: DataSourceService,
+    private readonly featureFlagService: FeatureFlagService,
    private readonly exceptionHandlerService: ExceptionHandlerService,
    private readonly jwtWrapperService: JwtWrapperService,
  ) {}
@@ -106,13 +110,24 @@ export class MiddlewareService {
        )
      : undefined;

-    const dataSourcesMetadata = data.workspace
-      ? await this.dataSourceService.getDataSourcesMetadataFromWorkspaceId(
-          data.workspace.id,
-        )
-      : undefined;
+    if (!data.workspace) {
+      throw new Error('No data sources found');
+    }

-    if (!dataSourcesMetadata || dataSourcesMetadata.length === 0) {
+    const isDataSourceMigrated = await this.featureFlagService.isFeatureEnabled(
+      FeatureFlagKey.IS_DATASOURCE_MIGRATED,
+      data.workspace.id,
+    );
+
+    const hasSchema = isDataSourceMigrated
+      ? isNonEmptyString(data.workspace.databaseSchema)
+      : (
+          await this.dataSourceService.getDataSourcesMetadataFromWorkspaceId(
+            data.workspace.id,
+          )
+        ).length > 0;
+
+    if (!hasSchema) {
      throw new Error('No data sources found');
    }

@@ -256,6 +256,7 @@ describe('WorkspaceEntityManager', () => {
        IS_CONNECTED_ACCOUNT_MIGRATED: false,
        IS_GRAPHQL_QUERY_TIMING_ENABLED: false,
        IS_RECORD_TABLE_WIDGET_ENABLED: false,
+        IS_DATASOURCE_MIGRATED: false,
      },
      userWorkspaceRoleMap: {},
      eventEmitterService: {
@@ -1,12 +1,20 @@
 import { Module } from '@nestjs/common';
+import { TypeOrmModule } from '@nestjs/typeorm';

 import { TypeORMModule } from 'src/database/typeorm/typeorm.module';
+import { FeatureFlagModule } from 'src/engine/core-modules/feature-flag/feature-flag.module';
+import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { DataSourceModule } from 'src/engine/metadata-modules/data-source/data-source.module';

 import { WorkspaceDataSourceService } from './workspace-datasource.service';

@Module({
-  imports: [DataSourceModule, TypeORMModule],
+  imports: [
+    TypeOrmModule.forFeature([WorkspaceEntity]),
+    TypeORMModule,
+    FeatureFlagModule,
+    DataSourceModule,
+  ],
  exports: [WorkspaceDataSourceService],
  providers: [WorkspaceDataSourceService],
 })
@@ -1,9 +1,13 @@
 import { Injectable } from '@nestjs/common';
-import { InjectDataSource } from '@nestjs/typeorm';
+import { InjectDataSource, InjectRepository } from '@nestjs/typeorm';

 import { msg } from '@lingui/core/macro';
-import { type DataSource, type EntityManager } from 'typeorm';
+import { isNonEmptyString } from '@sniptt/guards';
+import { FeatureFlagKey } from 'twenty-shared/types';
+import { type DataSource, type EntityManager, Repository } from 'typeorm';

+import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
+import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { DataSourceService } from 'src/engine/metadata-modules/data-source/data-source.service';
 import {
  PermissionsException,
@@ -14,18 +18,35 @@ import { getWorkspaceSchemaName } from 'src/engine/workspace-datasource/utils/ge
@Injectable()
 export class WorkspaceDataSourceService {
  constructor(
-    private readonly dataSourceService: DataSourceService,
+    @InjectRepository(WorkspaceEntity)
+    private readonly workspaceRepository: Repository<WorkspaceEntity>,
    @InjectDataSource()
    private readonly coreDataSource: DataSource,
+    private readonly featureFlagService: FeatureFlagService,
+    private readonly dataSourceService: DataSourceService,
  ) {}

  public async checkSchemaExists(workspaceId: string) {
-    const dataSource =
+    const isDataSourceMigrated = await this.featureFlagService.isFeatureEnabled(
+      FeatureFlagKey.IS_DATASOURCE_MIGRATED,
+      workspaceId,
+    );
+
+    if (isDataSourceMigrated) {
+      const workspace = await this.workspaceRepository.findOne({
+        select: ['databaseSchema'],
+        where: { id: workspaceId },
+      });
+
+      return isNonEmptyString(workspace?.databaseSchema);
+    }
+
+    const dataSources =
      await this.dataSourceService.getDataSourcesMetadataFromWorkspaceId(
        workspaceId,
      );

-    return dataSource.length > 0;
+    return dataSources.length > 0;
  }

  /**
@@ -25,4 +25,5 @@ export enum FeatureFlagKey {
  IS_CONNECTED_ACCOUNT_MIGRATED = 'IS_CONNECTED_ACCOUNT_MIGRATED',
  IS_GRAPHQL_QUERY_TIMING_ENABLED = 'IS_GRAPHQL_QUERY_TIMING_ENABLED',
  IS_RECORD_TABLE_WIDGET_ENABLED = 'IS_RECORD_TABLE_WIDGET_ENABLED',
+  IS_DATASOURCE_MIGRATED = 'IS_DATASOURCE_MIGRATED',
 }