From 356a707bfc7e4f9945868fbbcd94a03277d3e601 Mon Sep 17 00:00:00 2001
From: Jinjing Zhou <VoVAllen@users.noreply.github.com>
Date: Thu, 27 Oct 2022 19:44:15 +0800
Subject: [PATCH] feat: add config.repo in envd (#1101)

* add

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* fix lint

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* fix

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>
---
 envd/api/config/__init__.py                 |  9 +++++++++
 pkg/lang/frontend/starlark/config/config.go | 14 ++++++++++++++
 pkg/lang/frontend/starlark/config/const.go  |  1 +
 pkg/lang/ir/compile.go                      |  6 ++++++
 pkg/lang/ir/interface.go                    |  8 ++++++++
 pkg/lang/ir/types.go                        |  3 +++
 pkg/types/envd.go                           |  5 +++++
 pkg/types/label.go                          |  1 +
 8 files changed, 47 insertions(+)

diff --git a/envd/api/config/__init__.py b/envd/api/config/__init__.py
index 6b611cf4a..c2a5c0cd1 100644
--- a/envd/api/config/__init__.py
+++ b/envd/api/config/__init__.py
@@ -136,3 +136,12 @@ def rstudio_server():
     """
     Enable the RStudio Server (only work for `base(os="ubuntu20.04", language="r")`)
     """
+
+
+def repo(url: str, description: str):
+    """Setup repo related information. Will save to the image labels.
+
+    Args:
+        url (str): repo URL
+        description (str): repo description
+    """
diff --git a/pkg/lang/frontend/starlark/config/config.go b/pkg/lang/frontend/starlark/config/config.go
index 6df66da16..6436036f7 100644
--- a/pkg/lang/frontend/starlark/config/config.go
+++ b/pkg/lang/frontend/starlark/config/config.go
@@ -44,6 +44,7 @@ var Module = &starlarkstruct.Module{
 			ruleJuliaPackageServer, ruleFuncJuliaPackageServer),
 		"rstudio_server": starlark.NewBuiltin(ruleRStudioServer, ruleFuncRStudioServer),
 		"entrypoint":     starlark.NewBuiltin(ruleEntrypoint, ruleFuncEntrypoint),
+		"repo":           starlark.NewBuiltin(ruleRepo, ruleFuncRepo),
 	},
 }
 
@@ -212,3 +213,16 @@ func ruleFuncEntrypoint(thread *starlark.Thread, _ *starlark.Builtin,
 	ir.Entrypoint(argList)
 	return starlark.None, nil
 }
+
+func ruleFuncRepo(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var url, description string
+
+	if err := starlark.UnpackArgs(ruleRepo, args, kwargs, "url", &url, "description?", &description); err != nil {
+		return nil, err
+	}
+
+	logger.Debugf("repo info: url=%s, description=%s", url, description)
+	ir.Repo(url, description)
+	return starlark.None, nil
+}
diff --git a/pkg/lang/frontend/starlark/config/const.go b/pkg/lang/frontend/starlark/config/const.go
index a3edc4dc2..3e084f794 100644
--- a/pkg/lang/frontend/starlark/config/const.go
+++ b/pkg/lang/frontend/starlark/config/const.go
@@ -24,4 +24,5 @@ const (
 	ruleJuliaPackageServer = "config.julia_pkg_server"
 	ruleRStudioServer      = "config.rstudio_server"
 	ruleEntrypoint         = "config.entrypoint"
+	ruleRepo               = "config.repo"
 )
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
index 2b0c4be75..cc01690bd 100644
--- a/pkg/lang/ir/compile.go
+++ b/pkg/lang/ir/compile.go
@@ -181,6 +181,12 @@ func (g Graph) Labels() (map[string]string, error) {
 	}
 	labels[types.ImageLabelPorts] = string(portsData)
 
+	repoInfo, err := json.Marshal(g.Repo)
+	if err != nil {
+		return labels, err
+	}
+	labels[types.ImageLabelRepo] = string(repoInfo)
+
 	return labels, nil
 }
 
diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
index 58bd18948..5c47a9d30 100644
--- a/pkg/lang/ir/interface.go
+++ b/pkg/lang/ir/interface.go
@@ -19,6 +19,7 @@ import (
 	"github.com/opencontainers/go-digest"
 
 	"github.com/tensorchord/envd/pkg/editor/vscode"
+	"github.com/tensorchord/envd/pkg/types"
 )
 
 func Base(os, language, image string) error {
@@ -229,3 +230,10 @@ func RuntimeEnviron(env map[string]string) {
 func RuntimeInitScript(commands []string) {
 	DefaultGraph.RuntimeInitScript = append(DefaultGraph.RuntimeInitScript, commands)
 }
+
+func Repo(url, description string) {
+	DefaultGraph.Repo = types.RepoInfo{
+		Description: description,
+		URL:         url,
+	}
+}
diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
index 0355a7168..1872af215 100644
--- a/pkg/lang/ir/types.go
+++ b/pkg/lang/ir/types.go
@@ -19,6 +19,7 @@ import (
 
 	"github.com/tensorchord/envd/pkg/editor/vscode"
 	"github.com/tensorchord/envd/pkg/progress/compileui"
+	"github.com/tensorchord/envd/pkg/types"
 )
 
 // A Graph contains the state,
@@ -61,6 +62,8 @@ type Graph struct {
 	HTTP       []HTTPInfo
 	Entrypoint []string
 
+	Repo types.RepoInfo
+
 	*JupyterConfig
 	*GitConfig
 	*CondaConfig
diff --git a/pkg/types/envd.go b/pkg/types/envd.go
index 0b443f7a0..20b358490 100644
--- a/pkg/types/envd.go
+++ b/pkg/types/envd.go
@@ -145,6 +145,11 @@ type Dependency struct {
 	PyPIPackages []string `json:"pypi_packages,omitempty"`
 }
 
+type RepoInfo struct {
+	URL         string `json:"url,omitempty"`
+	Description string `json:"description,omitempty"`
+}
+
 type PortBinding struct {
 	Port     string
 	Protocol string
diff --git a/pkg/types/label.go b/pkg/types/label.go
index 061554a4d..baca480db 100644
--- a/pkg/types/label.go
+++ b/pkg/types/label.go
@@ -22,6 +22,7 @@ const (
 
 	ImageLabelVendor    = "ai.tensorchord.envd.vendor"
 	ImageLabelGPU       = "ai.tensorchord.envd.gpu"
+	ImageLabelRepo      = "ai.tensorchord.envd.repo"
 	ImageLabelPorts     = "ai.tensorchord.envd.ports"
 	ImageLabelAPT       = "ai.tensorchord.envd.apt.packages"
 	ImageLabelPyPI      = "ai.tensorchord.envd.pypi.commands"