Skip to content

Commit

Permalink
Introduce Plan & Subscription models (#2204)
Browse files Browse the repository at this point in the history
  • Loading branch information
PopDaph authored Oct 20, 2023
1 parent 750a113 commit 5fdb0bb
Show file tree
Hide file tree
Showing 7 changed files with 305 additions and 0 deletions.
5 changes: 5 additions & 0 deletions front/admin/db.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,12 @@ import {
Mention,
Message,
MessageReaction,
Plan,
Provider,
RetrievalDocument,
RetrievalDocumentChunk,
Run,
Subscription,
TrackedDocument,
User,
UserMessage,
Expand Down Expand Up @@ -55,6 +57,9 @@ async function main() {
await ExtractedEvent.sync({ alter: true });
await DocumentTrackerChangeSuggestion.sync({ alter: true });

await Plan.sync({ alter: true });
await Subscription.sync({ alter: true });

await AgentDustAppRunConfiguration.sync({ alter: true });
await AgentDustAppRunAction.sync({ alter: true });

Expand Down
3 changes: 3 additions & 0 deletions front/admin/init_plans.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
#!/bin/sh
env $(cat .env.local) npx tsx admin/init_plans.ts

16 changes: 16 additions & 0 deletions front/admin/init_plans.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
import { upsertFreePlans } from "@app/lib/plans/free_plans";

async function main() {
await upsertFreePlans();
process.exit(0);
}

main()
.then(() => {
console.log("Done");
process.exit(0);
})
.catch((err) => {
console.error(err);
process.exit(1);
});
3 changes: 3 additions & 0 deletions front/lib/models/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import {
TrackedDocument,
} from "@app/lib/models/doc_tracker";
import { EventSchema, ExtractedEvent } from "@app/lib/models/extract";
import { Plan, Subscription } from "@app/lib/models/plan";
import { User, UserMetadata } from "@app/lib/models/user";
import {
Key,
Expand Down Expand Up @@ -64,10 +65,12 @@ export {
Mention,
Message,
MessageReaction,
Plan,
Provider,
RetrievalDocument,
RetrievalDocumentChunk,
Run,
Subscription,
TrackedDocument,
User,
UserMessage,
Expand Down
213 changes: 213 additions & 0 deletions front/lib/models/plan.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,213 @@
import {
CreationOptional,
DataTypes,
ForeignKey,
InferAttributes,
InferCreationAttributes,
Model,
NonAttribute,
Transaction,
} from "sequelize";

import { front_sequelize } from "@app/lib/databases";
import { Workspace } from "@app/lib/models/workspace";

export class Plan extends Model<
InferAttributes<Plan>,
InferCreationAttributes<Plan>
> {
declare id: CreationOptional<number>;
declare createdAt: CreationOptional<Date>;
declare updatedAt: CreationOptional<Date>;

declare code: string; // unique
declare name: string;

// workspace limitations
declare maxWeeklyMessages: number;
declare maxUsersInWorkspace: number;
declare isSlackbotAllowed: boolean;
declare isManagedSlackAllowed: boolean;
declare isManagedNotionAllowed: boolean;
declare isManagedGoogleDriveAllowed: boolean;
declare isManagedGithubAllowed: boolean;
declare maxNbStaticDataSources: number;
declare maxNbStaticDocuments: number;
declare maxSizeStaticDataSources: number;
}
Plan.init(
{
id: {
type: DataTypes.INTEGER,
autoIncrement: true,
primaryKey: true,
},
createdAt: {
type: DataTypes.DATE,
allowNull: false,
defaultValue: DataTypes.NOW,
},
updatedAt: {
type: DataTypes.DATE,
allowNull: false,
defaultValue: DataTypes.NOW,
},
code: {
type: DataTypes.STRING,
allowNull: false,
},
name: {
type: DataTypes.STRING,
allowNull: false,
},
maxWeeklyMessages: {
type: DataTypes.INTEGER,
allowNull: false,
},
maxUsersInWorkspace: {
type: DataTypes.INTEGER,
allowNull: false,
},
isSlackbotAllowed: {
type: DataTypes.BOOLEAN,
defaultValue: false,
},
isManagedSlackAllowed: {
type: DataTypes.BOOLEAN,
defaultValue: false,
},
isManagedNotionAllowed: {
type: DataTypes.BOOLEAN,
defaultValue: false,
},
isManagedGoogleDriveAllowed: {
type: DataTypes.BOOLEAN,
defaultValue: false,
},
isManagedGithubAllowed: {
type: DataTypes.BOOLEAN,
defaultValue: false,
},
maxNbStaticDataSources: {
type: DataTypes.INTEGER,
allowNull: false,
},
maxNbStaticDocuments: {
type: DataTypes.INTEGER,
allowNull: false,
},
maxSizeStaticDataSources: {
type: DataTypes.INTEGER,
allowNull: false,
},
},
{
modelName: "plan",
sequelize: front_sequelize,
indexes: [{ unique: true, fields: ["code"] }],
}
);

export class Subscription extends Model<
InferAttributes<Subscription>,
InferCreationAttributes<Subscription>
> {
declare id: CreationOptional<number>;
declare createdAt: CreationOptional<Date>;
declare updatedAt: CreationOptional<Date>;

declare sId: string; // unique
declare status: string; // "active" | "ended"
declare startDate: Date;
declare endDate: Date | null;

declare workspaceId: ForeignKey<Workspace["id"]>;
declare workspace: NonAttribute<Workspace>;

declare planId: ForeignKey<Plan["id"]>;
declare plan: NonAttribute<Plan>;
}
Subscription.init(
{
id: {
type: DataTypes.INTEGER,
autoIncrement: true,
primaryKey: true,
},
createdAt: {
type: DataTypes.DATE,
allowNull: false,
defaultValue: DataTypes.NOW,
},
updatedAt: {
type: DataTypes.DATE,
allowNull: false,
defaultValue: DataTypes.NOW,
},
sId: {
type: DataTypes.STRING,
allowNull: false,
},
status: {
type: DataTypes.STRING,
allowNull: false,
validate: {
isIn: [["active", "ended"]],
},
},
startDate: {
type: DataTypes.DATE,
allowNull: false,
},
endDate: {
type: DataTypes.DATE,
allowNull: true,
},
},
{
modelName: "subscription",
sequelize: front_sequelize,
indexes: [{ unique: true, fields: ["sId"] }],
}
);
// Define a hook to ensure there's only one active subscription for each workspace
Subscription.addHook(
"beforeCreate",
"enforce_single_active_subscription",
async (subscription: Subscription, options: { transaction: Transaction }) => {
if (subscription.status === "active") {
// Check if there's already an active subscription for the same workspace
const existingActiveSubscription = await Subscription.findOne({
where: {
workspaceId: subscription.workspaceId,
status: "active",
},
transaction: options.transaction, // Include the transaction in your query
});

if (existingActiveSubscription) {
throw new Error(
"An active subscription already exists for this workspace."
);
}
}
}
);

// Plan <> Subscription relationship: attribute "planId" in Subscription
Plan.hasMany(Subscription, {
foreignKey: { name: "workspaceId", allowNull: false },
onDelete: "CASCADE",
});
Subscription.belongsTo(Plan, {
foreignKey: { name: "planId", allowNull: false },
});

// Subscription <> Workspace relationship: attribute "workspaceId" in Subscription
Workspace.hasMany(Subscription, {
foreignKey: { name: "workspaceId", allowNull: false },
onDelete: "CASCADE",
});
Subscription.belongsTo(Workspace, {
foreignKey: { name: "workspaceId", allowNull: false },
});
4 changes: 4 additions & 0 deletions front/lib/models/workspace.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,11 @@ import {
InferAttributes,
InferCreationAttributes,
Model,
NonAttribute,
} from "sequelize";

import { front_sequelize } from "@app/lib/databases";
import { Subscription } from "@app/lib/models/plan";
import { User } from "@app/lib/models/user";

export class Workspace extends Model<
Expand All @@ -24,6 +26,8 @@ export class Workspace extends Model<
declare description: string | null;
declare allowedDomain: string | null;
declare plan: string | null;

declare subscriptions: NonAttribute<Subscription[]>;
}
Workspace.init(
{
Expand Down
61 changes: 61 additions & 0 deletions front/lib/plans/free_plans.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import { Attributes } from "sequelize";

import { Plan } from "@app/lib/models";

type PlanAttributes = Omit<Attributes<Plan>, "id" | "createdAt" | "updatedAt">;

/**
* This file is used to store the plans data.
* We never delete plans, we only add new ones.
*/
const FREE_PLANS_DATA: PlanAttributes[] = [
{
code: "TEST_PLAN_V0",
name: "Test",
maxWeeklyMessages: 50,
maxUsersInWorkspace: 1,
isSlackbotAllowed: false,
isManagedSlackAllowed: false,
isManagedNotionAllowed: false,
isManagedGoogleDriveAllowed: false,
isManagedGithubAllowed: false,
maxNbStaticDataSources: 10,
maxNbStaticDocuments: 10,
maxSizeStaticDataSources: 1000,
},
{
code: "TRIAL_PLAN_V0",
name: "Free Trial",
maxWeeklyMessages: -1,
maxUsersInWorkspace: -1,
isSlackbotAllowed: true,
isManagedSlackAllowed: true,
isManagedNotionAllowed: true,
isManagedGoogleDriveAllowed: true,
isManagedGithubAllowed: true,
maxNbStaticDataSources: -1,
maxNbStaticDocuments: -1,
maxSizeStaticDataSources: -1,
},
];

export const upsertFreePlans = async () => {
for (const planData of FREE_PLANS_DATA) {
await _upsertFreePlan(planData);
}
};

const _upsertFreePlan = async (planData: PlanAttributes) => {
const plan = await Plan.findOne({
where: {
code: planData.code,
},
});
if (plan === null) {
await Plan.create(planData);
console.log(`Free plan ${planData.code} created.`);
} else {
await plan.update(planData);
console.log(`Free plan ${planData.code} updated.`);
}
};

0 comments on commit 5fdb0bb

Please sign in to comment.