feat: add basic storage usage tracking

apps/drive-web/src/files/use-upload-file.ts

@@ -9,7 +9,7 @@ function useUploadFile({
 	targetDirectory: Doc<"directories">
 }) {
 	const generateUploadUrl = useConvexMutation(api.files.generateUploadUrl)
-	const saveFile = useConvexMutation(api.files.saveFile)
+	const saveFile = useConvexMutation(api.filesystem.saveFile)
 
 	async function upload({
 		file,
@@ -44,8 +44,6 @@ function useUploadFile({
 			saveFile({
 				storageId,
 				name: file.name,
-				size: file.size,
-				mimeType: file.type,
 				directoryId: targetDirectory._id,
 			}),
 		)

packages/convex/auth.ts

@@ -19,12 +19,19 @@ export const authComponent = createClient<DataModel, typeof authSchema>(
 			user: {
 				onCreate: async (ctx, user) => {
 					const now = Date.now()
-					await ctx.db.insert("directories", {
+					await Promise.all([
+						ctx.db.insert("userInfo", {
+							userId: user._id,
+							storageUsageBytes: 0,
+							storageQuotaBytes: 1024 * 1024 * 1024 * 5, // 5GB
+						}),
+						ctx.db.insert("directories", {
 							name: "",
 							userId: user._id,
 							createdAt: now,
 							updatedAt: now,
-					})
+						}),
+					])
 				},
 			},
 		},

packages/convex/files.ts

@@ -7,9 +7,21 @@ import {
 } from "./functions"
 import * as Directories from "./model/directories"
 import * as Files from "./model/files"
+import * as User from "./model/user"
+import * as Err from "./shared/error"
 
 export const generateUploadUrl = authenticatedMutation({
 	handler: async (ctx) => {
+		const usageStatistics = await User.queryCachedUsageStatistics(ctx)
+		if (!usageStatistics) {
+			throw Err.create(Err.Code.Internal, "Internal server error")
+		}
+		if (
+			usageStatistics.storageUsageBytes >=
+			usageStatistics.storageQuotaBytes
+		) {
+			throw Err.create(Err.Code.Forbidden, "Storage quota exceeded")
+		}
 		return await ctx.storage.generateUploadUrl()
 	},
 })
@@ -68,12 +80,10 @@ export const createDirectory = authenticatedMutation({
 export const saveFile = authenticatedMutation({
 	args: {
 		name: v.string(),
-		size: v.number(),
 		directoryId: v.id("directories"),
 		storageId: v.id("_storage"),
-		mimeType: v.optional(v.string()),
 	},
-	handler: async (ctx, { name, storageId, directoryId, size, mimeType }) => {
+	handler: async (ctx, { name, storageId, directoryId }) => {
 		const directory = await authorizedGet(ctx, directoryId)
 		if (!directory) {
 			throw new Error("Directory not found")
@@ -81,16 +91,26 @@ export const saveFile = authenticatedMutation({
 
 		const now = Date.now()
 
-		await ctx.db.insert("files", {
+		const fileMetadata = await Promise.all([
+			ctx.db.system.get(storageId),
+			ctx.user.queryUsageStatistics(),
+		])
+		if (!fileMetadata) {
+			throw Err.create(Err.Code.Internal, "Internal server error")
+		}
+
+		await Promise.all([
+			ctx.db.insert("files", {
 				name,
-			size,
+				size: fileMetadata.size,
 				storageId,
 				directoryId,
 				userId: ctx.user._id,
 				mimeType,
 				createdAt: now,
 				updatedAt: now,
-		})
+			}),
+		])
 	},
 })
 

packages/convex/filesystem.ts

@@ -189,6 +189,17 @@ export const openFile = authenticatedMutation({
 	},
 })
 
+export const saveFile = authenticatedMutation({
+	args: {
+		name: v.string(),
+		directoryId: v.id("directories"),
+		storageId: v.id("_storage"),
+	},
+	handler: async (ctx, { name, directoryId, storageId }) => {
+		return await FileSystem.saveFile(ctx, { name, directoryId, storageId })
+	},
+})
+
 export const fetchRecentFiles = authenticatedQuery({
 	args: {
 		limit: v.number(),

packages/convex/model/filesystem.ts

@@ -20,6 +20,7 @@ import * as Directories from "./directories"
 import * as FilePreview from "./filepreview"
 import * as Files from "./files"
 import * as FileShare from "./fileshare"
+import * as User from "./user"
 
 export const VDirectoryHandle = v.object({
 	kind: v.literal(FileType.Directory),
@@ -266,6 +267,60 @@ export async function openFile(
 	}
 }
 
+export async function saveFile(
+	ctx: AuthenticatedMutationCtx,
+	{
+		name,
+		storageId,
+		directoryId,
+	}: {
+		name: string
+		storageId: Id<"_storage">
+		directoryId: Id<"directories">
+	},
+) {
+	const directory = await authorizedGet(ctx, directoryId)
+	if (!directory) {
+		throw Err.create(Err.Code.NotFound, "directory not found")
+	}
+
+	const [fileMetadata, userInfo] = await Promise.all([
+		ctx.db.system.get(storageId),
+		User.queryInfo(ctx),
+	])
+	if (!fileMetadata || !userInfo) {
+		throw Err.create(Err.Code.Internal, "Internal server error")
+	}
+
+	if (
+		userInfo.storageUsageBytes + fileMetadata.size >
+		userInfo.storageQuotaBytes
+	) {
+		await ctx.storage.delete(storageId)
+		throw Err.create(Err.Code.StorageQuotaExceeded, "Storage quota exceeded")
+	}
+
+	const now = Date.now()
+
+	const [fileId] = await Promise.all([
+		ctx.db.insert("files", {
+			name,
+			userId: ctx.user._id,
+			createdAt: now,
+			updatedAt: now,
+			storageId,
+			directoryId,
+			size: fileMetadata.size,
+			mimeType: fileMetadata.contentType,
+		}),
+		ctx.db.patch(userInfo._id, {
+			storageUsageBytes: userInfo.storageUsageBytes + fileMetadata.size,
+		}),
+	])
+
+	return fileId
+}
+
 export async function fetchRecentFiles(
 	ctx: AuthenticatedQueryCtx,
 	{ limit }: { limit: number },

packages/convex/model/user.ts

@@ -1,5 +1,7 @@
 import type { MutationCtx, QueryCtx } from "@fileone/convex/server"
+import type { Doc } from "../_generated/dataModel"
 import { authComponent } from "../auth"
+import { type AuthenticatedQueryCtx, authorizedGet } from "../functions"
 import * as Err from "../shared/error"
 
 export type AuthUser = Awaited<ReturnType<typeof authComponent.getAuthUser>>
@@ -23,3 +25,10 @@ export async function userOrThrow(ctx: QueryCtx | MutationCtx) {
 	const user = await authComponent.getAuthUser(ctx)
 	return user
 }
+
+export async function queryInfo(ctx: AuthenticatedQueryCtx) {
+	return await ctx.db
+		.query("userInfo")
+		.withIndex("byUserId", (q) => q.eq("userId", ctx.user._id))
+		.first()
+}

packages/convex/schema.ts

@@ -2,6 +2,12 @@ import { defineSchema, defineTable } from "convex/server"
 import { v } from "convex/values"
 
 const schema = defineSchema({
+	userInfo: defineTable({
+		userId: v.string(),
+		storageUsageBytes: v.number(),
+		storageQuotaBytes: v.number(),
+	}).index("byUserId", ["userId"]),
+
 	files: defineTable({
 		storageId: v.id("_storage"),
 		userId: v.string(), // BetterAuth user IDs are strings, not Convex Ids

packages/convex/shared/error.ts

@@ -9,6 +9,7 @@ export enum Code {
 	Internal = "Internal",
 	Unauthenticated = "Unauthenticated",
 	NotFound = "NotFound",
+	StorageQuotaExceeded = "StorageQuotaExceeded",
 }
 
 export type ApplicationErrorData = { code: Code; message?: string }