feat: multi-workspace followup (twentyhq#4197)

* Seed UserWorkspace for existing demo/dev users * add workspaces field to currentUser * new token generation endpoint for switching workspace * lint fix * include dependency * requested fixes * resolver test pass * changing defaultWorkspace and workspaceMember when switching workspaces * tests fix * requested changes * delete user/workspace edge case handled * after merge * requested changes * :wq! * workspace manytoone relation * lint fix / import fix * gql codegen * Fix migrations and generateJWT * migration fix * relations fix --------- Co-authored-by: martmull <martmull@hotmail.fr>
arnavsaxena17 · Mar 4, 2024 · 401375d · 401375d
1 parent 2813c01
commit 401375d
Show file tree

Hide file tree

Showing 25 changed files with 363 additions and 22 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -4,21 +4,21 @@
   "[typescript]": {
     "editor.formatOnSave": false,
     "editor.codeActionsOnSave": {
-      "source.fixAll.eslint": true,
+      "source.fixAll.eslint": "explicit",
       "source.addMissingImports": "always"
     }
   },
   "[javascript]": {
     "editor.formatOnSave": false,
     "editor.codeActionsOnSave": {
-      "source.fixAll.eslint": true,
+      "source.fixAll.eslint": "explicit",
       "source.addMissingImports": "always"
     }
   },
   "[typescriptreact]": {
     "editor.formatOnSave": false,
     "editor.codeActionsOnSave": {
-      "source.fixAll.eslint": true,
+      "source.fixAll.eslint": "explicit",
       "source.addMissingImports": "always"
     }
   },

diff --git a/packages/twenty-front/src/generated/graphql.tsx b/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?: Maybe<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
+      }
+    }
   }
 }
     `;

diff --git a/packages/twenty-front/src/modules/users/graphql/queries/getCurrentUser.ts b/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
+        }
+      }
     }
   }
 `;
diff --git a/packages/twenty-server/src/core/auth/auth.resolver.spec.ts b/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();
 

diff --git a/packages/twenty-server/src/core/auth/auth.resolver.ts b/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) {

diff --git a/packages/twenty-server/src/core/auth/dto/generate-jwt.input.ts b/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;
+}
diff --git a/packages/twenty-server/src/core/auth/services/auth.service.ts b/packages/twenty-server/src/core/auth/services/auth.service.ts
@@ -98,7 +98,7 @@ export class AuthService {
       where: {
         email,
       },
-      relations: ['defaultWorkspace'],
+      relations: ['defaultWorkspace', 'workspaces', 'workspaces.workspace'],
     });
 
     assert(user, "This user doesn't exist", NotFoundException);

diff --git a/packages/twenty-server/src/core/auth/services/token.service.spec.ts b/packages/twenty-server/src/core/auth/services/token.service.spec.ts
@@ -7,6 +7,8 @@ import { RefreshToken } from 'src/core/refresh-token/refresh-token.entity';
 import { User } from 'src/core/user/user.entity';
 import { JwtAuthStrategy } from 'src/core/auth/strategies/jwt.auth.strategy';
 import { EmailService } from 'src/integrations/email/email.service';
+import { UserWorkspaceService } from 'src/core/user-workspace/user-workspace.service';
+import { Workspace } from 'src/core/workspace/workspace.entity';
 
 import { TokenService } from './token.service';
 
@@ -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();
 

diff --git a/packages/twenty-server/src/core/auth/services/token.service.ts b/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,45 @@ 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.findOne({
+      where: { id: workspaceId },
+      relations: ['workspaceUsers'],
+    });
+
+    assert(workspace, 'workspace doesnt exist', NotFoundException);
+
+    assert(
+      workspace.workspaceUsers
+        .map((userWorkspace) => userWorkspace.userId)
+        .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();

diff --git a/packages/twenty-server/src/core/user-workspace/user-workspace.entity.ts b/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.workspaceUsers, {
+    onDelete: 'CASCADE',
+  })
   @JoinColumn({ name: 'workspaceId' })
   workspace: Workspace;
 
+  @Field({ nullable: false })
   @Column()
   workspaceId: string;
 

diff --git a/packages/twenty-server/src/core/user-workspace/user-workspace.service.ts b/packages/twenty-server/src/core/user-workspace/user-workspace.service.ts
@@ -81,14 +81,6 @@ export class UserWorkspaceService extends TypeOrmQueryService<UserWorkspace> {
     ).length;
   }
 
-  async findUserWorkspaces(userId: string): Promise<UserWorkspace[]> {
-    return this.userWorkspaceRepository.find({
-      where: {
-        userId,
-      },
-    });
-  }
-
   async checkUserWorkspaceExists(
     userId: string,
     workspaceId: string,