From 1bf193d54351e55bb48b001ce7454acef8b4ec46 Mon Sep 17 00:00:00 2001
From: Abhishek Shroff <shroff@kudosoft.net>
Date: Tue, 8 Oct 2024 20:02:28 +0530
Subject: [PATCH] [server] Add disk usage endpoint

---
 server/internal/api/routes/resources.go | 33 +++++++++++++++++++++++
 server/internal/core/filesystem.go      | 21 +++++++++++++++
 server/internal/db/resources.sql.go     | 36 +++++++++++++++++++++++++
 server/sql/queries/resources.sql        | 15 +++++++++++
 4 files changed, 105 insertions(+)

diff --git a/server/internal/api/routes/resources.go b/server/internal/api/routes/resources.go
index c6971a8e..a62ce224 100644
--- a/server/internal/api/routes/resources.go
+++ b/server/internal/api/routes/resources.go
@@ -38,6 +38,12 @@ type resourceResponse struct {
 	Permissions string     `json:"permissions,omitempty"`
 }
 
+type diskUsageResponse struct {
+	TotalSize int64 `json:"size"`
+	Files     int64 `json:"files"`
+	Dirs      int64 `json:"dirs"`
+}
+
 type resourceDetailResponse struct {
 	Metadata             resourceResponse   `json:"metadata"`
 	Children             []resourceResponse `json:"children,omitempty"`
@@ -57,6 +63,7 @@ type resourceMoveParams struct {
 func SetupResourceRoutes(r *gin.RouterGroup) {
 	group := r.Group("/resources")
 	group.GET("/ls/:id", handleResourceLsRoute)
+	group.GET("/du/:id", handleResourceDiskUsageRoute)
 	group.GET("/metadata/:id", handleResourceMetadataRoute)
 	group.GET("/details/:id", handleResourceDetailsRoute)
 	group.GET("/contents/:id", handleResourceContentsRoute)
@@ -113,6 +120,32 @@ func handleResourceLsRoute(c *gin.Context) {
 	c.JSON(200, results)
 }
 
+func handleResourceDiskUsageRoute(c *gin.Context) {
+	resourceId, err := uuid.Parse(c.Param("id"))
+	if err != nil {
+		panic(errResourceIDInvalid)
+	}
+
+	fs := auth.GetFileSystem(c)
+	resource, err := fs.ResourceByID(resourceId)
+	if err != nil {
+		if errors.Is(err, iofs.ErrNotExist) {
+			err = errResourceNotFound
+		}
+		panic(err)
+	}
+
+	info, err := fs.DiskUsage(resource)
+	if err != nil {
+		panic(err)
+	}
+	c.JSON(200, diskUsageResponse{
+		TotalSize: info.TotalSize,
+		Files:     info.Files,
+		Dirs:      info.Dirs,
+	})
+}
+
 func handleResourceDetailsRoute(c *gin.Context) {
 	resourceID, err := uuid.Parse(c.Param("id"))
 	if err != nil {
diff --git a/server/internal/core/filesystem.go b/server/internal/core/filesystem.go
index d8255068..58cc995c 100644
--- a/server/internal/core/filesystem.go
+++ b/server/internal/core/filesystem.go
@@ -28,6 +28,13 @@ var (
 	ErrMoveTargetSubdirectory     = errors.New("move target subdirectory")
 )
 
+type DiskUsageInfo struct {
+	TotalSize int64
+	Entities  int64
+	Files     int64
+	Dirs      int64
+}
+
 type FileSystem interface {
 	WithDb(*db.DbHandler) FileSystem
 	RunInTx(fn func(FileSystem) error) error
@@ -41,6 +48,7 @@ type FileSystem interface {
 	DeleteRecursive(r Resource, hardDelete bool) (uuid.UUIDs, error)
 	Move(r Resource, name string, parent *uuid.UUID) (Resource, error)
 	UpdatePermissions(r Resource, username string, permission Permission) error
+	DiskUsage(r Resource) (DiskUsageInfo, error)
 }
 
 type filesystem struct {
@@ -340,6 +348,19 @@ func (f filesystem) Move(r Resource, name string, parent *uuid.UUID) (Resource,
 	}
 }
 
+func (f filesystem) DiskUsage(r Resource) (DiskUsageInfo, error) {
+	if info, err := f.db.DiskUsage(f.ctx, r.ID()); err != nil {
+		return DiskUsageInfo{}, err
+	} else {
+		return DiskUsageInfo{
+			TotalSize: info.Size,
+			Entities:  info.Entities,
+			Files:     info.Files,
+			Dirs:      info.Dirs,
+		}, nil
+	}
+}
+
 func (f filesystem) UpdatePermissions(r Resource, username string, permission Permission) error {
 	if r.UserPermission() < PermissionReadWriteShare {
 		return ErrInsufficientPermissions
diff --git a/server/internal/db/resources.sql.go b/server/internal/db/resources.sql.go
index 3de83b5e..768d81b4 100644
--- a/server/internal/db/resources.sql.go
+++ b/server/internal/db/resources.sql.go
@@ -85,6 +85,42 @@ func (q *Queries) DeleteRecursive(ctx context.Context, id uuid.UUID) ([]uuid.UUI
 	return items, nil
 }
 
+const diskUsage = `-- name: DiskUsage :one
+WITH RECURSIVE nodes(id, parent, size, dir) AS (
+    SELECT r.id, r.parent, r.size, r.dir
+    FROM resources r WHERE r.id = $1::uuid
+  UNION ALL
+    SELECT r.id, r.parent, r.size, r.dir
+    FROM resources r JOIN nodes n on r.parent = n.id
+    WHERE deleted IS NULL
+)
+SELECT
+  SUM(size) AS size,
+  COUNT(*) AS entities,
+  COUNT(CASE dir WHEN true THEN 1 ELSE NULL END) AS dirs,
+  COUNT(CASE dir WHEN false THEN 1 ELSE NULL END) AS files
+FROM nodes
+`
+
+type DiskUsageRow struct {
+	Size     int64
+	Entities int64
+	Dirs     int64
+	Files    int64
+}
+
+func (q *Queries) DiskUsage(ctx context.Context, id uuid.UUID) (DiskUsageRow, error) {
+	row := q.db.QueryRow(ctx, diskUsage, id)
+	var i DiskUsageRow
+	err := row.Scan(
+		&i.Size,
+		&i.Entities,
+		&i.Dirs,
+		&i.Files,
+	)
+	return i, err
+}
+
 const hardDeleteRecursive = `-- name: HardDeleteRecursive :many
 WITH RECURSIVE nodes(id, parent) AS (
     SELECT r.id, r.parent
diff --git a/server/sql/queries/resources.sql b/server/sql/queries/resources.sql
index 4c53a1c3..cc72a215 100644
--- a/server/sql/queries/resources.sql
+++ b/server/sql/queries/resources.sql
@@ -44,6 +44,21 @@ SET
 WHERE id = @id::uuid
 RETURNING *;
 
+-- name: DiskUsage :one
+WITH RECURSIVE nodes(id, parent, size, dir) AS (
+    SELECT r.id, r.parent, r.size, r.dir
+    FROM resources r WHERE r.id = @id::uuid
+  UNION ALL
+    SELECT r.id, r.parent, r.size, r.dir
+    FROM resources r JOIN nodes n on r.parent = n.id
+    WHERE deleted IS NULL
+)
+SELECT
+  SUM(size) AS size,
+  COUNT(*) AS entities,
+  COUNT(CASE dir WHEN true THEN 1 ELSE NULL END) AS dirs,
+  COUNT(CASE dir WHEN false THEN 1 ELSE NULL END) AS files
+FROM nodes;
 
 
 -- name: DeleteRecursive :many