feat(server): workspace model (#9714)

packages/backend/server/src/models/common/index.ts

@@ -1 +1,2 @@
 export * from './feature';
+export * from './permission';

packages/backend/server/src/models/common/permission.ts

@@ -0,0 +1,6 @@
+export enum Permission {
+  Read = 0,
+  Write = 1,
+  Admin = 10,
+  Owner = 99,
+}

packages/backend/server/src/models/index.ts

@@ -12,12 +12,14 @@ import { MODELS_SYMBOL } from './provider';
 import { SessionModel } from './session';
 import { UserModel } from './user';
 import { VerificationTokenModel } from './verification-token';
+import { WorkspaceModel } from './workspace';
 
 const MODELS = {
   user: UserModel,
   session: SessionModel,
   verificationToken: VerificationTokenModel,
   feature: FeatureModel,
+  workspace: WorkspaceModel,
 };
 
 type ModelsType = {
@@ -73,3 +75,4 @@ export * from './feature';
 export * from './session';
 export * from './user';
 export * from './verification-token';
+export * from './workspace';

packages/backend/server/src/models/workspace.ts

@@ -0,0 +1,495 @@
+import { Injectable } from '@nestjs/common';
+import {
+  type Workspace,
+  WorkspaceMemberStatus,
+  type WorkspaceUserPermission,
+} from '@prisma/client';
+import { groupBy } from 'lodash-es';
+
+import { EventEmitter } from '../base';
+import { BaseModel } from './base';
+import { Permission } from './common';
+
+export { WorkspaceMemberStatus };
+export type { Workspace };
+export type UpdateWorkspaceInput = Pick<
+  Partial<Workspace>,
+  'public' | 'enableAi' | 'enableUrlPreview'
+>;
+
+export interface FindWorkspaceMembersOptions {
+  skip?: number;
+  /**
+   * Default to `8`
+   */
+  take?: number;
+}
+
+@Injectable()
+export class WorkspaceModel extends BaseModel {
+  constructor(private readonly event: EventEmitter) {
+    super();
+  }
+
+  // #region workspace
+
+  /**
+   * Create a new workspace for the user, default to private.
+   */
+  async create(userId: string) {
+    const workspace = await this.db.workspace.create({
+      data: {
+        public: false,
+        permissions: {
+          create: {
+            type: Permission.Owner,
+            userId: userId,
+            accepted: true,
+            status: WorkspaceMemberStatus.Accepted,
+          },
+        },
+      },
+    });
+    this.logger.log(`Created workspace ${workspace.id} for user ${userId}`);
+    return workspace;
+  }
+
+  /**
+   * Update the workspace with the given data.
+   */
+  async update(workspaceId: string, data: UpdateWorkspaceInput) {
+    await this.db.workspace.update({
+      where: {
+        id: workspaceId,
+      },
+      data,
+    });
+    this.logger.log(
+      `Updated workspace ${workspaceId} with data ${JSON.stringify(data)}`
+    );
+  }
+
+  async get(workspaceId: string) {
+    return await this.db.workspace.findUnique({
+      where: {
+        id: workspaceId,
+      },
+    });
+  }
+
+  async delete(workspaceId: string) {
+    await this.db.workspace.deleteMany({
+      where: {
+        id: workspaceId,
+      },
+    });
+    this.logger.log(`Deleted workspace ${workspaceId}`);
+  }
+
+  /**
+   * Find the workspace ids that the user is a member of owner.
+   */
+  async findOwnedIds(userId: string) {
+    const rows = await this.db.workspaceUserPermission.findMany({
+      where: {
+        userId,
+        type: Permission.Owner,
+        OR: this.acceptedCondition,
+      },
+      select: {
+        workspaceId: true,
+      },
+    });
+    return rows.map(row => row.workspaceId);
+  }
+
+  /**
+   * Find the accessible workspaces for the user.
+   */
+  async findAccessibleWorkspaces(userId: string) {
+    return await this.db.workspaceUserPermission.findMany({
+      where: {
+        userId,
+        OR: this.acceptedCondition,
+      },
+      include: {
+        workspace: true,
+      },
+    });
+  }
+
+  // #endregion
+
+  // #region workspace member and permission
+
+  /**
+   * Grant the workspace member with the given permission and status.
+   */
+  async grantMember(
+    workspaceId: string,
+    userId: string,
+    permission: Permission = Permission.Read,
+    status: WorkspaceMemberStatus = WorkspaceMemberStatus.Pending
+  ): Promise<WorkspaceUserPermission> {
+    const data = await this.db.workspaceUserPermission.findUnique({
+      where: {
+        workspaceId_userId: {
+          workspaceId,
+          userId,
+        },
+      },
+    });
+
+    if (!data) {
+      // Create a new permission
+      // TODO(fengmk2): should we check the permission here? Like owner can't be pending?
+      const created = await this.db.workspaceUserPermission.create({
+        data: {
+          workspaceId,
+          userId,
+          type: permission,
+          status,
+        },
+      });
+      await this.notifyMembersUpdated(workspaceId);
+      return created;
+    }
+
+    // If the user is already accepted and the new permission is owner, we need to revoke old owner
+    if (data.status === WorkspaceMemberStatus.Accepted || data.accepted) {
+      return await this.db.$transaction(async tx => {
+        const updated = await tx.workspaceUserPermission.update({
+          where: {
+            workspaceId_userId: { workspaceId, userId },
+          },
+          data: { type: permission },
+        });
+        // If the new permission is owner, we need to revoke old owner
+        if (permission === Permission.Owner) {
+          await tx.workspaceUserPermission.updateMany({
+            where: {
+              workspaceId,
+              type: Permission.Owner,
+              userId: { not: userId },
+            },
+            data: { type: Permission.Admin },
+          });
+          this.logger.log(
+            `Change owner of workspace ${workspaceId} to ${userId}`
+          );
+        }
+        return updated;
+      });
+    }
+
+    // If the user is not accepted, we can update the status directly
+    const allowedStatus = this.getAllowedStatusSource(data.status);
+    if (allowedStatus.includes(status)) {
+      const updated = await this.db.workspaceUserPermission.update({
+        where: { workspaceId_userId: { workspaceId, userId } },
+        data: {
+          status,
+          // TODO(fengmk2): should we update the permission here?
+          // type: permission,
+        },
+      });
+      return updated;
+    }
+
+    // nothing to do
+    return data;
+  }
+
+  /**
+   * Get the workspace member invitation.
+   */
+  async getMemberInvitation(invitationId: string) {
+    return await this.db.workspaceUserPermission.findUnique({
+      where: {
+        id: invitationId,
+      },
+    });
+  }
+
+  /**
+   * Accept the workspace member invitation.
+   * @param status: the status to update to, default to `Accepted`. Can be `Accepted` or `UnderReview`.
+   */
+  async acceptMemberInvitation(
+    invitationId: string,
+    workspaceId: string,
+    status: WorkspaceMemberStatus = WorkspaceMemberStatus.Accepted
+  ) {
+    const { count } = await this.db.workspaceUserPermission.updateMany({
+      where: {
+        id: invitationId,
+        workspaceId: workspaceId,
+        // TODO(fengmk2): should we check the status here?
+        AND: [{ accepted: false }, { status: WorkspaceMemberStatus.Pending }],
+      },
+      data: { accepted: true, status },
+    });
+    return count > 0;
+  }
+
+  /**
+   * Get a workspace member in accepted status by workspace id and user id.
+   */
+  async getMember(workspaceId: string, userId: string) {
+    return await this.db.workspaceUserPermission.findUnique({
+      where: {
+        workspaceId_userId: {
+          workspaceId,
+          userId,
+        },
+        OR: this.acceptedCondition,
+      },
+    });
+  }
+
+  /**
+   * Get a workspace member in any status by workspace id and user id.
+   */
+  async getMemberInAnyStatus(workspaceId: string, userId: string) {
+    return await this.db.workspaceUserPermission.findUnique({
+      where: {
+        workspaceId_userId: {
+          workspaceId,
+          userId,
+        },
+      },
+    });
+  }
+
+  /**
+   * Returns whether a given user is a member of a workspace and has the given or higher permission.
+   * Default to read permission.
+   */
+  async isMember(
+    workspaceId: string,
+    userId: string,
+    permission: Permission = Permission.Read
+  ) {
+    const count = await this.db.workspaceUserPermission.count({
+      where: {
+        workspaceId,
+        userId,
+        OR: this.acceptedCondition,
+        type: {
+          gte: permission,
+        },
+      },
+    });
+    return count > 0;
+  }
+
+  /**
+   * Get the workspace owner.
+   */
+  async getOwner(workspaceId: string) {
+    return await this.db.workspaceUserPermission.findFirst({
+      where: {
+        workspaceId,
+        type: Permission.Owner,
+        OR: this.acceptedCondition,
+      },
+      include: {
+        user: true,
+      },
+    });
+  }
+
+  /**
+   * Find the workspace admins.
+   */
+  async findAdmins(workspaceId: string) {
+    return await this.db.workspaceUserPermission.findMany({
+      where: {
+        workspaceId,
+        type: Permission.Admin,
+        OR: this.acceptedCondition,
+      },
+      include: {
+        user: true,
+      },
+    });
+  }
+
+  /**
+   * Find the workspace members.
+   */
+  async findMembers(
+    workspaceId: string,
+    options: FindWorkspaceMembersOptions = {}
+  ) {
+    return await this.db.workspaceUserPermission.findMany({
+      where: {
+        workspaceId,
+      },
+      skip: options.skip,
+      take: options.take || 8,
+      orderBy: [{ createdAt: 'asc' }, { type: 'desc' }],
+      include: {
+        user: true,
+      },
+    });
+  }
+
+  /**
+   * Delete a workspace member by workspace id and user id.
+   * Except the owner, the owner can't be deleted.
+   */
+  async deleteMember(workspaceId: string, userId: string) {
+    const member = await this.getMemberInAnyStatus(workspaceId, userId);
+
+    // We shouldn't revoke owner permission
+    // should auto deleted by workspace/user delete cascading
+    if (!member || member.type === Permission.Owner) {
+      return false;
+    }
+
+    await this.db.workspaceUserPermission.deleteMany({
+      where: {
+        workspaceId,
+        userId,
+      },
+    });
+    this.logger.log(
+      `Deleted workspace member ${userId} from workspace ${workspaceId}`
+    );
+
+    await this.notifyMembersUpdated(workspaceId);
+
+    if (
+      member.status === WorkspaceMemberStatus.UnderReview ||
+      member.status === WorkspaceMemberStatus.NeedMoreSeatAndReview
+    ) {
+      this.event.emit('workspace.members.requestDeclined', {
+        workspaceId,
+        userId,
+      });
+    }
+
+    return true;
+  }
+
+  private async notifyMembersUpdated(workspaceId: string) {
+    const count = await this.getMemberTotalCount(workspaceId);
+    this.event.emit('workspace.members.updated', {
+      workspaceId,
+      count,
+    });
+  }
+
+  /**
+   * Get the workspace member total count, including pending and accepted.
+   */
+  async getMemberTotalCount(workspaceId: string) {
+    return await this.db.workspaceUserPermission.count({
+      where: {
+        workspaceId,
+      },
+    });
+  }
+
+  /**
+   * Get the workspace member used count, only count the accepted member
+   */
+  async getMemberUsedCount(workspaceId: string) {
+    return await this.db.workspaceUserPermission.count({
+      where: {
+        workspaceId,
+        OR: this.acceptedCondition,
+      },
+    });
+  }
+
+  /**
+   * Refresh the workspace member seat status.
+   */
+  async refreshMemberSeatStatus(workspaceId: string, memberLimit: number) {
+    const usedCount = await this.getMemberUsedCount(workspaceId);
+    const availableCount = memberLimit - usedCount;
+    if (availableCount <= 0) {
+      return;
+    }
+
+    return await this.db.$transaction(async tx => {
+      const members = await tx.workspaceUserPermission.findMany({
+        select: { id: true, status: true },
+        where: {
+          workspaceId,
+          status: {
+            in: [
+              WorkspaceMemberStatus.NeedMoreSeat,
+              WorkspaceMemberStatus.NeedMoreSeatAndReview,
+            ],
+          },
+        },
+        // find the oldest members first
+        orderBy: { createdAt: 'asc' },
+      });
+
+      const needChange = members.slice(0, availableCount);
+      const groups = groupBy(needChange, m => m.status);
+
+      const toPendings = groups.NeedMoreSeat;
+      if (toPendings) {
+        // NeedMoreSeat => Pending
+        await tx.workspaceUserPermission.updateMany({
+          where: { id: { in: toPendings.map(m => m.id) } },
+          data: { status: WorkspaceMemberStatus.Pending },
+        });
+      }
+
+      const toUnderReviews = groups.NeedMoreSeatAndReview;
+      if (toUnderReviews) {
+        // NeedMoreSeatAndReview => UnderReview
+        await tx.workspaceUserPermission.updateMany({
+          where: { id: { in: toUnderReviews.map(m => m.id) } },
+          data: { status: WorkspaceMemberStatus.UnderReview },
+        });
+      }
+    });
+  }
+
+  /**
+   * Accepted condition for workspace member.
+   */
+  private get acceptedCondition() {
+    return [
+      {
+        // keep compatibility with old data
+        accepted: true,
+      },
+      {
+        status: WorkspaceMemberStatus.Accepted,
+      },
+    ];
+  }
+
+  /**
+   * NeedMoreSeat => Pending
+   *
+   * NeedMoreSeatAndReview => UnderReview
+   *
+   * Pending | UnderReview => Accepted
+   */
+  private getAllowedStatusSource(
+    to: WorkspaceMemberStatus
+  ): WorkspaceMemberStatus[] {
+    switch (to) {
+      case WorkspaceMemberStatus.NeedMoreSeat:
+        return [WorkspaceMemberStatus.Pending];
+      case WorkspaceMemberStatus.NeedMoreSeatAndReview:
+        return [WorkspaceMemberStatus.UnderReview];
+      case WorkspaceMemberStatus.Pending:
+      case WorkspaceMemberStatus.UnderReview: // need admin to review in team workspace
+        return [WorkspaceMemberStatus.Accepted];
+      default:
+        return [];
+    }
+  }
+
+  // #endregion
+}