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
|
import { eq, sum } from "drizzle-orm";
import type { DB, KarakeepDBTransaction } from "@karakeep/db";
import { assets, 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";
}
}
export async function 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 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);
}
export async function 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);
}
|
