Fix workspace export for indirect core entity relations

Include indirectly linked core records in workspace exports such as users and index field metadata

packages/twenty-server/src/database/commands/workspace-export/__tests__/workspace-export.service.spec.ts

@@ -0,0 +1,271 @@
+import { DataSource, type EntityMetadata, Repository } from 'typeorm';
+
+import { WorkspaceExportService } from 'src/database/commands/workspace-export/workspace-export.service';
+
+type MockColumnMetadata = {
+  databaseName: string;
+  isNullable?: boolean;
+  referencedColumn?: MockColumnMetadata;
+};
+
+type MockRelationMetadata = {
+  inverseEntityMetadata: EntityMetadata;
+  isNullable?: boolean;
+  joinColumns: MockColumnMetadata[];
+  inverseRelation?: MockRelationMetadata;
+};
+
+function makeColumn({
+  databaseName,
+  isNullable = false,
+  referencedColumnName,
+}: {
+  databaseName: string;
+  isNullable?: boolean;
+  referencedColumnName?: string;
+}): MockColumnMetadata {
+  return {
+    databaseName,
+    isNullable,
+    referencedColumn: referencedColumnName
+      ? { databaseName: referencedColumnName }
+      : undefined,
+  };
+}
+
+function makeEntityMetadata({
+  tableName,
+  columns = [],
+  manyToOneRelations = [],
+}: {
+  tableName: string;
+  columns?: MockColumnMetadata[];
+  manyToOneRelations?: MockRelationMetadata[];
+}): EntityMetadata {
+  return {
+    tableName,
+    schema: 'core',
+    columns,
+    manyToOneRelations,
+  } as EntityMetadata;
+}
+
+function makeRelationMetadata({
+  inverseEntityMetadata,
+  isNullable = false,
+  joinColumns,
+  inverseRelation,
+}: {
+  inverseEntityMetadata: EntityMetadata;
+  isNullable?: boolean;
+  joinColumns: MockColumnMetadata[];
+  inverseRelation?: MockRelationMetadata;
+}): MockRelationMetadata {
+  return {
+    inverseEntityMetadata,
+    isNullable,
+    joinColumns,
+    inverseRelation,
+  };
+}
+
+function discover(entityMetadatas: EntityMetadata[]): Map<string, string> {
+  const service = new WorkspaceExportService(
+    { entityMetadatas } as DataSource,
+    {} as Repository<any>,
+    {} as Repository<any>,
+  );
+
+  return new Map(
+    (
+      (service as any).discoverWorkspaceScopedEntities() as Array<{
+        entityMetadata: EntityMetadata;
+        whereClause: string;
+      }>
+    ).map(({ entityMetadata, whereClause }) => [
+      entityMetadata.tableName,
+      whereClause,
+    ]),
+  );
+}
+
+describe('WorkspaceExportService', () => {
+  it('discovers direct, forward, and required reverse relations', () => {
+    const workspace = makeEntityMetadata({
+      tableName: 'workspace',
+      columns: [makeColumn({ databaseName: 'id' })],
+    });
+    const user = makeEntityMetadata({
+      tableName: 'user',
+      columns: [makeColumn({ databaseName: 'id' })],
+    });
+    const userWorkspace = makeEntityMetadata({
+      tableName: 'userWorkspace',
+      columns: [
+        makeColumn({ databaseName: 'id' }),
+        makeColumn({ databaseName: 'workspaceId' }),
+        makeColumn({ databaseName: 'userId' }),
+      ],
+    });
+    const indexMetadata = makeEntityMetadata({
+      tableName: 'indexMetadata',
+      columns: [
+        makeColumn({ databaseName: 'id' }),
+        makeColumn({ databaseName: 'workspaceId' }),
+      ],
+    });
+    const indexFieldMetadata = makeEntityMetadata({
+      tableName: 'indexFieldMetadata',
+      columns: [
+        makeColumn({ databaseName: 'id' }),
+        makeColumn({ databaseName: 'indexMetadataId' }),
+      ],
+    });
+
+    userWorkspace.manyToOneRelations = [
+      makeRelationMetadata({
+        inverseEntityMetadata: user,
+        isNullable: true,
+        joinColumns: [
+          makeColumn({
+            databaseName: 'userId',
+            referencedColumnName: 'id',
+          }),
+        ],
+        inverseRelation: {} as MockRelationMetadata,
+      }),
+    ] as EntityMetadata['manyToOneRelations'];
+    indexFieldMetadata.manyToOneRelations = [
+      makeRelationMetadata({
+        inverseEntityMetadata: indexMetadata,
+        joinColumns: [
+          makeColumn({
+            databaseName: 'indexMetadataId',
+            referencedColumnName: 'id',
+          }),
+        ],
+      }),
+    ] as EntityMetadata['manyToOneRelations'];
+
+    const scopedEntities = discover([
+      workspace,
+      user,
+      userWorkspace,
+      indexMetadata,
+      indexFieldMetadata,
+    ]);
+
+    expect(scopedEntities.get('userWorkspace')).toBe('"workspaceId" = $1');
+    expect(scopedEntities.get('user')).toBe(
+      '"id" IN (SELECT "userId" FROM "core"."userWorkspace" WHERE "workspaceId" = $1)',
+    );
+    expect(scopedEntities.get('indexMetadata')).toBe('"workspaceId" = $1');
+    expect(scopedEntities.get('indexFieldMetadata')).toBe(
+      '"indexMetadataId" IN (SELECT "id" FROM "core"."indexMetadata" WHERE "workspaceId" = $1)',
+    );
+  });
+
+  it('does not reverse-include optional bidirectional relations', () => {
+    const workspace = makeEntityMetadata({
+      tableName: 'workspace',
+      columns: [makeColumn({ databaseName: 'id' })],
+    });
+    const user = makeEntityMetadata({
+      tableName: 'user',
+      columns: [makeColumn({ databaseName: 'id' })],
+    });
+    const appToken = makeEntityMetadata({
+      tableName: 'appToken',
+      columns: [
+        makeColumn({ databaseName: 'id' }),
+        makeColumn({ databaseName: 'workspaceId' }),
+        makeColumn({ databaseName: 'userId', isNullable: true }),
+      ],
+    });
+
+    appToken.manyToOneRelations = [
+      makeRelationMetadata({
+        inverseEntityMetadata: user,
+        isNullable: true,
+        joinColumns: [
+          makeColumn({
+            databaseName: 'userId',
+            isNullable: true,
+            referencedColumnName: 'id',
+          }),
+        ],
+        inverseRelation: {} as MockRelationMetadata,
+      }),
+    ] as EntityMetadata['manyToOneRelations'];
+
+    const scopedEntities = discover([workspace, user, appToken]);
+
+    expect(scopedEntities.get('appToken')).toBe('"workspaceId" = $1');
+    expect(scopedEntities.has('user')).toBe(false);
+  });
+
+  it('does not reverse-include unidirectional relations', () => {
+    const workspace = makeEntityMetadata({
+      tableName: 'workspace',
+      columns: [makeColumn({ databaseName: 'id' })],
+    });
+    const billingSubscription = makeEntityMetadata({
+      tableName: 'billingSubscription',
+      columns: [
+        makeColumn({ databaseName: 'id' }),
+        makeColumn({ databaseName: 'workspaceId' }),
+      ],
+    });
+    const billingSubscriptionItem = makeEntityMetadata({
+      tableName: 'billingSubscriptionItem',
+      columns: [
+        makeColumn({ databaseName: 'id' }),
+        makeColumn({ databaseName: 'billingSubscriptionId' }),
+        makeColumn({ databaseName: 'stripeProductId' }),
+      ],
+    });
+    const billingProduct = makeEntityMetadata({
+      tableName: 'billingProduct',
+      columns: [
+        makeColumn({ databaseName: 'id' }),
+        makeColumn({ databaseName: 'stripeProductId' }),
+      ],
+    });
+
+    billingSubscriptionItem.manyToOneRelations = [
+      makeRelationMetadata({
+        inverseEntityMetadata: billingSubscription,
+        joinColumns: [
+          makeColumn({
+            databaseName: 'billingSubscriptionId',
+            referencedColumnName: 'id',
+          }),
+        ],
+      }),
+      makeRelationMetadata({
+        inverseEntityMetadata: billingProduct,
+        joinColumns: [
+          makeColumn({
+            databaseName: 'stripeProductId',
+            referencedColumnName: 'stripeProductId',
+          }),
+        ],
+      }),
+    ] as EntityMetadata['manyToOneRelations'];
+
+    const scopedEntities = discover([
+      workspace,
+      billingSubscription,
+      billingSubscriptionItem,
+      billingProduct,
+    ]);
+
+    expect(scopedEntities.get('billingSubscription')).toBe(
+      '"workspaceId" = $1',
+    );
+    expect(scopedEntities.get('billingSubscriptionItem')).toBe(
+      '"billingSubscriptionId" IN (SELECT "id" FROM "core"."billingSubscription" WHERE "workspaceId" = $1)',
+    );
+    expect(scopedEntities.has('billingProduct')).toBe(false);
+  });
+});

