Skip to content

Commit c59161b

Browse files
QinYuuuuDev Agent
QinYuuuu
and
Dev Agent
authored
add filter to index space resource (#634)
Co-authored-by: Dev Agent <[email protected]>
1 parent 06b766d commit c59161b

File tree

8 files changed

+140
-46
lines changed

8 files changed

+140
-46
lines changed

‎_mocks/opencsg.com/csghub-server/builder/store/database/mock_SpaceResourceStore.go‎

Lines changed: 19 additions & 17 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎api/handler/space_resource.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,17 @@ func (h *SpaceResourceHandler) Index(ctx *gin.Context) {
7070
Page: page,
7171
},
7272
}
73+
if ctx.Query("resource_type") != "" {
74+
req.ResourceType = types.ResourceType(ctx.Query("resource_type"))
75+
if !types.ResourceTypeValid(req.ResourceType) {
76+
slog.ErrorContext(ctx.Request.Context(), "Invalid resource type", "resource_type", req.ResourceType)
77+
httpbase.BadRequest(ctx, "Invalid resource type")
78+
return
79+
}
80+
}
81+
if ctx.Query("hardware_type") != "" {
82+
req.HardwareType = ctx.Query("hardware_type")
83+
}
7384
spaceResources, total, err := h.spaceResource.Index(ctx.Request.Context(), req)
7485
if err != nil {
7586
slog.ErrorContext(ctx.Request.Context(), "Failed to get space resources", slog.String("cluster_id", clusterId), slog.String("deploy_type", deployTypeStr), slog.Any("error", err))

‎api/handler/space_resource_test.go‎

Lines changed: 33 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
"github.com/gin-gonic/gin"
88
mockcomponent "opencsg.com/csghub-server/_mocks/opencsg.com/csghub-server/component"
9+
"opencsg.com/csghub-server/api/httpbase"
910
"opencsg.com/csghub-server/builder/testutil"
1011
"opencsg.com/csghub-server/common/types"
1112
)
@@ -36,26 +37,41 @@ func (t *SpaceResourceTester) WithHandleFunc(fn func(h *SpaceResourceHandler) gi
3637
}
3738

3839
func TestSpaceResourceHandler_Index(t *testing.T) {
39-
tester := NewSpaceResourceTester(t).WithHandleFunc(func(h *SpaceResourceHandler) gin.HandlerFunc {
40-
return h.Index
41-
})
40+
t.Run("200", func(t *testing.T) {
41+
tester := NewSpaceResourceTester(t).WithHandleFunc(func(h *SpaceResourceHandler) gin.HandlerFunc {
42+
return h.Index
43+
})
4244

43-
req := &types.SpaceResourceIndexReq{
44-
ClusterID: "c1",
45-
DeployType: types.InferenceType,
46-
CurrentUser: "u",
47-
PageOpts: types.PageOpts{
48-
PageSize: 50,
49-
Page: 1,
50-
},
51-
}
45+
req := &types.SpaceResourceIndexReq{
46+
ClusterID: "c1",
47+
DeployType: types.InferenceType,
48+
CurrentUser: "u",
49+
ResourceType: types.ResourceTypeCPU,
50+
HardwareType: "intel",
51+
PageOpts: types.PageOpts{
52+
PageSize: 50,
53+
Page: 1,
54+
},
55+
}
56+
57+
tester.mocks.spaceResource.EXPECT().Index(tester.Ctx(), req).Return(
58+
[]types.SpaceResource{{Name: "sp"}}, 0, nil,
59+
)
60+
tester.WithQuery("cluster_id", "c1").WithQuery("deploy_type", "").
61+
WithQuery("resource_type", "cpu").WithQuery("hardware_type", "intel").WithUser().Execute()
5262

53-
tester.mocks.spaceResource.EXPECT().Index(tester.Ctx(), req).Return(
54-
[]types.SpaceResource{{Name: "sp"}}, 0, nil,
55-
)
56-
tester.WithQuery("cluster_id", "c1").WithQuery("deploy_type", "").WithUser().Execute()
63+
tester.ResponseEq(t, 200, tester.OKText, []types.SpaceResource{{Name: "sp"}})
64+
})
65+
t.Run("invalid resource type", func(t *testing.T) {
66+
tester := NewSpaceResourceTester(t).WithHandleFunc(func(h *SpaceResourceHandler) gin.HandlerFunc {
67+
return h.Index
68+
})
69+
70+
tester.WithQuery("cluster_id", "c1").WithQuery("deploy_type", "").
71+
WithQuery("resource_type", "invalid").WithQuery("hardware_type", "intel").WithUser().Execute()
5772

58-
tester.ResponseEq(t, 200, tester.OKText, []types.SpaceResource{{Name: "sp"}})
73+
tester.ResponseEqSimple(t, 400, httpbase.R{Msg: "Invalid resource type"})
74+
})
5975
}
6076

