tensorchord · kemingy · Dec 10, 2022 · Dec 8, 2022 · Dec 9, 2022
@@ -45,6 +45,7 @@ var Module = &starlarkstruct.Module{
 		"rstudio_server": starlark.NewBuiltin(ruleRStudioServer, ruleFuncRStudioServer),
 		"entrypoint":     starlark.NewBuiltin(ruleEntrypoint, ruleFuncEntrypoint),
 		"repo":           starlark.NewBuiltin(ruleRepo, ruleFuncRepo),
+		"owner":          starlark.NewBuiltin
10000
(ruleOwner, ruleFuncOwner),
 	},
 }
 
@@ -225,3 +226,24 @@ func ruleFuncRepo(thread *starlark.Thread, _ *starlark.Builtin,
 	ir.Repo(url, description)
 	return starlark.None, nil
 }
+
+func ruleFuncOwner(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var (
+		envdUid = starlark.MakeInt(-1)
+		envdGid = starlark.MakeInt(-1)
+	)
+
+	if err := starlark.UnpackArgs(ruleOwner, args, kwargs, "uid", &envdUid, "gid", &envdGid); err != nil {
+		return nil, err
+	}
+	uid, _ := envdUid.Int64()
+	gid, _ := envdGid.Int64()
+	if uid < 0 || uid > 65535 || gid < 0 || gid > 65535 {
+		err := errors.New("get a wrong uid or gid")
+		return nil, err
+	}
+	logger.Debugf("owner info: uid=%d, gid=%d", uid, gid)
+	ir.Owner(int(uid), int(gid))
+	return starlark.None, nil
+}
@@ -25,4 +25,5 @@ const (
 	ruleRStudioServer      = "config.rstudio_server"
 	ruleEntrypoint         = "config.entrypoint"
 	ruleRepo               = "config.repo"
+	ruleOwner              = "config.owner"
 )
@@ -112,7 +112,7 @@ func (g *generalGraph) Compile(ctx context.Context, envName string, pub string)
 	g.EnvironmentName = envName
 	g.PublicKeyPath = pub
 
-	uid, gid, err := getUIDGID()
+	uid, gid, err := g.getUIDGID()
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get uid/gid")
 	}

@@ -330,3 +330,9 @@ func Repo(url, description string) {
 		URL:         url,
 	}
 }
+
+func Owner(uid, gid int) {
+	g := DefaultGraph.(*generalGraph)
+	g.uid = uid
+	g.gid = gid
+}
@@ -25,8 +25,8 @@ import (
 // such as its call stack and thread-local storage.
 // TODO(gaocegeg): Refactor it to support order.
 type generalGraph struct {
-	uid int
-	gid int
+	uid int `default:"-1"`
+	gid int `default:"-1"`
 
 	ir.Language
 	Image string

@@ -65,7 +65,10 @@ func parseLanguage(l string) (string, *string, error) {
 	}
 }
 
-func getUIDGID() (int, int, error) {
+func (g generalGraph) getUIDGID() (int, int, error) {
+	// Firstly check cli flag set uid and git.
+	// if not set, then use the config.owner statements in envd.build
+	// The current user is the last choice
 	owner := viper.GetString(flag.FlagBuildOwner)
 	if len(owner) > 0 {
 		logrus.WithField
BEBD
("flag", owner).Info("use owner")
@@ -88,6 +91,9 @@ func getUIDGID() (int, int, error) {
 		}
 		return uid, gid, nil
 	}
+	if g.uid != -1 && g.gid != -1 {
+		return g.gid, g.gid, nil
+	}
 	user, err := user.Current()
 	if err != nil {
 		return 0, 0, errors.Wrap(err, "failed to get uid/gid")

@@ -147,6 +147,11 @@ type RepoInfo struct {
 	Description string `json:"description,omitempty"`
 }
 
+type OwnerInfo struct {
+	Uid int64 `json:"uid,omitempty"`
+	Gid int64 `json:"gid,omitempty"`
+}
+
 type PortBinding struct {
 	Name     string
 	Port     string