packages/twenty-server/src/database/commands/workspace-export/utils/get-core-entity-metadatas-with-workspace-id.util.ts

@@ -1,11 +0,0 @@
-import { DataSource } from 'typeorm';
-
-export const getCoreEntityMetadatasWithWorkspaceId = (
-  dataSource: DataSource,
-) => {
-  return dataSource.entityMetadatas.filter((entityMetadata) =>
-    entityMetadata.columns.some(
-      (column) => column.propertyName === 'workspaceId',
-    ),
-  );
-};

packages/twenty-server/src/database/commands/workspace-export/workspace-export.service.ts

@@ -17,7 +17,6 @@ import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.ent
 import { getWorkspaceSchemaName } from 'src/engine/workspace-datasource/utils/get-workspace-schema-name.util';
 import { computeTableName } from 'src/engine/utils/compute-table-name.util';
 import { escapeIdentifier } from 'src/engine/workspace-manager/workspace-migration/utils/remove-sql-injection.util';
-import { getCoreEntityMetadatasWithWorkspaceId } from 'src/database/commands/workspace-export/utils/get-core-entity-metadatas-with-workspace-id.util';
 import { generateWorkspaceSchemaDdl } from 'src/database/commands/workspace-export/utils/generate-workspace-schema-ddl.util';
 import { buildInsertPrefix } from 'src/database/commands/workspace-export/utils/build-insert-prefix.util';
 import { buildWorkspaceTableColumnSets } from 'src/database/commands/workspace-export/utils/build-workspace-table-column-sets.util';
