anything-llm/server/models/workspace.js

239 lines
7.5 KiB
JavaScript
Raw Normal View History

2023-06-08 06:31:35 +02:00
const slugify = require("slugify");
const { Document } = require("./documents");
const { checkForMigrations } = require("../utils/database");
const { WorkspaceUser } = require("./workspaceUsers");
2023-06-04 04:28:07 +02:00
const Workspace = {
2023-06-08 06:31:35 +02:00
tablename: "workspaces",
writable: [
// Used for generic updates so we can validate keys in request body
"name",
"slug",
"vectorTag",
"openAiTemp",
"openAiHistory",
"lastUpdatedAt",
"openAiPrompt",
],
2023-06-04 04:28:07 +02:00
colsInit: `
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
2023-06-04 04:28:07 +02:00
slug TEXT NOT NULL UNIQUE,
vectorTag TEXT DEFAULT NULL,
createdAt TEXT DEFAULT CURRENT_TIMESTAMP,
openAiTemp REAL DEFAULT NULL,
openAiHistory INTEGER DEFAULT 20,
lastUpdatedAt TEXT DEFAULT CURRENT_TIMESTAMP,
openAiPrompt TEXT DEFAULT NULL
2023-06-04 04:28:07 +02:00
`,
migrateTable: async function () {
console.log(`\x1b[34m[MIGRATING]\x1b[0m Checking for Workspace migrations`);
const db = await this.db(false);
await checkForMigrations(this, db);
},
migrations: function () {
return [
{
colName: "openAiTemp",
execCmd: `ALTER TABLE ${this.tablename} ADD COLUMN openAiTemp REAL DEFAULT NULL`,
doif: false,
},
{
colName: "openAiPrompt",
execCmd: `ALTER TABLE ${this.tablename} ADD COLUMN openAiPrompt TEXT DEFAULT NULL`,
doif: false,
},
{
colName: "id",
execCmd: `CREATE TRIGGER IF NOT EXISTS Trg_LastUpdated AFTER UPDATE ON ${this.tablename}
FOR EACH ROW
BEGIN
UPDATE ${this.tablename} SET lastUpdatedAt = CURRENT_TIMESTAMP WHERE id = old.id;
END`,
doif: true,
},
{
colName: "openAiHistory",
execCmd: `ALTER TABLE ${this.tablename} ADD COLUMN openAiHistory INTEGER DEFAULT 20`,
doif: false,
},
];
},
db: async function (tracing = true) {
2023-06-08 06:31:35 +02:00
const sqlite3 = require("sqlite3").verbose();
const { open } = require("sqlite");
2023-06-04 04:28:07 +02:00
const db = await open({
filename: `${
!!process.env.STORAGE_DIR ? `${process.env.STORAGE_DIR}/` : "storage/"
}anythingllm.db`,
2023-06-08 06:31:35 +02:00
driver: sqlite3.Database,
});
2023-06-04 04:28:07 +02:00
2023-06-08 06:31:35 +02:00
await db.exec(
`PRAGMA foreign_keys = ON;CREATE TABLE IF NOT EXISTS ${this.tablename} (${this.colsInit})`
2023-06-08 06:31:35 +02:00
);
if (tracing) db.on("trace", (sql) => console.log(sql));
2023-06-08 06:31:35 +02:00
return db;
2023-06-04 04:28:07 +02:00
},
new: async function (name = null, creatorId = null) {
2023-06-08 06:31:35 +02:00
if (!name) return { result: null, message: "name cannot be null" };
var slug = slugify(name, { lower: true });
const existingBySlug = await this.get(`slug = '${slug}'`);
if (existingBySlug !== null) {
const slugSeed = Math.floor(10000000 + Math.random() * 90000000);
slug = slugify(`${name}-${slugSeed}`, { lower: true });
}
2023-06-04 04:28:07 +02:00
2023-06-08 06:31:35 +02:00
const db = await this.db();
const { id, success, message } = await db
.run(`INSERT INTO ${this.tablename} (name, slug) VALUES (?, ?)`, [
name,
slug,
2023-06-08 06:31:35 +02:00
])
2023-06-04 04:28:07 +02:00
.then((res) => {
2023-06-08 06:31:35 +02:00
return { id: res.lastID, success: true, message: null };
2023-06-04 04:28:07 +02:00
})
.catch((error) => {
2023-06-08 06:31:35 +02:00
return { id: null, success: false, message: error.message };
});
if (!success) {
db.close();
return { workspace: null, message };
}
2023-06-04 04:28:07 +02:00
2023-06-08 06:31:35 +02:00
const workspace = await db.get(
`SELECT * FROM ${this.tablename} WHERE id = ${id}`
);
db.close();
// If created with a user then we need to create the relationship as well.
// If creating with an admin User it wont change anything because admins can
// view all workspaces anyway.
if (!!creatorId) await WorkspaceUser.create(creatorId, workspace.id);
2023-06-08 06:31:35 +02:00
return { workspace, message: null };
2023-06-04 04:28:07 +02:00
},
update: async function (id = null, data = {}) {
if (!id) throw new Error("No workspace id provided for update");
const validKeys = Object.keys(data).filter((key) =>
this.writable.includes(key)
);
const values = Object.values(data);
if (validKeys.length === 0 || validKeys.length !== values.length)
return { workspace: { id }, message: "No valid fields to update!" };
2023-06-28 21:54:18 +02:00
const template = `UPDATE ${this.tablename} SET ${validKeys.map((key) => {
return `${key}=?`;
})} WHERE id = ?`;
const db = await this.db();
const { success, message } = await db
.run(template, [...values, id])
.then(() => {
return { success: true, message: null };
})
.catch((error) => {
return { success: false, message: error.message };
});
2023-06-28 21:54:18 +02:00
db.close();
if (!success) {
return { workspace: null, message };
}
const updatedWorkspace = await this.get(`id = ${id}`);
return { workspace: updatedWorkspace, message: null };
},
getWithUser: async function (user = null, clause = "") {
if (user.role === "admin") return this.get(clause);
const db = await this.db();
const result = await db
.get(
`SELECT * FROM ${this.tablename} as workspace
LEFT JOIN workspace_users as ws_users
ON ws_users.workspace_id = workspace.id
WHERE ws_users.user_id = ${user?.id} AND ${clause}`
)
.then((res) => res || null);
if (!result) return null;
db.close();
const workspace = { ...result, id: result.workspace_id };
const documents = await Document.forWorkspace(workspace.id);
return { ...workspace, documents };
},
2023-06-08 06:31:35 +02:00
get: async function (clause = "") {
const db = await this.db();
const result = await db
.get(`SELECT * FROM ${this.tablename} WHERE ${clause}`)
.then((res) => res || null);
2023-06-04 04:28:07 +02:00
if (!result) return null;
db.close();
2023-06-04 04:28:07 +02:00
const documents = await Document.forWorkspace(result.id);
2023-06-08 06:31:35 +02:00
return { ...result, documents };
2023-06-04 04:28:07 +02:00
},
2023-06-08 06:31:35 +02:00
delete: async function (clause = "") {
const db = await this.db();
await db.get(`DELETE FROM ${this.tablename} WHERE ${clause}`);
db.close();
2023-06-08 06:31:35 +02:00
return true;
2023-06-04 04:28:07 +02:00
},
where: async function (clause = "", limit = null, orderBy = null) {
2023-06-08 06:31:35 +02:00
const db = await this.db();
const results = await db.all(
`SELECT * FROM ${this.tablename} ${clause ? `WHERE ${clause}` : ""} ${
!!limit ? `LIMIT ${limit}` : ""
} ${!!orderBy ? orderBy : ""}`
2023-06-08 06:31:35 +02:00
);
db.close();
2023-06-08 06:31:35 +02:00
return results;
2023-06-04 04:28:07 +02:00
},
whereWithUser: async function (
user,
clause = null,
limit = null,
orderBy = null
) {
if (user.role === "admin") return await this.where(clause, limit);
const db = await this.db();
const results = await db.all(
`SELECT * FROM ${this.tablename} as workspace
LEFT JOIN workspace_users as ws_users
ON ws_users.workspace_id = workspace.id
WHERE ws_users.user_id = ${user.id} ${clause ? `AND ${clause}` : ""} ${
!!limit ? `LIMIT ${limit}` : ""
} ${!!orderBy ? orderBy : ""}`
);
db.close();
const workspaces = results.map((ws) => {
return { ...ws, id: ws.workspace_id };
});
return workspaces;
},
whereWithUsers: async function (clause = "", limit = null, orderBy = null) {
const workspaces = await this.where(clause, limit, orderBy);
for (const workspace of workspaces) {
const userIds = (
await WorkspaceUser.where(`workspace_id = ${workspace.id}`)
).map((rel) => rel.user_id);
workspace.userIds = userIds;
}
return workspaces;
},
updateUsers: async function (workspaceId, userIds = []) {
await WorkspaceUser.delete(`workspace_id = ${workspaceId}`);
await WorkspaceUser.createManyUsers(userIds, workspaceId);
return { success: true, error: null };
},
2023-06-08 06:31:35 +02:00
};
2023-06-04 04:28:07 +02:00
2023-06-08 06:31:35 +02:00
module.exports = { Workspace };