ce,ee allow create public dataset

OpenCSGs · Jan 8, 2025 · 87c11be · 87c11be
1 parent 6419f2d
commit 87c11be
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 26 deletions.
diff --git a/api/handler/dataset.go b/api/handler/dataset.go
@@ -76,6 +76,10 @@ func (h *DatasetHandler) Create(ctx *gin.Context) {
 		httpbase.BadRequest(ctx, fmt.Errorf("sensitive check failed: %w", err).Error())
 		return
 	}
+	if !req.Private && !h.allowCreatePublic() {
+		httpbase.BadRequest(ctx, "creating public dataset is not allowed")
+		return
+	}
 	req.Username = currentUser
 
 	dataset, err := h.dataset.Create(ctx, req)

diff --git a/api/handler/dataset_ce.go b/api/handler/dataset_ce.go
@@ -0,0 +1,7 @@
+//go:build !saas
+
+package handler
+
+func (h *DatasetHandler) allowCreatePublic() bool {
+	return true
+}
diff --git a/api/handler/dataset_ce_test.go b/api/handler/dataset_ce_test.go
@@ -0,0 +1,58 @@
+//go:build !saas
+
+package handler
+
+import (
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"opencsg.com/csghub-server/common/types"
+)
+
+func TestDatasetHandler_Create(t *testing.T) {
+	t.Run("no public", func(t *testing.T) {
+
+		tester := NewDatasetTester(t).WithHandleFunc(func(h *DatasetHandler) gin.HandlerFunc {
+			return h.Create
+		})
+		tester.RequireUser(t)
+
+		tester.mocks.sensitive.EXPECT().CheckRequestV2(tester.ctx, &types.CreateDatasetReq{
+			CreateRepoReq: types.CreateRepoReq{Private: true},
+		}).Return(true, nil)
+		tester.mocks.dataset.EXPECT().Create(tester.ctx, &types.CreateDatasetReq{
+			CreateRepoReq: types.CreateRepoReq{Private: true, Username: "u"},
+		}).Return(&types.Dataset{Name: "d"}, nil)
+		tester.WithBody(t, &types.CreateDatasetReq{
+			CreateRepoReq: types.CreateRepoReq{Private: true},
+		}).Execute()
+
+		tester.ResponseEqSimple(t, 200, gin.H{
+			"data": &types.Dataset{Name: "d"},
+		})
+
+	})
+
+	t.Run("public", func(t *testing.T) {
+
+		tester := NewDatasetTester(t).WithHandleFunc(func(h *DatasetHandler) gin.HandlerFunc {
+			return h.Create
+		})
+		tester.RequireUser(t)
+
+		tester.mocks.sensitive.EXPECT().CheckRequestV2(tester.ctx, &types.CreateDatasetReq{
+			CreateRepoReq: types.CreateRepoReq{Private: false},
+		}).Return(true, nil)
+		tester.mocks.dataset.EXPECT().Create(tester.ctx, &types.CreateDatasetReq{
+			CreateRepoReq: types.CreateRepoReq{Private: false, Username: "u"},
+		}).Return(&types.Dataset{Name: "d"}, nil)
+		tester.WithBody(t, &types.CreateDatasetReq{
+			CreateRepoReq: types.CreateRepoReq{Private: false},
+		}).Execute()
+
+		tester.ResponseEqSimple(t, 200, gin.H{
+			"data": &types.Dataset{Name: "d"},
+		})
+	})
+
+}
diff --git a/api/handler/dataset_test.go b/api/handler/dataset_test.go
@@ -41,32 +41,6 @@ func (t *DatasetTester) WithHandleFunc(fn func(h *DatasetHandler) gin.HandlerFun
 	return t
 }
 
-func TestDatasetHandler_Create(t *testing.T) {
-
-	t.Run("public", func(t *testing.T) {
-
-		tester := NewDatasetTester(t).WithHandleFunc(func(h *DatasetHandler) gin.HandlerFunc {
-			return h.Create
-		})
-		tester.RequireUser(t)
-
-		tester.mocks.sensitive.EXPECT().CheckRequestV2(tester.ctx, &types.CreateDatasetReq{
-			CreateRepoReq: types.CreateRepoReq{Private: true},
-		}).Return(true, nil)
-		tester.mocks.dataset.EXPECT().Create(tester.ctx, &types.CreateDatasetReq{
-			CreateRepoReq: types.CreateRepoReq{Private: true, Username: "u"},
-		}).Return(&types.Dataset{Name: "d"}, nil)
-		tester.WithBody(t, &types.CreateDatasetReq{
-			CreateRepoReq: types.CreateRepoReq{Private: true},
-		}).Execute()
-
-		tester.ResponseEqSimple(t, 200, gin.H{
-			"data": &types.Dataset{Name: "d"},
-		})
-	})
-
-}
-
 func TestDatasetHandler_Index(t *testing.T) {
 	cases := []struct {
 		sort   string