@@ -32,22 +31,23 @@ type WorkspaceExportParams = {
   tableFilter?: string[];
 };
 
-type RowFilter = {
-  filterColumn: string;
-  filterValue: string;
-};
-
 type WriteRowsOptions = {
   schemaName: string;
   tableName: string;
   displayName: string;
   queryRunner: QueryRunner;
   stream: WriteStream;
-  rowFilter?: RowFilter;
+  whereClause: string;
+  queryParameters: unknown[];
   jsonColumns?: Set<string>;
   excludedColumns?: Set<string>;
 };
 
+type WorkspaceScopedEntity = {
+  entityMetadata: EntityMetadata;
+  whereClause: string;
+};
+
 @Injectable()
 export class WorkspaceExportService {
   private readonly logger = new Logger(WorkspaceExportService.name);
@@ -113,7 +113,7 @@ export class WorkspaceExportService {
         `\nCREATE SCHEMA IF NOT EXISTS ${escapeIdentifier(schemaName)};\n\n`,
       );
 
-      this.writeWorkspaceSchemaDdl(
+      this.writeWorkspaceSchemaDDL(
         workspaceId,
         schemaName,
         objectMetadatas,
@@ -153,20 +153,20 @@ export class WorkspaceExportService {
     if (workspaceEntityMetadata) {
       await this.writeRows({
         schemaName: workspaceEntityMetadata.schema || 'core',
-        tableName: workspaceEntityMetadata.tableName,
-        displayName: workspaceEntityMetadata.tableName,
+        tableName: 'workspace',
+        displayName: 'workspace',
         queryRunner,
         stream,
-        rowFilter: { filterColumn: 'id', filterValue: workspaceId },
+        whereClause: '"id" = $1',
+        queryParameters: [workspaceId],
         jsonColumns: this.buildJsonColumnSet(workspaceEntityMetadata),
       });
     }
 
-    const coreEntityMetadatas = getCoreEntityMetadatasWithWorkspaceId(
-      this.dataSource,
-    );
-
-    for (const entityMetadata of coreEntityMetadatas) {
+    for (const {
+      entityMetadata,
+      whereClause,
+    } of this.discoverWorkspaceScopedEntities()) {
       try {
         await this.writeRows({
           schemaName: entityMetadata.schema || 'core',
@@ -174,10 +174,8 @@ export class WorkspaceExportService {
           displayName: entityMetadata.tableName,
           queryRunner,
           stream,
-          rowFilter: {
-            filterColumn: 'workspaceId',
-            filterValue: workspaceId,
-          },
+          whereClause,
+          queryParameters: [workspaceId],
           jsonColumns: this.buildJsonColumnSet(entityMetadata),
         });
       } catch (error) {
@@ -186,6 +184,113 @@ export class WorkspaceExportService {
     }
   }
 
+  private discoverWorkspaceScopedEntities(): WorkspaceScopedEntity[] {
+    const allEntities = this.dataSource.entityMetadatas;
+    const scopedEntities: WorkspaceScopedEntity[] = [];
+    const scopedWhere = new Map<string, string>([['workspace', '"id" = $1']]);
+
+    for (const entityMetadata of allEntities) {
+      if (
+        entityMetadata.columns.some(
+          (column) => column.databaseName === 'workspaceId',
+        )
+      ) {
+        scopedEntities.push({
+          entityMetadata,
+          whereClause: '"workspaceId" = $1',
+        });
+        scopedWhere.set(entityMetadata.tableName, '"workspaceId" = $1');
+      }
+    }
+
+    let changed = true;
+
+    while (changed) {
+      changed = false;
+
+      for (const entityMetadata of allEntities) {
+        if (scopedWhere.has(entityMetadata.tableName)) continue;
+
+        const whereClause = this.resolveWhereClause(
+          entityMetadata,
+          scopedWhere,
+          scopedEntities,
+        );
+
+        if (!whereClause) continue;
+
+        scopedWhere.set(entityMetadata.tableName, whereClause);
+        scopedEntities.push({ entityMetadata, whereClause });
+        changed = true;
+      }
+    }
+
+    return scopedEntities;
+  }
+
+  private resolveWhereClause(
+    entityMetadata: EntityMetadata,
+    scopedWhere: Map<string, string>,
+    scopedEntities: WorkspaceScopedEntity[],
+  ): string | null {
+    for (const relation of entityMetadata.manyToOneRelations) {
+      const parent = relation.inverseEntityMetadata;
+      const parentWhere = scopedWhere.get(parent.tableName);
+
+      if (!parentWhere) continue;
+
+      const joinColumn = relation.joinColumns[0];
+
+      if (!joinColumn?.databaseName) continue;
+
+      return this.buildInSubqueryWhereClause(
+        joinColumn.databaseName,
+        parent.schema || 'core',
+        parent.tableName,
+        joinColumn.referencedColumn?.databaseName ?? 'id',
+        parentWhere,
+      );
+    }
+
+    for (const {
+      entityMetadata: scopedEntity,
+      whereClause,
+    } of scopedEntities) {
+      for (const relation of scopedEntity.manyToOneRelations) {
+        if (relation.inverseEntityMetadata !== entityMetadata) continue;
+        if (!relation.inverseRelation) continue;
+        if (relation.joinColumns.some((joinColumn) => joinColumn.isNullable))
+          continue;
+
+        const joinColumn = relation.joinColumns[0];
+
+        if (!joinColumn?.databaseName) continue;
+
+        return this.buildInSubqueryWhereClause(
+          joinColumn.referencedColumn?.databaseName ?? 'id',
+          scopedEntity.schema || 'core',
+          scopedEntity.tableName,
+          joinColumn.databaseName,
+          whereClause,
+        );
+      }
+    }
+
+    return null;
+  }
+
+  private buildInSubqueryWhereClause(
+    column: string,
+    parentSchema: string,
+    parentTable: string,
+    parentColumn: string,
+    parentWhere: string,
+  ): string {
+    const qualifiedTable = `${escapeIdentifier(parentSchema)}.${escapeIdentifier(parentTable)}`;
+
+    return `${escapeIdentifier(column)} IN (SELECT ${escapeIdentifier(parentColumn)} FROM ${qualifiedTable} WHERE ${parentWhere})`;
+  }
+
   private buildJsonColumnSet(entityMetadata: EntityMetadata): Set<string> {
     return new Set(
       entityMetadata.columns
@@ -200,32 +305,26 @@ export class WorkspaceExportService {
     displayName,
     queryRunner,
     stream,
-    rowFilter,
+    whereClause,
+    queryParameters,
     jsonColumns,
     excludedColumns,
   }: WriteRowsOptions): Promise<void> {
-    const whereClause = rowFilter
-      ? ` WHERE "${rowFilter.filterColumn}" = $1`
-      : '';
-    const queryParameters = rowFilter ? [rowFilter.filterValue] : [];
-
-    const [{ count: totalCount }] = await queryRunner.query(
-      `SELECT COUNT(*)::int as count FROM "${schemaName}"."${tableName}"${whereClause}`,
-      queryParameters,
-    );
-
-    if (totalCount === 0) return;
-
-    this.logger.log(`  ${displayName}: ${totalCount} rows`);
+    const qualifiedTable = `${escapeIdentifier(schemaName)}.${escapeIdentifier(tableName)}`;
 
     let insertPrefix: string | undefined;
+    let totalCount = 0;
 
-    for (let offset = 0; offset < totalCount; offset += BATCH_SIZE) {
+    for (let offset = 0; ; offset += BATCH_SIZE) {
       const rows: Record<string, unknown>[] = await queryRunner.query(
-        `SELECT * FROM "${schemaName}"."${tableName}"${whereClause} ORDER BY "id" LIMIT ${BATCH_SIZE} OFFSET ${offset}`,
+        `SELECT * FROM ${qualifiedTable} WHERE ${whereClause} ORDER BY "id" LIMIT ${BATCH_SIZE} OFFSET ${offset}`,
         queryParameters,
       );
 
+      if (rows.length === 0) break;
+
+      totalCount += rows.length;
+
       const batchStatements: string[] = [];
 
       for (const row of rows) {
@@ -249,10 +348,16 @@ export class WorkspaceExportService {
       if (!stream.write(batchStatements.join(''))) {
         await once(stream, 'drain');
       }
+
+      if (rows.length < BATCH_SIZE) break;
+    }
+
+    if (totalCount > 0) {
+      this.logger.log(`  ${displayName}: ${totalCount} rows`);
     }
   }
 
-  private writeWorkspaceSchemaDdl(
+  private writeWorkspaceSchemaDDL(
     workspaceId: string,
     schemaName: string,
     objectMetadatas: ObjectMetadataEntity[],
@@ -314,6 +419,8 @@ export class WorkspaceExportService {
           displayName: objectMetadata.nameSingular,
           queryRunner,
           stream,
+          whereClause: 'TRUE',
+          queryParameters: [],
           jsonColumns,
           excludedColumns: generatedColumns,
         });