Skip to content

Commit

Permalink
feat: historicized membership data model (#4542)
Browse files Browse the repository at this point in the history
* make startAt non-nullable

* make startAt non-nullable

* membership resource

* fix bug

* more checks

* misc BL fixes

* migration to delete the revoked role

* remove comment

* review: use lightWorkspaceType only as interface

* support tx everywhere

* use resource for poke as well

* enh: isRevoked

* active membership of user in ws

* add check to other function

* remove serialization methods

* cleanup

* pr syntax

---------

Co-authored-by: Henry Fontanier <[email protected]>
  • Loading branch information
2 people authored and flvndvd committed May 26, 2024
1 parent 89588b3 commit 9aac5da
Show file tree
Hide file tree
Showing 37 changed files with 802 additions and 523 deletions.
209 changes: 8 additions & 201 deletions front/admin/cli.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,13 @@ import { Storage } from "@google-cloud/storage";
import parseArgs from "minimist";
import readline from "readline";

import { subscriptionForWorkspace } from "@app/lib/auth";
import {
DataSource,
EventSchema,
Membership,
User,
Workspace,
} from "@app/lib/models";
import { DataSource, EventSchema, User, Workspace } from "@app/lib/models";
import { FREE_UPGRADED_PLAN_CODE } from "@app/lib/plans/plan_codes";
import {
internalSubscribeWorkspaceToFreeNoPlan,
internalSubscribeWorkspaceToFreePlan,
} from "@app/lib/plans/subscription";
import { MembershipResource } from "@app/lib/resources/membership_resource";
import { generateModelSId } from "@app/lib/utils";
import logger from "@app/logger/logger";

Expand All @@ -25,105 +19,6 @@ const { DUST_DATA_SOURCES_BUCKET = "", SERVICE_ACCOUNT } = process.env;
// `cli` takes an object type and a command as first two arguments and then a list of arguments.
const workspace = async (command: string, args: parseArgs.ParsedArgs) => {
switch (command) {
case "find": {
if (!args.name) {
throw new Error("Missing --name argument");
}

const workspaces = await Workspace.findAll({
where: {
name: args.name,
},
});

workspaces.forEach((w) => {
console.log(`> wId='${w.sId}' name='${w.name}'`);
});
return;
}

case "show": {
if (!args.wId) {
throw new Error("Missing --wId argument");
}

const w = await Workspace.findOne({
where: {
sId: args.wId,
},
});

if (!w) {
throw new Error(`Workspace not found: wId='${args.wId}'`);
}

console.log(`workspace:`);
console.log(` wId: ${w.sId}`);
console.log(` name: ${w.name}`);

const subscription = await subscriptionForWorkspace(w.sId);
const plan = subscription.plan;
console.log(` plan:`);
console.log(` limits:`);
console.log(` dataSources:`);
console.log(` count: ${plan.limits.dataSources.count}`);
console.log(` documents:`);
console.log(
` count: ${plan.limits.dataSources.documents.count}`
);
console.log(
` sizeMb: ${plan.limits.dataSources.documents.sizeMb}`
);
console.log(
` managed Slack: ${plan.limits.connections.isSlackAllowed}`
);
console.log(
` managed Notion: ${plan.limits.connections.isNotionAllowed}`
);
console.log(
` managed Github: ${plan.limits.connections.isGithubAllowed}`
);
console.log(
` managed Intercom: ${plan.limits.connections.isIntercomAllowed}`
);
console.log(
` managed Google Drive: ${plan.limits.connections.isGoogleDriveAllowed}`
);

const dataSources = await DataSource.findAll({
where: {
workspaceId: w.id,
},
});

console.log("Data sources:");
dataSources.forEach((ds) => {
console.log(` - name: ${ds.name} provider: ${ds.connectorProvider}`);
});

const memberships = await Membership.findAll({
where: {
workspaceId: w.id,
},
});
const users = await User.findAll({
where: {
id: memberships.map((m) => m.userId),
},
});

console.log("Users:");
users.forEach((u) => {
console.log(
` - userId: ${u.id} email: ${u.email} role: ${
memberships.find((m) => m.userId === u.id)?.role
}`
);
});

return;
}

case "create": {
if (!args.name) {
throw new Error("Missing --name argument");
Expand Down Expand Up @@ -182,98 +77,10 @@ const workspace = async (command: string, args: parseArgs.ParsedArgs) => {
return;
}

case "add-user": {
if (!args.wId) {
throw new Error("Missing --wId argument");
}
if (!args.userId) {
throw new Error("Missing --userId argument");
}
if (!args.role) {
throw new Error("Missing --role argument");
}
if (!["admin", "builder", "user"].includes(args.role)) {
throw new Error(`Invalid --role: ${args.role}`);
}
const role = args.role as "admin" | "builder" | "user";

const w = await Workspace.findOne({
where: {
sId: args.wId,
},
});
if (!w) {
throw new Error(`Workspace not found: wId='${args.wId}'`);
}
const u = await User.findOne({
where: {
id: args.userId,
},
});
if (!u) {
throw new Error(`User not found: userId='${args.userId}'`);
}
await Membership.create({
role,
workspaceId: w.id,
userId: u.id,
startAt: new Date(),
});
return;
}

case "change-role": {
if (!args.wId) {
throw new Error("Missing --wId argument");
}
if (!args.userId) {
throw new Error("Missing --userId argument");
}
if (!args.role) {
throw new Error("Missing --role argument");
}
if (!["admin", "builder", "user", "revoked"].includes(args.role)) {
throw new Error(`Invalid --role: ${args.role}`);
}
const role = args.role as "admin" | "builder" | "user" | "revoked";

const w = await Workspace.findOne({
where: {
sId: args.wId,
},
});
if (!w) {
throw new Error(`Workspace not found: wId='${args.wId}'`);
}
const u = await User.findOne({
where: {
id: args.userId,
},
});
if (!u) {
throw new Error(`User not found: userId='${args.userId}'`);
}
const m = await Membership.findOne({
where: {
workspaceId: w.id,
userId: u.id,
},
});
if (!m) {
throw new Error(
`User is not a member of workspace: userId='${args.userId}' wId='${args.wId}'`
);
}

m.role = role;
await m.save();
return;
}

default:
console.log(`Unknown workspace command: ${command}`);
console.log(
"Possible values: `find`, `show`, `create`, `set-limits`, `add-user`, `change-role`, `upgrade`, `downgrade`"
"Possible values: `find`, `show`, `create`, `set-limits`, `upgrade`, `downgrade`"
);
}
};
Expand Down Expand Up @@ -319,10 +126,8 @@ const user = async (command: string, args: parseArgs.ParsedArgs) => {
console.log(` name: ${u.name}`);
console.log(` email: ${u.email}`);

const memberships = await Membership.findAll({
where: {
userId: u.id,
},
const memberships = await MembershipResource.getLatestMemberships({
userIds: [u.id],
});

const workspaces = await Workspace.findAll({
Expand All @@ -338,7 +143,9 @@ const user = async (command: string, args: parseArgs.ParsedArgs) => {
console.log(` - wId: ${w.sId}`);
console.log(` name: ${w.name}`);
if (m) {
console.log(` role: ${m.role}`);
console.log(` role: ${m.isRevoked() ? "revoked" : m.role}`);
console.log(` startAt: ${m.startAt}`);
console.log(` endAt: ${m.endAt}`);
}
});

Expand Down
4 changes: 2 additions & 2 deletions front/admin/db.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ import {
EventSchema,
ExtractedEvent,
Key,
Membership,
MembershipInvitation,
Mention,
Message,
Expand Down Expand Up @@ -47,14 +46,15 @@ import {
import { ConversationClassification } from "@app/lib/models/conversation_classification";
import { FeatureFlag } from "@app/lib/models/feature_flag";
import { ContentFragmentModel } from "@app/lib/resources/storage/models/content_fragment";
import { MembershipModel } from "@app/lib/resources/storage/models/membership";
import { TemplateModel } from "@app/lib/resources/storage/models/templates";

async function main() {
await User.sync({ alter: true });
await UserMetadata.sync({ alter: true });
await Workspace.sync({ alter: true });
await WorkspaceHasDomain.sync({ alter: true });
await Membership.sync({ alter: true });
await MembershipModel.sync({ alter: true });
await MembershipInvitation.sync({ alter: true });
await App.sync({ alter: true });
await Dataset.sync({ alter: true });
Expand Down
8 changes: 3 additions & 5 deletions front/lib/amplitude/node/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@ import {
import { isGlobalAgentId } from "@app/lib/api/assistant/global_agents";
import type { Authenticator } from "@app/lib/auth";
import { subscriptionForWorkspace } from "@app/lib/auth";
import { Membership } from "@app/lib/models";
import { User, Workspace } from "@app/lib/models";
import { countActiveSeatsInWorkspace } from "@app/lib/plans/workspace_usage";
import { MembershipResource } from "@app/lib/resources/membership_resource";

let BACKEND_CLIENT: Ampli | null = null;

Expand Down Expand Up @@ -55,10 +55,8 @@ export function getBackendClient() {
export async function trackUserMemberships(userId: ModelId) {
const amplitude = getBackendClient();
const user = await User.findByPk(userId);
const memberships = await Membership.findAll({
where: {
userId: userId,
},
const memberships = await MembershipResource.getActiveMemberships({
userIds: [userId],
});
const groups: string[] = [];
for (const membership of memberships) {
Expand Down
Loading

0 comments on commit 9aac5da

Please sign in to comment.