Refactor runtime framework router (#249)

* use admin role to operate runtime framework * refactor runtime framework router --------- Co-authored-by: James <[email protected]>
OpenCSGs · Jan 23, 2025 · 358b126 · 358b126
1 parent 9a83786
commit 358b126
Show file tree

Hide file tree

Showing 16 changed files with 31,603 additions and 15,384 deletions.
diff --git a/_mocks/opencsg.com/csghub-server/component/mock_ModelComponent.go b/_mocks/opencsg.com/csghub-server/component/mock_ModelComponent.go
diff --git a/_mocks/opencsg.com/csghub-server/component/mock_RepoComponent.go b/_mocks/opencsg.com/csghub-server/component/mock_RepoComponent.go
diff --git a/api/handler/model.go b/api/handler/model.go
@@ -1196,8 +1196,9 @@ func (h *ModelHandler) ListAllRuntimeFramework(ctx *gin.Context) {
 // @Param        body body types.RuntimeFrameworkModels true "body"
 // @Success      200  {object}  types.Response{} "OK"
 // @Failure      400  {object}  types.APIBadRequest "Bad request"
+// @Failure      403  {object}  types.APIForbidden "Forbidden"
 // @Failure      500  {object}  types.APIInternalServerError "Internal server error"
-// @Router       /runtime_framework/{id} [post]
+// @Router       /runtime_framework/{id}/models [put]
 func (h *ModelHandler) UpdateModelRuntimeFrameworks(ctx *gin.Context) {
 	var req types.RuntimeFrameworkModels
 	if err := ctx.ShouldBindJSON(&req); err != nil {
@@ -1206,6 +1207,12 @@ func (h *ModelHandler) UpdateModelRuntimeFrameworks(ctx *gin.Context) {
 		return
 	}
 
+	currentUser := httpbase.GetCurrentUser(ctx)
+	if currentUser == "" {
+		httpbase.UnauthorizedError(ctx, component.ErrUserNotFound)
+		return
+	}
+
 	id, err := strconv.ParseInt(ctx.Param("id"), 10, 64)
 	if err != nil {
 		slog.Error("Bad request format", "error", err)
@@ -1227,8 +1234,12 @@ func (h *ModelHandler) UpdateModelRuntimeFrameworks(ctx *gin.Context) {
 
 	slog.Info("update runtime frameworks models", slog.Any("req", req), slog.Any("runtime framework id", id), slog.Any("deployType", deployType))
 
-	list, err := h.model.SetRuntimeFrameworkModes(ctx.Request.Context(), deployType, id, req.Models)
+	list, err := h.model.SetRuntimeFrameworkModes(ctx.Request.Context(), currentUser, deployType, id, req.Models)
 	if err != nil {
+		if errors.Is(err, component.ErrForbidden) {
+			httpbase.ForbiddenError(ctx, err)
+			return
+		}
 		slog.Error("Failed to set models runtime framework", slog.Any("error", err))
 		httpbase.ServerError(ctx, err)
 		return
@@ -1249,8 +1260,9 @@ func (h *ModelHandler) UpdateModelRuntimeFrameworks(ctx *gin.Context) {
 // @Param        body body types.RuntimeFrameworkModels true "body"
 // @Success      200  {object}  types.Response{} "OK"
 // @Failure      400  {object}  types.APIBadRequest "Bad request"
+// @Failure      403  {object}  types.APIForbidden "Forbidden"
 // @Failure      500  {object}  types.APIInternalServerError "Internal server error"
-// @Router       /runtime_framework/{id} [delete]
+// @Router       /runtime_framework/{id}/models [delete]
 func (h *ModelHandler) DeleteModelRuntimeFrameworks(ctx *gin.Context) {
 	var req types.RuntimeFrameworkModels
 	if err := ctx.ShouldBindJSON(&req); err != nil {
@@ -1259,6 +1271,12 @@ func (h *ModelHandler) DeleteModelRuntimeFrameworks(ctx *gin.Context) {
 		return
 	}
 
+	currentUser := httpbase.GetCurrentUser(ctx)
+	if currentUser == "" {
+		httpbase.UnauthorizedError(ctx, component.ErrUserNotFound)
+		return
+	}
+
 	id, err := strconv.ParseInt(ctx.Param("id"), 10, 64)
 	if err != nil {
 		slog.Error("Bad request format", "error", err)
@@ -1280,8 +1298,12 @@ func (h *ModelHandler) DeleteModelRuntimeFrameworks(ctx *gin.Context) {
 
 	slog.Info("update runtime frameworks models", slog.Any("req", req), slog.Any("runtime framework id", id), slog.Any("deployType", deployType))
 
-	list, err := h.model.DeleteRuntimeFrameworkModes(ctx.Request.Context(), deployType, id, req.Models)
+	list, err := h.model.DeleteRuntimeFrameworkModes(ctx.Request.Context(), currentUser, deployType, id, req.Models)
 	if err != nil {
+		if errors.Is(err, component.ErrForbidden) {
+			httpbase.ForbiddenError(ctx, err)
+			return
+		}
 		slog.Error("Failed to set models runtime framework", slog.Any("error", err))
 		httpbase.ServerError(ctx, err)
 		return

diff --git a/api/handler/model_test.go b/api/handler/model_test.go
@@ -432,7 +432,7 @@ func TestModelHandler_UpdateModelRuntimeFramework(t *testing.T) {
 
 	tester.WithUser().WithQuery("deploy_type", "").AddPagination(1, 10).WithParam("id", "1")
 	tester.mocks.model.EXPECT().SetRuntimeFrameworkModes(
-		tester.ctx, types.InferenceType, int64(1), []string{"foo"},
+		tester.ctx, "u", types.InferenceType, int64(1), []string{"foo"},
 	).Return([]string{"bar"}, nil)
 	tester.WithBody(t, types.RuntimeFrameworkModels{
 		Models: []string{"foo"},
@@ -447,7 +447,7 @@ func TestModelHandler_DeleteModelRuntimeFramework(t *testing.T) {
 
 	tester.WithUser().WithQuery("deploy_type", "").AddPagination(1, 10).WithParam("id", "1")
 	tester.mocks.model.EXPECT().DeleteRuntimeFrameworkModes(
-		tester.ctx, types.InferenceType, int64(1), []string{"foo"},
+		tester.ctx, "u", types.InferenceType, int64(1), []string{"foo"},
 	).Return([]string{"bar"}, nil)
 	tester.WithBody(t, types.RuntimeFrameworkModels{
 		Models: []string{"foo"},