feat: inject current organization

api-server/controllers/controllersv1/organization.go

@@ -24,23 +24,10 @@ type organizationController struct {
 
 var OrganizationController = organizationController{}
 
-type GetOrganizationSchema struct {
-	OrgName string `header:"X-Yatai-Organization"`
-}
+type GetOrganizationSchema struct{}
 
 func (s *GetOrganizationSchema) GetOrganization(ctx context.Context) (*models.Organization, error) {
-	if s.OrgName == "" {
-		user, err := services.GetCurrentUser(ctx)
-		if err != nil {
-			return nil, err
-		}
-		return services.OrganizationService.GetUserOrganization(ctx, user.ID)
-	}
-	organization, err := services.OrganizationService.GetByName(ctx, s.OrgName)
-	if err != nil {
-		return nil, errors.Wrapf(err, "get organization %s", s.OrgName)
-	}
-	return organization, nil
+	return services.GetCurrentOrganization(ctx)
 }
 
 func (c *organizationController) canView(ctx context.Context, organization *models.Organization) error {

api-server/routes/routes.go

@@ -16,13 +16,14 @@ import (
 	"github.com/wI2L/fizz"
 	"github.com/wI2L/fizz/openapi"
 
+	"github.com/bentoml/yatai-common/consts"
+	"github.com/bentoml/yatai-common/utils"
 	"github.com/bentoml/yatai-schemas/schemasv1"
 	"github.com/bentoml/yatai/api-server/config"
 	"github.com/bentoml/yatai/api-server/controllers/controllersv1"
 	"github.com/bentoml/yatai/api-server/controllers/web"
 	"github.com/bentoml/yatai/api-server/models"
 	"github.com/bentoml/yatai/api-server/services"
-	"github.com/bentoml/yatai/common/consts"
 	"github.com/bentoml/yatai/common/scookie"
 	"github.com/bentoml/yatai/common/yataicontext"
 )
@@ -41,6 +42,21 @@ var staticFiles = map[string]string{
 
 const WebsocketConnectContextKey = "websocket-connect"
 
+func injectCurrentOrganization(c *gin.Context) {
+	orgName := c.GetHeader(consts.YataiOrganizationHeaderName)
+	if orgName == "" {
+		c.Next()
+		return
+	}
+	org, err := services.OrganizationService.GetByName(c, orgName)
+	if err != nil {
+		c.AbortWithError(http.StatusInternalServerError, err)
+		return
+	}
+	services.SetCurrentOrganization(c, org)
+	c.Next()
+}
+
 func NewRouter() (*fizz.Fizz, error) {
 	tonic.SetRenderHook(func(c *gin.Context, statusCode int, payload interface{}) {
 		if _, exists := c.Get(WebsocketConnectContextKey); exists {
@@ -66,6 +82,7 @@ func NewRouter() (*fizz.Fizz, error) {
 	engine := gin.New()
 
 	store := cookie.NewStore([]byte(config.YataiConfig.Server.SessionSecretKey))
+	engine.Use(injectCurrentOrganization)
 	engine.Use(sessions.Sessions("yatai-session-v1", store))
 
 	engine.GET("/logout", web.Logout)
@@ -279,7 +296,21 @@ func getLoginUser(ctx *gin.Context) (user *models.User, err error) {
 	}
 
 	yataicontext.SetUserName(ctx, user.Name)
-	services.SetLoginUser(ctx, user)
+	services.SetCurrentUser(ctx, user)
+	org, err := services.GetCurrentOrganization(ctx)
+	isNotFound := utils.IsNotFound(err)
+	if err != nil && !isNotFound {
+		err = errors.Wrap(err, "get current organization")
+		return
+	}
+	if isNotFound {
+		org, err = services.OrganizationService.GetUserOrganization(ctx, user.ID)
+		if err != nil {
+			err = errors.Wrap(err, "get user organization")
+			return
+		}
+		services.SetCurrentOrganization(ctx, org)
+	}
 	return
 }
 

api-server/services/api_token.go

@@ -162,12 +162,12 @@ func (s *apiTokenService) GetByToken(ctx context.Context, token string) (*models
 	if len(pieces) == 3 && pieces[0] == consts.YataiApiTokenPrefixYataiDeploymentOperator {
 		clusterName := pieces[1]
 		token_ := pieces[2]
-		defaultOrg, err := OrganizationService.GetDefault(ctx)
+		org, err := GetCurrentOrganization(ctx)
 		if err != nil {
 			err = errors.Wrap(err, "failed to get default organization")
 			return nil, err
 		}
-		cluster, err := ClusterService.GetByName(ctx, defaultOrg.ID, clusterName)
+		cluster, err := ClusterService.GetByName(ctx, org.ID, clusterName)
 		if err != nil {
 			err = errors.Wrapf(err, "failed to get cluster %s", clusterName)
 			return nil, err
@@ -192,7 +192,7 @@ func (s *apiTokenService) GetByToken(ctx context.Context, token string) (*models
 			return nil, err
 		}
 		var apiToken *models.ApiToken
-		apiToken, err = ApiTokenService.GetByName(ctx, defaultOrg.ID, adminUser.ID, consts.YataiK8sBotApiTokenName)
+		apiToken, err = ApiTokenService.GetByName(ctx, org.ID, adminUser.ID, consts.YataiK8sBotApiTokenName)
 		apiTokenIsNotFound := utils.IsNotFound(err)
 		if err != nil && !apiTokenIsNotFound {
 			err = errors.Wrapf(err, "get api token")
@@ -204,14 +204,14 @@ func (s *apiTokenService) GetByToken(ctx context.Context, token string) (*models
 			}
 			apiToken, err = ApiTokenService.Create(ctx, CreateApiTokenOption{
 				Name:           consts.YataiK8sBotApiTokenName,
-				OrganizationId: defaultOrg.ID,
+				OrganizationId: org.ID,
 				UserId:         adminUser.ID,
 				Description:    "yatai k8s bot api token",
 				Scopes:         &scopes,
 			})
 			if err != nil {
 				var err_ error
-				apiToken, err_ = ApiTokenService.GetByName(ctx, defaultOrg.ID, adminUser.ID, consts.YataiK8sBotApiTokenName)
+				apiToken, err_ = ApiTokenService.GetByName(ctx, org.ID, adminUser.ID, consts.YataiK8sBotApiTokenName)
 				if err_ != nil {
 					err = errors.Wrapf(err, "create api token %s", consts.YataiK8sBotApiTokenName)
 					return nil, err

api-server/services/organization.go

@@ -6,17 +6,18 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/gin-gonic/gin"
 	"github.com/minio/minio-go/v7"
 	"github.com/minio/minio-go/v7/pkg/credentials"
 	"github.com/pkg/errors"
 	"gorm.io/gorm"
 	"k8s.io/apimachinery/pkg/util/validation"
 
+	"github.com/bentoml/yatai-common/consts"
+	"github.com/bentoml/yatai-common/utils"
 	"github.com/bentoml/yatai-schemas/modelschemas"
 	"github.com/bentoml/yatai/api-server/config"
 	"github.com/bentoml/yatai/api-server/models"
-	"github.com/bentoml/yatai/common/consts"
-	"github.com/bentoml/yatai/common/utils"
 )
 
 type organizationService struct{}
@@ -407,3 +408,24 @@ func (s *organizationService) GetS3Config(ctx context.Context, org *models.Organ
 	err = errors.New("no s3 config")
 	return
 }
+
+const CurrentOrganizationKey = "currentOrganization"
+
+func SetCurrentOrganization(ctx *gin.Context, org *models.Organization) {
+	if org == nil {
+		return
+	}
+	ctx.Set(CurrentOrganizationKey, org)
+}
+
+func GetCurrentOrganization(ctx context.Context) (*models.Organization, error) {
+	org_ := ctx.Value(CurrentOrganizationKey)
+	if org_ == nil {
+		return nil, errors.Wrap(consts.ErrNotFound, "current organization")
+	}
+	org, ok := org_.(*models.Organization)
+	if !ok {
+		return nil, errors.New("current organization is not a organization")
+	}
+	return org, nil
+}

api-server/services/user.go

@@ -11,10 +11,10 @@ import (
 	"golang.org/x/crypto/bcrypt"
 	"gorm.io/gorm"
 
+	"github.com/bentoml/yatai-common/consts"
+	"github.com/bentoml/yatai-common/utils"
 	"github.com/bentoml/yatai-schemas/modelschemas"
 	"github.com/bentoml/yatai/api-server/models"
-	"github.com/bentoml/yatai/common/consts"
-	"github.com/bentoml/yatai/common/utils"
 )
 
 type userService struct{}
@@ -25,7 +25,7 @@ func (*userService) getBaseDB(ctx context.Context) *gorm.DB {
 	return mustGetSession(ctx).Model(&models.User{})
 }
 
-const LoginUserKey = "loginUser"
+const CurrentUserKey = "currentUser"
 
 type CreateUserOption struct {
 	Name      string
@@ -333,21 +333,21 @@ func (s *userService) GetAssociatedCreator(ctx context.Context, associate ICreat
 	return user, err
 }
 
-func SetLoginUser(ctx *gin.Context, user *models.User) {
+func SetCurrentUser(ctx *gin.Context, user *models.User) {
 	if user == nil {
 		return
 	}
-	ctx.Set(LoginUserKey, user)
+	ctx.Set(CurrentUserKey, user)
 }
 
 func GetCurrentUser(ctx context.Context) (*models.User, error) {
-	user_ := ctx.Value(LoginUserKey)
+	user_ := ctx.Value(CurrentUserKey)
 	if user_ == nil {
 		return nil, errors.Wrap(consts.ErrNotFound, "cannot find current user")
 	}
 	user, ok := user_.(*models.User)
 	if !ok {
-		return nil, errors.New("get login user err")
+		return nil, errors.New("get current user err, the type is not *models.User")
 	}
 	return user, nil
 }

go.mod

@@ -5,7 +5,7 @@ go 1.17
 require (
 	github.com/GoogleContainerTools/kaniko v1.8.1
 	github.com/bentoml/grafana-operator v1.4.1-0.20210927064226-14795530b647
-	github.com/bentoml/yatai-common v0.0.0-20220812134652-80086685eb60
+	github.com/bentoml/yatai-common v0.0.0-20220816152244-197bb805112a
 	github.com/bentoml/yatai-deployment-operator v0.0.0-20220509150025-8aae33c8b98b
 	github.com/bentoml/yatai-schemas v0.0.0-20220816014402-e1efcc5a8e9f
 	github.com/ghodss/yaml v1.0.0

go.sum

@@ -354,8 +354,8 @@ github.com/benbjohnson/clock v1.0.3/go.mod h1:bGMdMPoPVvcYyt1gHDf4J2KE153Yf9BuiU
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/bentoml/grafana-operator v1.4.1-0.20210927064226-14795530b647 h1:fKWmU8aM6tda4yThxj8SGKWJYCxjLW9gSM6WIIFFgvo=
 github.com/bentoml/grafana-operator v1.4.1-0.20210927064226-14795530b647/go.mod h1:w8sk7HbPMnOlZislz8Ch3pRw15FNMY6o1Q+53fepnhc=
-github.com/bentoml/yatai-common v0.0.0-20220812134652-80086685eb60 h1:n4rBerBySNm7z9uN42PLkHVKBIRT2JhMamX450AyrOA=
-github.com/bentoml/yatai-common v0.0.0-20220812134652-80086685eb60/go.mod h1:pox0XYk/bVUwKkadn0XwWHEbJmxSEeN3+HwGA4a8uOQ=
+github.com/bentoml/yatai-common v0.0.0-20220816152244-197bb805112a h1:AfHdjq19EQzRbPG10meuhD1xnQ5O5B4r7azz4Xm9fDo=
+github.com/bentoml/yatai-common v0.0.0-20220816152244-197bb805112a/go.mod h1:pox0XYk/bVUwKkadn0XwWHEbJmxSEeN3+HwGA4a8uOQ=
 github.com/bentoml/yatai-deployment-operator v0.0.0-20220509150025-8aae33c8b98b h1:FTDJ+5TMgBWGB8EYRvLyJKtYDnWKvE6CqqYnhY+TEe4=
 github.com/bentoml/yatai-deployment-operator v0.0.0-20220509150025-8aae33c8b98b/go.mod h1:vu1jltSCPUq+nKmO/ixQxAVKrYEcT/XXGIhluCMtLs8=
 github.com/bentoml/yatai-schemas v0.0.0-20220425180144-46a9d8e948f8/go.mod h1:OT5T7rrtdjJIkJpUazp1fQBbAVzl55YXyUJNVRLqskw=