feat: multi-workspace followup

packages/twenty-front/src/generated/graphql.tsx

@@ -235,6 +235,7 @@ export type Mutation = {
   deleteUser: User;
   emailPasswordResetLink: EmailPasswordResetLink;
   generateApiKeyToken: ApiKeyToken;
+  generateJWT: AuthTokens;
   generateTransientToken: TransientToken;
   impersonate: Verify;
   renewToken: AuthTokens;
@@ -289,6 +290,11 @@ export type MutationGenerateApiKeyTokenArgs = {
 };
 
 
+export type MutationGenerateJwtArgs = {
+  workspaceId: Scalars['String'];
+};
+
+
 export type MutationImpersonateArgs = {
   userId: Scalars['String'];
 };
@@ -568,6 +574,7 @@ export type User = {
   createdAt: Scalars['DateTime'];
   defaultAvatarUrl?: Maybe<Scalars['String']>;
   defaultWorkspace: Workspace;
+  defaultWorkspaceId: Scalars['String'];
   deletedAt?: Maybe<Scalars['DateTime']>;
   disabled?: Maybe<Scalars['Boolean']>;
   email: Scalars['String'];
@@ -581,6 +588,7 @@ export type User = {
   supportUserHash?: Maybe<Scalars['String']>;
   updatedAt: Scalars['DateTime'];
   workspaceMember?: Maybe<WorkspaceMember>;
+  workspaces: Array<UserWorkspace>;
 };
 
 export type UserEdge = {
@@ -596,6 +604,18 @@ export type UserExists = {
   exists: Scalars['Boolean'];
 };
 
+export type UserWorkspace = {
+  __typename?: 'UserWorkspace';
+  createdAt: Scalars['DateTime'];
+  deletedAt: Scalars['DateTime'];
+  id: Scalars['ID'];
+  updatedAt: Scalars['DateTime'];
+  user: User;
+  userId: Scalars['String'];
+  workspace?: Maybe<Workspace>;
+  workspaceId: Scalars['String'];
+};
+
 export type ValidatePasswordResetToken = {
   __typename?: 'ValidatePasswordResetToken';
   email: Scalars['String'];
@@ -916,7 +936,7 @@ export type UploadProfilePictureMutation = { __typename?: 'Mutation', uploadProf
 export type GetCurrentUserQueryVariables = Exact<{ [key: string]: never; }>;
 
 
-export type GetCurrentUserQuery = { __typename?: 'Query', currentUser: { __typename?: 'User', id: string, firstName: string, lastName: string, email: string, canImpersonate: boolean, supportUserHash?: string | null, workspaceMember?: { __typename?: 'WorkspaceMember', id: string, colorScheme: string, avatarUrl?: string | null, locale: string, name: { __typename?: 'FullName', firstName: string, lastName: string } } | null, defaultWorkspace: { __typename?: 'Workspace', id: string, displayName?: string | null, logo?: string | null, domainName?: string | null, inviteHash?: string | null, allowImpersonation: boolean, subscriptionStatus: string, activationStatus: string, featureFlags?: Array<{ __typename?: 'FeatureFlag', id: string, key: string, value: boolean, workspaceId: string }> | null } } };
+export type GetCurrentUserQuery = { __typename?: 'Query', currentUser: { __typename?: 'User', id: string, firstName: string, lastName: string, email: string, canImpersonate: boolean, supportUserHash?: string | null, workspaceMember?: { __typename?: 'WorkspaceMember', id: string, colorScheme: string, avatarUrl?: string | null, locale: string, name: { __typename?: 'FullName', firstName: string, lastName: string } } | null, defaultWorkspace: { __typename?: 'Workspace', id: string, displayName?: string | null, logo?: string | null, domainName?: string | null, inviteHash?: string | null, allowImpersonation: boolean, subscriptionStatus: string, activationStatus: string, featureFlags?: Array<{ __typename?: 'FeatureFlag', id: string, key: string, value: boolean, workspaceId: string }> | null }, workspaces: Array<{ __typename?: 'UserWorkspace', workspace?: { __typename?: 'Workspace', id: string, displayName?: string | null, logo?: string | null, domainName?: string | null } | null }> } };
 
 export type ActivateWorkspaceMutationVariables = Exact<{
   input: ActivateWorkspaceInput;
@@ -1854,6 +1874,14 @@ export const GetCurrentUserDocument = gql`
         workspaceId
       }
     }
+    workspaces {
+      workspace {
+        id
+        displayName
+        logo
+        domainName
+      }
+    }
   }
 }
     `;

packages/twenty-front/src/modules/users/graphql/queries/getCurrentUser.ts

@@ -36,6 +36,14 @@ export const GET_CURRENT_USER = gql`
           workspaceId
         }
       }
+      workspaces {
+        workspace {
+          id
+          displayName
+          logo
+          domainName
+        }
+      }
     }
   }
 `;

packages/twenty-server/src/core/auth/auth.resolver.spec.ts

@@ -3,6 +3,8 @@ import { getRepositoryToken } from '@nestjs/typeorm';
 
 import { Workspace } from 'src/core/workspace/workspace.entity';
 import { UserService } from 'src/core/user/services/user.service';
+import { UserWorkspaceService } from 'src/core/user-workspace/user-workspace.service';
+import { User } from 'src/core/user/user.entity';
 
 import { AuthResolver } from './auth.resolver';
 
@@ -20,6 +22,10 @@ describe('AuthResolver', () => {
           provide: getRepositoryToken(Workspace, 'core'),
           useValue: {},
         },
+        {
+          provide: getRepositoryToken(User, 'core'),
+          useValue: {},
+        },
         {
           provide: AuthService,
           useValue: {},
@@ -32,6 +38,10 @@ describe('AuthResolver', () => {
           provide: UserService,
           useValue: {},
         },
+        {
+          provide: UserWorkspaceService,
+          useValue: {},
+        },
       ],
     }).compile();
 

packages/twenty-server/src/core/auth/auth.resolver.ts

@@ -25,6 +25,8 @@ import { UpdatePasswordViaResetTokenInput } from 'src/core/auth/dto/update-passw
 import { EmailPasswordResetLink } from 'src/core/auth/dto/email-password-reset-link.entity';
 import { InvalidatePassword } from 'src/core/auth/dto/invalidate-password.entity';
 import { EmailPasswordResetLinkInput } from 'src/core/auth/dto/email-password-reset-link.input';
+import { GenerateJwtInput } from 'src/core/auth/dto/generate-jwt.input';
+import { UserWorkspaceService } from 'src/core/user-workspace/user-workspace.service';
 
 import { ApiKeyToken, AuthTokens } from './dto/token.entity';
 import { TokenService } from './services/token.service';
@@ -49,6 +51,7 @@ export class AuthResolver {
     private authService: AuthService,
     private tokenService: TokenService,
     private userService: UserService,
+    private userWorkspaceService: UserWorkspaceService,
   ) {}
 
   @Query(() => UserExists)
@@ -128,6 +131,20 @@ export class AuthResolver {
     return result;
   }
 
+  @Mutation(() => AuthTokens)
+  @UseGuards(JwtAuthGuard)
+  async generateJWT(
+    @AuthUser() user: User,
+    @Args() args: GenerateJwtInput,
+  ): Promise<AuthTokens> {
+    const token = await this.tokenService.generateSwitchWorkspaceToken(
+      user,
+      args.workspaceId,
+    );
+
+    return token;
+  }
+
   @Mutation(() => AuthTokens)
   async renewToken(@Args() args: RefreshTokenInput): Promise<AuthTokens> {
     if (!args.refreshToken) {

packages/twenty-server/src/core/auth/dto/generate-jwt.input.ts

@@ -0,0 +1,11 @@
+import { ArgsType, Field } from '@nestjs/graphql';
+
+import { IsNotEmpty, IsString } from 'class-validator';
+
+@ArgsType()
+export class GenerateJwtInput {
+  @Field(() => String)
+  @IsNotEmpty()
+  @IsString()
+  workspaceId: string;
+}

packages/twenty-server/src/core/auth/services/auth.service.ts

@@ -98,7 +98,7 @@ export class AuthService {
       where: {
         email,
       },
-      relations: ['defaultWorkspace'],
+      relations: ['defaultWorkspace', 'workspaces'],
     });
 
     assert(user, "This user doesn't exist", NotFoundException);

packages/twenty-server/src/core/auth/services/token.service.spec.ts

@@ -9,6 +9,8 @@ import { JwtAuthStrategy } from 'src/core/auth/strategies/jwt.auth.strategy';
 import { EmailService } from 'src/integrations/email/email.service';
 
 import { TokenService } from './token.service';
+import { UserWorkspaceService } from 'src/core/user-workspace/user-workspace.service';
+import { Workspace } from 'src/core/workspace/workspace.entity';
 
 describe('TokenService', () => {
   let service: TokenService;
@@ -33,6 +35,10 @@ describe('TokenService', () => {
           provide: EmailService,
           useValue: {},
         },
+        {
+          provide: UserWorkspaceService,
+          useValue: {},
+        },
         {
           provide: getRepositoryToken(User, 'core'),
           useValue: {},
@@ -41,6 +47,10 @@ describe('TokenService', () => {
           provide: getRepositoryToken(RefreshToken, 'core'),
           useValue: {},
         },
+        {
+          provide: getRepositoryToken(Workspace, 'core'),
+          useValue: {},
+        },
       ],
     }).compile();
 

packages/twenty-server/src/core/auth/services/token.service.ts

@@ -29,6 +29,7 @@ import { assert } from 'src/utils/assert';
 import {
   ApiKeyToken,
   AuthToken,
+  AuthTokens,
   PasswordResetToken,
 } from 'src/core/auth/dto/token.entity';
 import { EnvironmentService } from 'src/integrations/environment/environment.service';
@@ -39,6 +40,8 @@ import { EmailService } from 'src/integrations/email/email.service';
 import { InvalidatePassword } from 'src/core/auth/dto/invalidate-password.entity';
 import { EmailPasswordResetLink } from 'src/core/auth/dto/email-password-reset-link.entity';
 import { JwtData } from 'src/core/auth/types/jwt-data.type';
+import { UserWorkspaceService } from 'src/core/user-workspace/user-workspace.service';
+import { Workspace } from 'src/core/workspace/workspace.entity';
 
 @Injectable()
 export class TokenService {
@@ -50,10 +53,16 @@ export class TokenService {
     private readonly userRepository: Repository<User>,
     @InjectRepository(RefreshToken, 'core')
     private readonly refreshTokenRepository: Repository<RefreshToken>,
+    @InjectRepository(Workspace, 'core')
+    private readonly workspaceRepository: Repository<Workspace>,
     private readonly emailService: EmailService,
+    private readonly userWorkspaceService: UserWorkspaceService,
   ) {}
 
-  async generateAccessToken(userId: string): Promise<AuthToken> {
+  async generateAccessToken(
+    userId: string,
+    workspaceId?: string,
+  ): Promise<AuthToken> {
     const expiresIn = this.environmentService.getAccessTokenExpiresIn();
 
     assert(expiresIn, '', InternalServerErrorException);
@@ -74,7 +83,7 @@ export class TokenService {
 
     const jwtPayload: JwtPayload = {
       sub: user.id,
-      workspaceId: user.defaultWorkspace.id,
+      workspaceId: workspaceId ? workspaceId : user.defaultWorkspace.id,
     };
 
     return {
@@ -232,6 +241,38 @@ export class TokenService {
     };
   }
 
+  async generateSwitchWorkspaceToken(
+    user: User,
+    workspaceId: string,
+  ): Promise<AuthTokens> {
+    const userExists = await this.userRepository.findBy({ id: user.id });
+
+    assert(userExists, 'User not found', NotFoundException);
+
+    const workspace = await this.workspaceRepository.findOneBy({
+      id: workspaceId,
+    });
+
+    assert(workspace, 'workspace doesnt exist', NotFoundException);
+
+    assert(workspace.users.map((u) => u.id).includes(user.id), 'user does not belong to workspace', ForbiddenException);
+
+    await this.userRepository.save({
+      id: user.id,
+      defaultWorkspace: workspace,
+    });
+
+    const token = await this.generateAccessToken(user.id, workspaceId);
+    const refreshToken = await this.generateRefreshToken(user.id);
+
+    return {
+      tokens: {
+        accessToken: token,
+        refreshToken,
+      },
+    };
+  }
+
   async verifyRefreshToken(refreshToken: string) {
     const secret = this.environmentService.getRefreshTokenSecret();
     const coolDown = this.environmentService.getRefreshTokenCoolDown();

packages/twenty-server/src/core/user-workspace/user-workspace.entity.ts

@@ -6,6 +6,7 @@ import {
   CreateDateColumn,
   Entity,
   JoinColumn,
+  ManyToOne,
   PrimaryGeneratedColumn,
   Unique,
   UpdateDateColumn,
@@ -22,15 +23,25 @@ export class UserWorkspace {
   @PrimaryGeneratedColumn('uuid')
   id: string;
 
+  @Field(() => User)
+  @ManyToOne(() => User, (user) => user.workspaces, {
+    onDelete: 'CASCADE',
+  })
   @JoinColumn({ name: 'userId' })
   user: User;
 
+  @Field({ nullable: false })
   @Column()
   userId: string;
 
+  @Field(() => Workspace, { nullable: true })
+  @ManyToOne(() => Workspace, (workspace) => workspace.users, {
+    onDelete: 'CASCADE',
+  })
   @JoinColumn({ name: 'workspaceId' })
   workspace: Workspace;
 
+  @Field({ nullable: false })
   @Column()
   workspaceId: string;
 

packages/twenty-server/src/core/user-workspace/user-workspace.service.ts

@@ -45,14 +45,6 @@ export class UserWorkspaceService extends TypeOrmQueryService<UserWorkspace> {
     );
   }
 
-  async findUserWorkspaces(userId: string): Promise<UserWorkspace[]> {
-    return this.userWorkspaceRepository.find({
-      where: {
-        userId,
-      },
-    });
-  }
-
   async checkUserWorkspaceExists(
     userId: string,
     workspaceId: string,

packages/twenty-server/src/core/user/services/user.service.spec.ts

@@ -4,6 +4,7 @@ import { getRepositoryToken } from '@nestjs/typeorm';
 import { User } from 'src/core/user/user.entity';
 import { DataSourceService } from 'src/metadata/data-source/data-source.service';
 import { TypeORMService } from 'src/database/typeorm/typeorm.service';
+import { UserWorkspace } from 'src/core/user-workspace/user-workspace.entity';
 
 import { UserService } from './user.service';
 
@@ -18,6 +19,10 @@ describe('UserService', () => {
           provide: getRepositoryToken(User, 'core'),
           useValue: {},
         },
+        {
+          provide: getRepositoryToken(UserWorkspace, 'core'),
+          useValue: {},
+        },
         {
           provide: DataSourceService,
           useValue: {},