1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
|
import { count, eq, sum } from "drizzle-orm";
import type { DB, KarakeepDBTransaction } from "@karakeep/db";
import { assets, bookmarks, users } from "@karakeep/db/schema";
import { QuotaApproved } from "@karakeep/shared/storageQuota";
export class StorageQuotaError extends Error {
constructor(
public readonly currentUsage: number,
public readonly quota: number,
public readonly requestedSize: number,
) {
super(
`Storage quota exceeded. Current usage: ${Math.round(currentUsage / 1024 / 1024)}MB, Quota: ${Math.round(quota / 1024 / 1024)}MB, Requested: ${Math.round(requestedSize / 1024 / 1024)}MB`,
);
this.name = "StorageQuotaError";
}
}
// TODO: Change the API of this class to either return a boolean
// or throw an exception on lack of quota because now, it's inconsistent.
export class QuotaService {
// TODO: Use quota approval tokens for bookmark creation when
// bookmark creation logic is in the model.
static async canCreateBookmark(db: DB, userId: string) {
const user = await db.query.users.findFirst({
where: eq(users.id, userId),
columns: {
bookmarkQuota: true,
},
});
if (user?.bookmarkQuota !== null && user?.bookmarkQuota !== undefined) {
const currentBookmarkCount = await db
.select({ count: count() })
.from(bookmarks)
.where(eq(bookmarks.userId, userId));
if (currentBookmarkCount[0].count >= user.bookmarkQuota) {
return {
result: false,
error: `Bookmark quota exceeded. You can only have ${user.bookmarkQuota} bookmarks.`,
} as const;
}
}
return {
result: true,
} as const;
}
static async checkStorageQuota(
db: DB | KarakeepDBTransaction,
userId: string,
requestedSize: number,
): Promise<QuotaApproved> {
const user = await db.query.users.findFirst({
where: eq(users.id, userId),
columns: {
storageQuota: true,
},
});
if (user?.storageQuota === null || user?.storageQuota === undefined) {
// No quota limit - approve the request
return QuotaApproved._create(userId, requestedSize);
}
const currentUsage = await this.getCurrentStorageUsage(db, userId);
if (currentUsage + requestedSize > user.storageQuota) {
throw new StorageQuotaError(
currentUsage,
user.storageQuota,
requestedSize,
);
}
// Quota check passed - return approval token
return QuotaApproved._create(userId, requestedSize);
}
static async getCurrentStorageUsage(
db: DB | KarakeepDBTransaction,
userId: string,
): Promise<number> {
const currentUsageResult = await db
.select({ totalSize: sum(assets.size) })
.from(assets)
.where(eq(assets.userId, userId));
return Number(currentUsageResult[0]?.totalSize ?? 0);
}
}
|