6177
func TestSpaceResourceHandler_Create(t *testing.T) {

‎builder/store/database/space_resource.go‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ type spaceResourceStoreImpl struct {
1414
}
1515

1616
type SpaceResourceStore interface {
17-
Index(ctx context.Context, clusterId string, per, page int) ([]SpaceResource, int, error)
17+
Index(ctx context.Context, filter types.SpaceResourceFilter, per, page int) ([]SpaceResource, int, error)
1818
Create(ctx context.Context, input SpaceResource) (*SpaceResource, error)
1919
Update(ctx context.Context, input SpaceResource) (*SpaceResource, error)
2020
Delete(ctx context.Context, input SpaceResource) error
@@ -40,9 +40,18 @@ type SpaceResource struct {
4040
times
4141
}
4242

43-
func (s *spaceResourceStoreImpl) Index(ctx context.Context, clusterId string, per, page int) ([]SpaceResource, int, error) {
43+
func (s *spaceResourceStoreImpl) Index(ctx context.Context, filter types.SpaceResourceFilter, per, page int) ([]SpaceResource, int, error) {
4444
var result []SpaceResource
45-
query := s.db.Operator.Core.NewSelect().Model(&result).Where("cluster_id = ?", clusterId)
45+
query := s.db.Operator.Core.NewSelect().Model(&result)
46+
if filter.ClusterID != "" {
47+
query = query.Where("cluster_id = ?", filter.ClusterID)
48+
}
49+
if filter.HardwareType != "" {
50+
query = query.Where("EXISTS (SELECT 1 FROM jsonb_each(resources::jsonb) WHERE value->>'type' = ?)", filter.HardwareType)
51+
}
52+
if filter.ResourceType != "" {
53+
query = query.Where("EXISTS (SELECT 1 FROM jsonb_each(resources::jsonb) WHERE key = ?)", filter.ResourceType)
54+
}
4655
query = query.Order("name asc").
4756
Limit(per).
4857
Offset((page - 1) * per)
@@ -106,14 +115,14 @@ func (s *spaceResourceStoreImpl) FindAll(ctx context.Context) ([]SpaceResource,
106115
func (s *spaceResourceStoreImpl) FindAllResourceTypes(ctx context.Context, clusterId string) ([]string, error) {
107116
typeSet := make(map[string]bool)
108117
var hardWareTypes []string
109-
118+
filter := types.SpaceResourceFilter{ClusterID: clusterId}
110119
// Use pagination to query resources
111120
page := 1
112121
per := 100 // Set a reasonable page size
113122

114123
for {
115124
// Get resources for current page
116-
resources, _, err := s.Index(ctx, clusterId, per, page)
125+
resources, _, err := s.Index(ctx, filter, per, page)
117126
if err != nil {
118127
return nil, err
119128
}

‎builder/store/database/space_resource_test.go‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"github.com/stretchr/testify/require"
88
"opencsg.com/csghub-server/builder/store/database"
99
"opencsg.com/csghub-server/common/tests"
10+
"opencsg.com/csghub-server/common/types"
1011
)
1112

1213
func TestSpaceResourceStore_CRUD(t *testing.T) {
@@ -39,7 +40,7 @@ func TestSpaceResourceStore_CRUD(t *testing.T) {
3940
require.Equal(t, 1, len(srs))
4041
require.Equal(t, "c1", srs[0].ClusterID)
4142

42-
srs, _, err = store.Index(ctx, "c1", 50, 1)
43+
srs, _, err = store.Index(ctx, types.SpaceResourceFilter{ClusterID: "c1"}, 50, 1)
4344
require.Nil(t, err)
4445
require.Equal(t, 1, len(srs))
4546
require.Equal(t, "c1", srs[0].ClusterID)
@@ -147,3 +148,30 @@ func TestSpaceResourceStore_FindAllResourceTypes_InvalidJSON(t *testing.T) {
147148
require.Nil(t, err)
148149
require.Empty(t, types)
149150
}
151+
152+
func TestSpaceResourceStore_Filter(t *testing.T) {
153+
db := tests.InitTestDB()
154+
defer db.Close()
155+
ctx := context.TODO()
156+
157+
store := database.NewSpaceResourceStoreWithDB(db)
158+
159+
_, err := store.Create(ctx, database.SpaceResource{
160+
Name: "r1",
161+
ClusterID: "c1",
162+
Resources: `{"cpu": { "type": "Intel","num": "200m"}}`,
163+
})
164+
require.Nil(t, err)
165+
sr := &database.SpaceResource{}
166+
err = db.Core.NewSelect().Model(sr).Where("name=?", "r1").Scan(ctx, sr)
167+
require.Nil(t, err)
168+
require.Equal(t, "c1", sr.ClusterID)
169+
170+
srs, total, err := store.Index(ctx, types.SpaceResourceFilter{
171+
ResourceType: types.ResourceTypeCPU,
172+
}, 50, 1)
173+
require.Nil(t, err)
174+
require.Equal(t, 1, len(srs))
175+
require.Equal(t, 1, total)
176+
require.Equal(t, "c1", srs[0].ClusterID)
177+
}

‎common/types/space_resource.go‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,16 @@ const (
1717
PayModeYear PayMode = "year"
1818
)
1919

20+
func ResourceTypeValid(resourceType ResourceType) bool {
21+
return resourceType == ResourceTypeCPU ||
22+
resourceType == ResourceTypeGPU ||
23+
resourceType == ResourceTypeNPU ||
24+
resourceType == ResourceTypeGCU ||
25+
resourceType == ResourceTypeGPGPU ||
26+
resourceType == ResourceTypeMLU ||
27+
resourceType == ResourceTypeDCU
28+
}
29+
2030
type SpaceResource struct {
2131
ID int64 `json:"id"`
2232
Name string `json:"name"`
@@ -43,8 +53,17 @@ type UpdateSpaceResourceReq struct {
4353
}
4454

4555
type SpaceResourceIndexReq struct {
46-
ClusterID string `json:"cluster_id"`
47-
DeployType int `json:"deploy_type"`
48-
CurrentUser string `json:"current_user"`
56+
ClusterID string `json:"cluster_id"`
57+
DeployType int `json:"deploy_type"`
58+
CurrentUser string `json:"current_user"`
59+
ResourceType ResourceType `json:"resource_type"`
60+
HardwareType string `json:"hardware_type"`
61+
IsAvailable bool `json:"is_available"`
4962
PageOpts
5063
}
64+
65+
type SpaceResourceFilter struct {
66+
ClusterID string `json:"cluster_id"`
67+
ResourceType ResourceType `json:"resource_type"`
68+
HardwareType string `json:"hardware_type"`
69+
}

‎component/space_resource.go‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,16 @@ func (c *spaceResourceComponentImpl) Index(ctx context.Context, req *types.Space
5050
var total int
5151
for _, clusterID := range clusterIDs {
5252
var singleClusterResult []types.SpaceResource
53-
databaseSpaceResources, currentTotal, err := c.spaceResourceStore.Index(ctx, clusterID, req.PageSize, req.Page)
53+
dbReq := types.SpaceResourceFilter{
54+
ClusterID: clusterID,
55+
}
56+
if req.HardwareType != "" {
57+
dbReq.HardwareType = req.HardwareType
58+
}
59+
if req.ResourceType != "" {
60+
dbReq.ResourceType = req.ResourceType
61+
}
62+
databaseSpaceResources, currentTotal, err := c.spaceResourceStore.Index(ctx, dbReq, req.PageSize, req.Page)
5463
if err != nil {
5564
slog.Error("failed to index space resource", slog.String("clusterID", clusterID), slog.Any("error", err))
5665
continue

‎component/space_resource_ce_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ func TestSpaceResourceComponent_Index(t *testing.T) {
1616
ctx := context.TODO()
1717
sc := initializeTestSpaceResourceComponent(ctx, t)
1818

19-
sc.mocks.stores.SpaceResourceMock().EXPECT().Index(ctx, "c1", 50, 1).Return(
19+
sc.mocks.stores.SpaceResourceMock().EXPECT().Index(ctx, types.SpaceResourceFilter{ClusterID: "c1"}, 50, 1).Return(
2020
[]database.SpaceResource{
2121
{ID: 1, Name: "sr", Resources: `{"memory": "1000", "gpu": {"num": "5"}}`},
2222
{ID: 2, Name: "sr2", Resources: `{"memory": "1000"}`},
@@ -54,7 +54,7 @@ func TestSpaceResourceComponent_Index_With_Status_Filter(t *testing.T) {
5454
sc.mocks.deployer.EXPECT().ListCluster(ctx).Return(clusters, nil)
5555
sc.mocks.deployer.EXPECT().CheckHeartbeatTimeout(ctx, "cluster1").Once().Return(true, nil)
5656
sc.mocks.deployer.EXPECT().CheckHeartbeatTimeout(ctx, "cluster2").Once().Return(false, nil)
57-
sc.mocks.stores.SpaceResourceMock().EXPECT().Index(ctx, "cluster2", 50, 1).Return([]database.SpaceResource{}, 0, nil)
57+
sc.mocks.stores.SpaceResourceMock().EXPECT().Index(ctx, types.SpaceResourceFilter{ClusterID: "cluster2"}, 50, 1).Return([]database.SpaceResource{}, 0, nil)
5858
sc.mocks.deployer.EXPECT().GetClusterById(ctx, "cluster2").Return(&types.ClusterRes{}, nil)
5959
req := &types.SpaceResourceIndexReq{
6060
CurrentUser: "user1",

0 commit comments

Comments
 (0)