From aada3aab0289228af20a36ab962e33a90bd00f64 Mon Sep 17 00:00:00 2001
From: Aaron Liu <abslant.liu@gmail.com>
Date: Fri, 4 Jul 2025 10:05:15 +0800
Subject: [PATCH] feat(storage): load balance storage policy (#2436)

---
 assets                  |  2 +-
 inventory/policy.go     |  9 +++++++++
 service/admin/policy.go | 13 ++++++++++++-
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/assets b/assets
index 8e2c2bc..27996dc 160000
--- a/assets
+++ b/assets
@@ -1 +1 @@
-Subproject commit 8e2c2bcff17d4728a01c2cabab8c3b639d72f428
+Subproject commit 27996dc3ea22ab3b7ae525841ad1d45098d375f4
diff --git a/inventory/policy.go b/inventory/policy.go
index 2343cf6..285c01b 100644
--- a/inventory/policy.go
+++ b/inventory/policy.go
@@ -27,6 +27,7 @@ type (
 	SkipStoragePolicyCache struct{}
 
 	StoragePolicyClient interface {
+		TxOperator
 		// GetByGroup returns the storage policies of the group.
 		GetByGroup(ctx context.Context, group *ent.Group) (*ent.StoragePolicy, error)
 		// GetPolicyByID returns the storage policy by id.
@@ -64,6 +65,14 @@ type storagePolicyClient struct {
 	cache  cache.Driver
 }
 
+func (c *storagePolicyClient) SetClient(newClient *ent.Client) TxOperator {
+	return &storagePolicyClient{client: newClient, cache: c.cache}
+}
+
+func (c *storagePolicyClient) GetClient() *ent.Client {
+	return c.client
+}
+
 func (c *storagePolicyClient) Delete(ctx context.Context, policy *ent.StoragePolicy) error {
 	if err := c.client.StoragePolicy.DeleteOne(policy).Exec(ctx); err != nil {
 		return fmt.Errorf("failed to delete storage policy: %w", err)
diff --git a/service/admin/policy.go b/service/admin/policy.go
index e73e5cd..aa6625c 100644
--- a/service/admin/policy.go
+++ b/service/admin/policy.go
@@ -294,11 +294,22 @@ func (service *UpdateStoragePolicyService) Update(c *gin.Context) (*GetStoragePo
 	}
 
 	service.Policy.ID = idInt
-	_, err = storagePolicyClient.Upsert(c, service.Policy)
+
+	sc, tx, ctx, err := inventory.WithTx(c, storagePolicyClient)
 	if err != nil {
+		return nil, serializer.NewError(serializer.CodeDBError, "Failed to create transaction", err)
+	}
+
+	_, err = sc.Upsert(ctx, service.Policy)
+	if err != nil {
+		_ = inventory.Rollback(tx)
 		return nil, serializer.NewError(serializer.CodeDBError, "Failed to update policy", err)
 	}
 
+	if err := inventory.Commit(tx); err != nil {
+		return nil, serializer.NewError(serializer.CodeDBError, "Failed to commit transaction", err)
+	}
+
 	_ = dep.KV().Delete(manager.EntityUrlCacheKeyPrefix)
 
 	s := SingleStoragePolicyService{ID: idInt}