From e93f30a6a53a3a5d1fb3088d12c2a4469c54f9a3 Mon Sep 17 00:00:00 2001
From: Sourabh Awashti <sourabh.awashti@harness.io>
Date: Tue, 18 Feb 2025 05:50:56 +0000
Subject: [PATCH] feat:[AH-945]: API to fetch storage details per account
 (#3373)

* feat:[AH-945]: API to fetch storage details per account
* feat:[AH-945]: API to fetch storage details per account
---
 .../app/api/controller/metadata/controller.go |  6 +++++
 registry/app/api/router/harness/route.go      |  2 ++
 registry/app/store/database.go                |  2 ++
 registry/app/store/database/blob.go           | 22 +++++++++++++++++++
 registry/app/store/database/generic_blob.go   | 22 +++++++++++++++++++
 5 files changed, 54 insertions(+)

diff --git a/registry/app/api/controller/metadata/controller.go b/registry/app/api/controller/metadata/controller.go
index 38888d144..ba32c3452 100644
--- a/registry/app/api/controller/metadata/controller.go
+++ b/registry/app/api/controller/metadata/controller.go
@@ -30,6 +30,8 @@ import (
 type APIController struct {
 	ImageStore         store.ImageRepository
 	fileManager        filemanager.FileManager
+	BlobStore          store.BlobRepository
+	GenericBlobStore   store.GenericBlobRepository
 	RegistryRepository store.RegistryRepository
 	UpstreamProxyStore store.UpstreamProxyConfigRepository
 	TagStore           store.TagRepository
@@ -49,6 +51,8 @@ type APIController struct {
 func NewAPIController(
 	repositoryStore store.RegistryRepository,
 	fileManager filemanager.FileManager,
+	blobStore store.BlobRepository,
+	genericBlobStore store.GenericBlobRepository,
 	upstreamProxyStore store.UpstreamProxyConfigRepository,
 	tagStore store.TagRepository,
 	manifestStore store.ManifestRepository,
@@ -66,6 +70,8 @@ func NewAPIController(
 ) *APIController {
 	return &APIController{
 		fileManager:        fileManager,
+		GenericBlobStore:   genericBlobStore,
+		BlobStore:          blobStore,
 		RegistryRepository: repositoryStore,
 		UpstreamProxyStore: upstreamProxyStore,
 		TagStore:           tagStore,
diff --git a/registry/app/api/router/harness/route.go b/registry/app/api/router/harness/route.go
index 9fe89882d..62efa461e 100644
--- a/registry/app/api/router/harness/route.go
+++ b/registry/app/api/router/harness/route.go
@@ -79,6 +79,8 @@ func NewAPIHandler(
 	apiController := metadata.NewAPIController(
 		repoDao,
 		fileManager,
+		nil,
+		nil,
 		upstreamproxyDao,
 		tagDao,
 		manifestDao,
diff --git a/registry/app/store/database.go b/registry/app/store/database.go
index ac5433f4f..41a91856b 100644
--- a/registry/app/store/database.go
+++ b/registry/app/store/database.go
@@ -46,6 +46,7 @@ type BlobRepository interface {
 		ctx context.Context, repoID int64, d digest.Digest,
 		image string,
 	) (bool, error)
+	TotalSizeByRootParentID(ctx context.Context, id int64) (int64, error)
 }
 
 type CleanupPolicyRepository interface {
@@ -556,6 +557,7 @@ type GenericBlobRepository interface {
 	) (*types.GenericBlob, error)
 	Create(ctx context.Context, gb *types.GenericBlob) error
 	DeleteByID(ctx context.Context, id string) error
+	TotalSizeByRootParentID(ctx context.Context, id int64) (int64, error)
 }
 
 type WebhooksRepository interface {
diff --git a/registry/app/store/database/blob.go b/registry/app/store/database/blob.go
index 2547457a5..56135f004 100644
--- a/registry/app/store/database/blob.go
+++ b/registry/app/store/database/blob.go
@@ -16,6 +16,7 @@ package database
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"time"
 
@@ -99,6 +100,27 @@ func (bd blobDao) FindByDigestAndRootParentID(ctx context.Context, d digest.Dige
 	return bd.mapToBlob(dst)
 }
 
+func (bd blobDao) TotalSizeByRootParentID(ctx context.Context, rootID int64) (int64, error) {
+	q := database.Builder.Select("SUM(blob_size) AS size").
+		From("blobs").
+		Where("blob_root_parent_id = ?", rootID)
+
+	db := dbtx.GetAccessor(ctx, bd.db)
+
+	var size int64
+	sqlQuery, args, err := q.ToSql()
+	if err != nil {
+		return 0, errors2.Wrap(err, "Failed to convert query to sql")
+	}
+
+	if err = db.QueryRowContext(ctx, sqlQuery, args...).Scan(&size); err != nil &&
+		!errors2.Is(err, sql.ErrNoRows) {
+		return 0,
+			database.ProcessSQLErrorf(ctx, err, "Failed to find total blob size for root parent with id %d", rootID)
+	}
+	return size, nil
+}
+
 func (bd blobDao) FindByID(ctx context.Context, id int64) (*types.Blob, error) {
 	stmt := PrimaryQuery.
 		Where("blob_id = ?", id)
diff --git a/registry/app/store/database/generic_blob.go b/registry/app/store/database/generic_blob.go
index af0d78b8f..706314856 100644
--- a/registry/app/store/database/generic_blob.go
+++ b/registry/app/store/database/generic_blob.go
@@ -57,6 +57,28 @@ func (g GenericBlobDao) FindByID(ctx context.Context, id string) (*types.Generic
 	return g.mapToGenericBlob(ctx, dst)
 }
 
+func (g GenericBlobDao) TotalSizeByRootParentID(ctx context.Context, rootID int64) (int64, error) {
+	q := databaseg.Builder.
+		Select("SUM(generic_blob_size) AS size").
+		From("generic_blobs").
+		Where("generic_blob_root_parent_id = ?", rootID)
+
+	db := dbtx.GetAccessor(ctx, g.sqlDB)
+
+	var size int64
+	sqlQuery, args, err := q.ToSql()
+	if err != nil {
+		return 0, errors.Wrap(err, "Failed to convert query to sql")
+	}
+
+	if err = db.QueryRowContext(ctx, sqlQuery, args...).Scan(&size); err != nil &&
+		!errors.Is(err, sql.ErrNoRows) {
+		return 0,
+			databaseg.ProcessSQLErrorf(ctx, err, "Failed to find total blob size for root parent with id %d", rootID)
+	}
+	return size, nil
+}
+
 func (g GenericBlobDao) FindBySha256AndRootParentID(ctx context.Context,
 	sha256 string, rootParentID int64) (
 	*types.GenericBlob, error) {