From 03497b36715590647e70ac0cc48e15c6b6928044 Mon Sep 17 00:00:00 2001 From: Neel Dalsania <92356010+neel-astro@users.noreply.github.com> Date: Mon, 22 Nov 2021 17:11:16 +0530 Subject: [PATCH] fixed default airflow image tag behaviour (#450) * fixed default airflow image tag behaviour * added test to cover new default airflow image tag logic * fixed the output string for selecting default airflow image * added code review changes --- airflow_versions/airflow_versions.go | 33 ++++++++++- airflow_versions/airflow_versions_test.go | 35 ++++++++++- cmd/utils.go | 15 ++--- cmd/utils_test.go | 72 +++++++++++++++++++++-- 4 files changed, 135 insertions(+), 20 deletions(-) diff --git a/airflow_versions/airflow_versions.go b/airflow_versions/airflow_versions.go index 9950b8854..5042212a2 100644 --- a/airflow_versions/airflow_versions.go +++ b/airflow_versions/airflow_versions.go @@ -3,8 +3,11 @@ package airflowversions import ( "fmt" "sort" + "strings" ) +var tagPrefixOrder = []string{"buster-onbuild", "onbuild", "buster"} + // GetDefaultImageTag returns default airflow image tag func GetDefaultImageTag(httpClient *Client, airflowVersion string) (string, error) { r := Request{} @@ -14,15 +17,39 @@ func GetDefaultImageTag(httpClient *Client, airflowVersion string) (string, erro return "", err } + availableTags := []string{} vs := make(AirflowVersions, len(resp.AvailableReleases)) for i, r := range resp.AvailableReleases { + if r.Version == airflowVersion { + availableTags = r.Tags + break + } v, err := NewAirflowVersion(r.Version, r.Tags) if err == nil { vs[i] = v } } - sort.Sort(vs) - maxAvailableVersion := vs[len(vs)-1] - return fmt.Sprintf("%s-buster-onbuild", maxAvailableVersion.Coerce()), nil + var selectedVersion *AirflowVersion + if airflowVersion == "" && len(vs) != 0 { + sort.Sort(vs) + selectedVersion = vs[len(vs)-1] + availableTags = selectedVersion.tags + } else { + selectedVersion, err = NewAirflowVersion(airflowVersion, availableTags) + if err != nil { + return "", err + } + } + + for tagIndex := range tagPrefixOrder { + for idx := range availableTags { + if strings.HasPrefix(availableTags[idx], selectedVersion.Coerce()) && strings.HasSuffix(availableTags[idx], tagPrefixOrder[tagIndex]) { + return fmt.Sprintf("%s-%s", selectedVersion.Coerce(), tagPrefixOrder[tagIndex]), nil + } + } + } + + // case when airflowVersion requested is not present in certified astronomer endpoint, but is valid version as per Houston configuration + return airflowVersion, nil } diff --git a/airflow_versions/airflow_versions_test.go b/airflow_versions/airflow_versions_test.go index 43d98303b..377cde8f3 100644 --- a/airflow_versions/airflow_versions_test.go +++ b/airflow_versions/airflow_versions_test.go @@ -8,6 +8,8 @@ import ( testUtil "github.com/astronomer/astro-cli/pkg/testing" "github.com/stretchr/testify/assert" + + "github.com/Masterminds/semver" ) func TestGetDefaultImageTag(t *testing.T) { @@ -53,6 +55,14 @@ func TestGetDefaultImageTag(t *testing.T) { "1.10.4-11-buster" ], "channel": "stable" + }, + { + "version": "2.2.0", + "level": "new_feature", + "url": "https://github.com/astronomer/airflow/releases/tag/v2.2.0%2Bastro.2", + "release_date": "2021-10-14T12:46:00+00:00", + "tags": ["2.2.0", "2.2.0-onbuild"], + "channel": "stable" } ] }` @@ -65,9 +75,28 @@ func TestGetDefaultImageTag(t *testing.T) { }) httpClient := NewClient(client) - defaultImageTag, err := GetDefaultImageTag(httpClient, "") - assert.NoError(t, err) - assert.Equal(t, "1.10.5-buster-onbuild", defaultImageTag) + tests := []struct { + airflowVersion string + output string + err error + }{ + {airflowVersion: "", output: "2.2.0-onbuild", err: nil}, + {airflowVersion: "1.10.5", output: "1.10.5-buster-onbuild", err: nil}, + {airflowVersion: "1.10.4-rc.1", output: "1.10.4-rc.1", err: nil}, + {airflowVersion: "2.2.1", output: "2.2.1", err: nil}, + {airflowVersion: "2.2.0", output: "2.2.0-onbuild", err: nil}, + {airflowVersion: "2.2.x", output: "", err: semver.ErrInvalidSemVer}, + } + + for _, tt := range tests { + defaultImageTag, err := GetDefaultImageTag(httpClient, tt.airflowVersion) + if tt.err == nil { + assert.NoError(t, err) + } else { + assert.EqualError(t, err, tt.err.Error()) + } + assert.Equal(t, tt.output, defaultImageTag) + } } func TestGetDefaultImageTagError(t *testing.T) { diff --git a/cmd/utils.go b/cmd/utils.go index ef16f66c4..6f78845f7 100644 --- a/cmd/utils.go +++ b/cmd/utils.go @@ -12,8 +12,6 @@ import ( ) func prepareDefaultAirflowImageTag(airflowVersion string, httpClient *airflowversions.Client, houstonClient *houston.Client, out io.Writer) (string, error) { - defaultImageTag, _ := airflowversions.GetDefaultImageTag(httpClient, "") - r := houston.Request{ Query: houston.DeploymentInfoRequest, } @@ -25,23 +23,20 @@ func prepareDefaultAirflowImageTag(airflowVersion string, httpClient *airflowver if airflowVersion != "" && !acceptableVersion(airflowVersion, acceptableAirflowVersions) { return "", errors.Errorf(messages.ErrInvalidAirflowVersion, strings.Join(acceptableAirflowVersions, ", ")) } - if airflowVersion == "" { - defaultImageTag = "" - } else { - defaultImageTag = fmt.Sprintf("%s-buster-onbuild", airflowVersion) - } } else if airflowVersion != "" { switch t := err; t { - default: - return "", err case houston.ErrVerboseInaptPermissions: return "", errors.New("the --airflow-version flag is not supported if you're not authenticated to Astronomer. Please authenticate and try again") + default: + return "", err } } + defaultImageTag, _ := airflowversions.GetDefaultImageTag(httpClient, airflowVersion) + if defaultImageTag == "" { defaultImageTag = "2.0.0-buster-onbuild" - fmt.Fprintf(out, "Initializing Airflow project\nNot connected to Astronomer, pulling Airflow development files from %s\n", defaultImageTag) + fmt.Fprintf(out, "Initializing Airflow project, pulling Airflow development files from %s\n", defaultImageTag) } return defaultImageTag, nil } diff --git a/cmd/utils_test.go b/cmd/utils_test.go index 500cfbdd9..910b4f591 100644 --- a/cmd/utils_test.go +++ b/cmd/utils_test.go @@ -78,10 +78,10 @@ func Test_prepareDefaultAirflowImageTag(t *testing.T) { expectedImageTag string expectedError string }{ - {airflowVersion: "1.10.14", expectedImageTag: "1.10.14-buster-onbuild", expectedError: ""}, - {airflowVersion: "1.10.15", expectedImageTag: "1.10.15-buster-onbuild", expectedError: ""}, - {airflowVersion: "", expectedImageTag: "2.0.0-buster-onbuild", expectedError: ""}, - {airflowVersion: "2.0.2", expectedImageTag: "2.0.2-buster-onbuild", expectedError: ""}, + {airflowVersion: "1.10.14", expectedImageTag: "1.10.14", expectedError: ""}, + {airflowVersion: "1.10.15", expectedImageTag: "1.10.15", expectedError: ""}, + {airflowVersion: "", expectedImageTag: "1.10.5-buster-onbuild", expectedError: ""}, + {airflowVersion: "2.0.2", expectedImageTag: "2.0.2", expectedError: ""}, {airflowVersion: "9.9.9", expectedImageTag: "", expectedError: "Unsupported Airflow Version specified. Please choose from: 2.1.0, 2.0.2, 2.0.0, 1.10.15, 1.10.14, 1.10.12, 1.10.10, 1.10.7, 1.10.5 \n"}, } for _, tt := range myTests { @@ -95,6 +95,70 @@ func Test_prepareDefaultAirflowImageTag(t *testing.T) { } } +func Test_fallbackDefaultAirflowImageTag(t *testing.T) { + testUtil.InitTestConfig() + + // prepare fake response from updates.astronomer.io + okResponse := `{ + "version": "1.0", + "available_releases": [] +}` + client := testUtil.NewTestClient(func(req *http.Request) *http.Response { + return &http.Response{ + StatusCode: 200, + Body: ioutil.NopCloser(bytes.NewBufferString(okResponse)), + Header: make(http.Header), + } + }) + httpClient := airflowversions.NewClient(client) + + // prepare fake response from houston + ok := `{ + "data": { + "deploymentConfig": { + "airflowVersions": [ + "2.1.0", + "2.0.2", + "2.0.0", + "1.10.15", + "1.10.14", + "1.10.12", + "1.10.10", + "1.10.7", + "1.10.5" + ] + } + } + }` + houstonClient := testUtil.NewTestClient(func(req *http.Request) *http.Response { + return &http.Response{ + StatusCode: 200, + Body: ioutil.NopCloser(bytes.NewBufferString(ok)), + Header: make(http.Header), + } + }) + api := houston.NewHoustonClient(houstonClient) + + output := new(bytes.Buffer) + + myTests := []struct { + airflowVersion string + expectedImageTag string + expectedError string + }{ + {airflowVersion: "", expectedImageTag: "2.0.0-buster-onbuild", expectedError: ""}, + } + for _, tt := range myTests { + defaultTag, err := prepareDefaultAirflowImageTag(tt.airflowVersion, httpClient, api, output) + if tt.expectedError != "" { + assert.EqualError(t, err, tt.expectedError) + } else { + assert.NoError(t, err) + } + assert.Equal(t, tt.expectedImageTag, defaultTag) + } +} + func Test_prepareDefaultAirflowImageTagHoustonBadRequest(t *testing.T) { testUtil.InitTestConfig() mockErrorResponse := `An error occurred`