From fa07f37840755f3a6e7fc7253cfee1cc1de08803 Mon Sep 17 00:00:00 2001
From: novahow <58504997+novahow@users.noreply.github.com>
Date: Tue, 2 Jul 2024 08:30:32 +0800
Subject: [PATCH] reimplement functions in go-mouff-update, use ghreposervice
 (#5470)

* reimplement functions in go-mouff-update, use ghreposervice

Signed-off-by: novahow <b08902047@ntu.edu.tw>

* add tests

Signed-off-by: novahow <b08902047@ntu.edu.tw>

* fix version error

Signed-off-by: novahow <b08902047@ntu.edu.tw>

* remove testmain

Signed-off-by: novahow <b08902047@ntu.edu.tw>

* Run make lint-fix

Signed-off-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>

* Call TearDown

Signed-off-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>

---------

Signed-off-by: novahow <b08902047@ntu.edu.tw>
Signed-off-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>
Co-authored-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>
Signed-off-by: Vladyslav Libov <vladyslav.libov@d-fine.de>
---
 flytectl/Makefile                    |   2 +-
 flytectl/cmd/upgrade/upgrade.go      |   2 +-
 flytectl/cmd/upgrade/upgrade_test.go |   5 -
 flytectl/cmd/version/version.go      |   2 +-
 flytectl/cmd/version/version_test.go |   3 +
 flytectl/pkg/github/githubutil.go    |   4 +-
 flytectl/pkg/github/provider.go      | 175 +++++++++++++++++++++++++++
 flytectl/pkg/github/provider_test.go |  68 +++++++++++
 8 files changed, 251 insertions(+), 10 deletions(-)
 create mode 100644 flytectl/pkg/github/provider.go
 create mode 100644 flytectl/pkg/github/provider_test.go

diff --git a/flytectl/Makefile b/flytectl/Makefile
index fd32a32415..17c8b5f2bc 100644
--- a/flytectl/Makefile
+++ b/flytectl/Makefile
@@ -4,7 +4,7 @@ include ../boilerplate/flyte/docker_build/Makefile
 include ../boilerplate/flyte/golang_test_targets/Makefile
 include ../boilerplate/flyte/end2end/Makefile
 
-GIT_VERSION := $(shell git describe --always --tags)
+GIT_VERSION := $(shell git describe --dirty --tags --long --match 'flytectl/*' --first-parent | sed 's/^flytectl\///')
 GIT_HASH := $(shell git rev-parse --short HEAD)
 TIMESTAMP := $(shell date '+%Y-%m-%d')
 PACKAGE ?=github.com/flyteorg/flyte/flytestdlib
diff --git a/flytectl/cmd/upgrade/upgrade.go b/flytectl/cmd/upgrade/upgrade.go
index c4ad132c30..485db6cc30 100644
--- a/flytectl/cmd/upgrade/upgrade.go
+++ b/flytectl/cmd/upgrade/upgrade.go
@@ -110,7 +110,7 @@ func upgrade(u *updater.Updater) (string, error) {
 }
 
 func isUpgradeSupported(goos platformutil.Platform) (bool, error) {
-	latest, err := github.FlytectlReleaseConfig.GetLatestVersion()
+	latest, err := github.FlytectlReleaseConfig.Provider.(*github.GHProvider).GetCleanLatestVersion()
 	if err != nil {
 		return false, err
 	}
diff --git a/flytectl/cmd/upgrade/upgrade_test.go b/flytectl/cmd/upgrade/upgrade_test.go
index a8a955639e..d4132f1df4 100644
--- a/flytectl/cmd/upgrade/upgrade_test.go
+++ b/flytectl/cmd/upgrade/upgrade_test.go
@@ -1,7 +1,6 @@
 package upgrade
 
 import (
-	"fmt"
 	"sort"
 	"testing"
 
@@ -172,7 +171,3 @@ func TestSelfUpgradeRollback(t *testing.T) {
 	})
 
 }
-
-func TestMain(_ *testing.M) {
-	fmt.Println("Skipping due to https://github.com/flyteorg/flyte/issues/5372")
-}
diff --git a/flytectl/cmd/version/version.go b/flytectl/cmd/version/version.go
index 38b61d5538..88da1330a2 100644
--- a/flytectl/cmd/version/version.go
+++ b/flytectl/cmd/version/version.go
@@ -51,7 +51,7 @@ func GetVersionCommand(rootCmd *cobra.Command) map[string]cmdCore.CommandEntry {
 
 func getVersion(ctx context.Context, args []string, cmdCtx cmdCore.CommandContext) error {
 	goos := platformutil.Platform(runtime.GOOS)
-	version, err := github.FlytectlReleaseConfig.GetLatestVersion()
+	version, err := github.FlytectlReleaseConfig.Provider.(*github.GHProvider).GetCleanLatestVersion()
 	if err != nil {
 		logger.Error(ctx, "Unable to get the latest version because %v", err)
 	} else {
diff --git a/flytectl/cmd/version/version_test.go b/flytectl/cmd/version/version_test.go
index a1abbd7883..791a895e46 100644
--- a/flytectl/cmd/version/version_test.go
+++ b/flytectl/cmd/version/version_test.go
@@ -55,6 +55,7 @@ func TestVersionCommand(t *testing.T) {
 func TestVersionCommandFunc(t *testing.T) {
 	ctx := context.Background()
 	s := testutils.Setup()
+	defer s.TearDown()
 	stdlibversion.Build = ""
 	stdlibversion.BuildTime = ""
 	stdlibversion.Version = testVersion
@@ -67,6 +68,7 @@ func TestVersionCommandFunc(t *testing.T) {
 func TestVersionCommandFuncError(t *testing.T) {
 	ctx := context.Background()
 	s := testutils.Setup()
+	defer s.TearDown()
 	stdlibversion.Build = ""
 	stdlibversion.BuildTime = ""
 	stdlibversion.Version = "v"
@@ -79,6 +81,7 @@ func TestVersionCommandFuncError(t *testing.T) {
 func TestVersionCommandFuncErr(t *testing.T) {
 	ctx := context.Background()
 	s := testutils.Setup()
+	defer s.TearDown()
 	stdlibversion.Build = ""
 	stdlibversion.BuildTime = ""
 	stdlibversion.Version = testVersion
diff --git a/flytectl/pkg/github/githubutil.go b/flytectl/pkg/github/githubutil.go
index 5e639006e5..4de769d352 100644
--- a/flytectl/pkg/github/githubutil.go
+++ b/flytectl/pkg/github/githubutil.go
@@ -14,7 +14,6 @@ import (
 	"github.com/flyteorg/flyte/flytestdlib/logger"
 	stdlibversion "github.com/flyteorg/flyte/flytestdlib/version"
 	"github.com/google/go-github/v42/github"
-	"github.com/mouuff/go-rocket-update/pkg/provider"
 	"github.com/mouuff/go-rocket-update/pkg/updater"
 	"golang.org/x/oauth2"
 	"golang.org/x/text/cases"
@@ -39,9 +38,10 @@ var Client GHRepoService
 
 // FlytectlReleaseConfig represent the updater config for flytectl binary
 var FlytectlReleaseConfig = &updater.Updater{
-	Provider: &provider.Github{
+	Provider: &GHProvider{
 		RepositoryURL: flytectlRepository,
 		ArchiveName:   getFlytectlAssetName(),
+		ghRepo:        GetGHRepoService(),
 	},
 	ExecutableName: flytectl,
 	Version:        stdlibversion.Version,
diff --git a/flytectl/pkg/github/provider.go b/flytectl/pkg/github/provider.go
new file mode 100644
index 0000000000..9c8407052b
--- /dev/null
+++ b/flytectl/pkg/github/provider.go
@@ -0,0 +1,175 @@
+package github
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strings"
+
+	go_github "github.com/google/go-github/v42/github"
+	"github.com/mouuff/go-rocket-update/pkg/provider"
+)
+
+// Github provider finds a archive file in the repository's releases to provide files
+type GHProvider struct {
+	RepositoryURL string // Repository URL, example github.com/mouuff/go-rocket-update
+	ArchiveName   string // Archive name (the zip/tar.gz you upload for a release on github), example: binaries.zip
+
+	tmpDir             string            // temporary directory this is used internally
+	decompressProvider provider.Provider // provider used to decompress the downloaded archive
+	archivePath        string            // path to the downloaded archive (should be in tmpDir)
+	ghRepo             GHRepoService     // github repository service
+}
+
+// githubRepositoryInfo is used to get the name of the project and the owner name
+// from this fields we are able to get other links (such as the release and tags link)
+type githubRepositoryInfo struct {
+	RepositoryOwner string
+	RepositoryName  string
+}
+
+// getRepositoryInfo parses the github repository URL
+func (c *GHProvider) repositoryInfo() (*githubRepositoryInfo, error) {
+	re := regexp.MustCompile(`github\.com/(.*?)/(.*?)$`)
+	submatches := re.FindAllStringSubmatch(c.RepositoryURL, 1)
+	if len(submatches) < 1 {
+		return nil, errors.New("Invalid github URL:" + c.RepositoryURL)
+	}
+	return &githubRepositoryInfo{
+		RepositoryOwner: submatches[0][1],
+		RepositoryName:  submatches[0][2],
+	}, nil
+}
+
+// getArchiveURL get the archive URL for the github repository
+// If no tag is provided then the latest version is selected
+func (c *GHProvider) getArchiveURL(tag string) (string, error) {
+	if len(tag) == 0 {
+		// Get latest version if no tag is provided
+		var err error
+		tag, err = c.GetLatestVersion()
+		if err != nil {
+			return "", err
+		}
+	}
+
+	info, err := c.repositoryInfo()
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/%s",
+		info.RepositoryOwner,
+		info.RepositoryName,
+		tag,
+		c.ArchiveName,
+	), nil
+}
+
+// Open opens the provider
+func (c *GHProvider) Open() (err error) {
+	archiveURL, err := c.getArchiveURL("") // get archive url for latest version
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequest("GET", archiveURL, nil)
+	if err != nil {
+		return err
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	c.tmpDir, err = os.MkdirTemp("", "rocket-update")
+	if err != nil {
+		return err
+	}
+
+	c.archivePath = filepath.Join(c.tmpDir, c.ArchiveName)
+	archiveFile, err := os.Create(c.archivePath)
+	if err != nil {
+		return err
+	}
+	_, err = io.Copy(archiveFile, resp.Body)
+	archiveFile.Close()
+	if err != nil {
+		return err
+	}
+	c.decompressProvider, err = provider.Decompress(c.archivePath)
+	if err != nil {
+		return err
+	}
+	return c.decompressProvider.Open()
+}
+
+// Close closes the provider
+func (c *GHProvider) Close() error {
+	if c.decompressProvider != nil {
+		c.decompressProvider.Close()
+		c.decompressProvider = nil
+	}
+
+	if len(c.tmpDir) > 0 {
+		os.RemoveAll(c.tmpDir)
+		c.tmpDir = ""
+		c.archivePath = ""
+	}
+	return nil
+}
+
+// GetLatestVersion gets the latest version
+func (c *GHProvider) GetLatestVersion() (string, error) {
+	tags, err := c.getReleases()
+	if err != nil {
+		return "", err
+	}
+	latestTag := tags[0].GetTagName()
+	return latestTag, err
+}
+
+// GetCleanLatestVersion gets the latest version without the "flytectl/" prefix
+func (c *GHProvider) GetCleanLatestVersion() (string, error) {
+	latest, err := c.GetLatestVersion()
+	if err != nil {
+		return "", err
+	}
+	clearVersion := strings.TrimPrefix(latest, fmt.Sprintf("%s/", flytectl))
+	return clearVersion, nil
+}
+
+func (c *GHProvider) getReleases() ([]*go_github.RepositoryRelease, error) {
+	g := c.ghRepo
+	releases, _, err := g.ListReleases(context.Background(), owner, flyte, &go_github.ListOptions{
+		PerPage: 100,
+	})
+	if err != nil {
+		return nil, err
+	}
+	var filteredReleases []*go_github.RepositoryRelease
+	for _, release := range releases {
+		if strings.HasPrefix(release.GetTagName(), flytectl) {
+			filteredReleases = append(filteredReleases, release)
+		}
+	}
+	return filteredReleases, err
+}
+
+// Walk walks all the files provided
+func (c *GHProvider) Walk(walkFn provider.WalkFunc) error {
+	if c.decompressProvider == nil {
+		// TODO specify error
+		return provider.ErrNotOpenned
+	}
+	return c.decompressProvider.Walk(walkFn)
+}
+
+// Retrieve file relative to "provider" to destination
+func (c *GHProvider) Retrieve(src string, dest string) error {
+	return c.decompressProvider.Retrieve(src, dest)
+}
diff --git a/flytectl/pkg/github/provider_test.go b/flytectl/pkg/github/provider_test.go
new file mode 100644
index 0000000000..e342ec04e9
--- /dev/null
+++ b/flytectl/pkg/github/provider_test.go
@@ -0,0 +1,68 @@
+package github
+
+import (
+	"testing"
+
+	"github.com/flyteorg/flyte/flytectl/pkg/github/mocks"
+	go_github "github.com/google/go-github/v42/github"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestGetLatestFlytectlVersion(t *testing.T) {
+	t.Run("Get latest release", func(t *testing.T) {
+		mockGh := &mocks.GHRepoService{}
+		// return a list of github releases
+		mockGh.OnListReleasesMatch(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
+			[]*go_github.RepositoryRelease{
+				{TagName: go_github.String("flytectl/1.2.4")},
+				{TagName: go_github.String("flytectl/1.2.3")},
+				{TagName: go_github.String("other-1.0.0")},
+			},
+			nil,
+			nil,
+		)
+		mockProvider := &GHProvider{
+			RepositoryURL: flytectlRepository,
+			ArchiveName:   getFlytectlAssetName(),
+			ghRepo:        mockGh,
+		}
+
+		latestVersion, err := mockProvider.GetLatestVersion()
+		assert.Nil(t, err)
+		assert.Equal(t, "flytectl/1.2.4", latestVersion)
+		cleanVersion, err := mockProvider.GetCleanLatestVersion()
+		assert.Nil(t, err)
+		assert.Equal(t, "1.2.4", cleanVersion)
+	})
+}
+
+func TestGetFlytectlReleases(t *testing.T) {
+	t.Run("Get releases", func(t *testing.T) {
+		mockGh := &mocks.GHRepoService{}
+		allReleases := []*go_github.RepositoryRelease{
+			{TagName: go_github.String("flytectl/1.2.4")},
+			{TagName: go_github.String("flytectl/1.2.3")},
+			{TagName: go_github.String("other-1.0.0")},
+		}
+		releases := []*go_github.RepositoryRelease{
+			{TagName: go_github.String("flytectl/1.2.4")},
+			{TagName: go_github.String("flytectl/1.2.3")},
+		}
+		// return a list of github releases
+		mockGh.OnListReleasesMatch(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
+			allReleases,
+			nil,
+			nil,
+		)
+		mockProvider := &GHProvider{
+			RepositoryURL: flytectlRepository,
+			ArchiveName:   getFlytectlAssetName(),
+			ghRepo:        mockGh,
+		}
+
+		flytectlReleases, err := mockProvider.getReleases()
+		assert.Nil(t, err)
+		assert.Equal(t, releases, flytectlReleases)
+	})
+}