From 9655db1a687724d773ae143496b1412f818acb45 Mon Sep 17 00:00:00 2001 From: Masselot Antoine Date: Fri, 29 Apr 2022 20:42:43 +0200 Subject: [PATCH] add flag --runtime-version to dev init command (#533) * add flag --runtime-version to dev init command * add warning message if user provided invalid runtime v in dev init * fix review --- airflow_versions/airflow_versions.go | 15 ++- airflow_versions/airflow_versions_test.go | 7 +- cmd/airflow.go | 25 ++++- cmd/utils.go | 4 +- cmd/utils_test.go | 128 ++++++++++++++++++++-- 5 files changed, 160 insertions(+), 19 deletions(-) diff --git a/airflow_versions/airflow_versions.go b/airflow_versions/airflow_versions.go index 19fe1f0f2..ae568f87d 100644 --- a/airflow_versions/airflow_versions.go +++ b/airflow_versions/airflow_versions.go @@ -19,7 +19,7 @@ func (e ErrNoTagAvailable) Error() string { } // GetDefaultImageTag returns default airflow image tag -func GetDefaultImageTag(httpClient *Client, airflowVersion string) (string, error) { +func GetDefaultImageTag(httpClient *Client, airflowVersion, userRuntimeVersion string) (string, error) { r := Request{} resp, err := r.DoWithClient(httpClient) @@ -31,15 +31,24 @@ func GetDefaultImageTag(httpClient *Client, airflowVersion string) (string, erro return getAstronomerCertifiedTag(resp.AvailableReleases, airflowVersion) } - return getAstroRuntimeTag(resp.RuntimeVersions, airflowVersion) + return getAstroRuntimeTag(resp.RuntimeVersions, airflowVersion, userRuntimeVersion) } // get latest runtime tag associated to provided airflow version // if no airflow version is provided, returns the latest astro runtime version available -func getAstroRuntimeTag(runtimeVersions map[string]RuntimeVersion, airflowVersion string) (string, error) { +func getAstroRuntimeTag(runtimeVersions map[string]RuntimeVersion, airflowVersion, userRuntimeVersion string) (string, error) { availableTags := []string{} availableVersions := []string{} + // If user wants a specific runtime version, check that it is a valid runtime released version + if userRuntimeVersion != "" { + if _, ok := runtimeVersions[userRuntimeVersion]; ok { + return userRuntimeVersion, nil + } + // user provided invalid runtime version, print warning + fmt.Printf("You provided an invalid runtime version %s, ignoring provided version...", userRuntimeVersion) + } + for runtimeVersion, r := range runtimeVersions { if r.Metadata.Channel != VersionChannelStable { continue diff --git a/airflow_versions/airflow_versions_test.go b/airflow_versions/airflow_versions_test.go index 771abc34d..82e08be27 100644 --- a/airflow_versions/airflow_versions_test.go +++ b/airflow_versions/airflow_versions_test.go @@ -145,16 +145,19 @@ func TestGetAstroRuntimeTag(t *testing.T) { tests := []struct { airflowVersion string + runtimeVersion string output string err error }{ {airflowVersion: "", output: "4.0.0", err: nil}, + {airflowVersion: "", runtimeVersion: "9.9.9", output: "4.0.0", err: nil}, {airflowVersion: "2.1.1", output: "3.0.1", err: nil}, + {airflowVersion: "", runtimeVersion: "3.1.0", output: "3.1.0", err: nil}, {airflowVersion: "2.2.2", output: "", err: ErrNoTagAvailable{airflowVersion: "2.2.2"}}, } for _, tt := range tests { - defaultImageTag, err := getAstroRuntimeTag(tagToRuntimeVersion, tt.airflowVersion) + defaultImageTag, err := getAstroRuntimeTag(tagToRuntimeVersion, tt.airflowVersion, tt.runtimeVersion) if tt.err == nil { assert.NoError(t, err) } else { @@ -176,7 +179,7 @@ func TestGetDefaultImageTagError(t *testing.T) { }) httpClient := NewClient(client, true) - defaultImageTag, err := GetDefaultImageTag(httpClient, "") + defaultImageTag, err := GetDefaultImageTag(httpClient, "", "") assert.Error(t, err) assert.Equal(t, "", defaultImageTag) } diff --git a/cmd/airflow.go b/cmd/airflow.go index 0379ab2a3..df37a000f 100644 --- a/cmd/airflow.go +++ b/cmd/airflow.go @@ -25,9 +25,10 @@ import ( ) var ( - errNotProjectDir = errors.New("not in a project directory") - errProjectConfigNotFound = errors.New("project config file does not exists") - errInvalidProjectName = errors.New(messages.ErrInvalidConfigProjectName) + errNotProjectDir = errors.New("not in a project directory") + errProjectConfigNotFound = errors.New("project config file does not exists") + errInvalidProjectName = errors.New(messages.ErrInvalidConfigProjectName) + errInvalidBothAirflowAndRuntimeVersions = errors.New("You provided both a Runtime version and an Airflow version. You have to provide only one of these to initialize your project.") //nolint ) var ( @@ -112,8 +113,12 @@ func newAirflowInitCmd(out io.Writer) *cobra.Command { }, } cmd.Flags().StringVarP(&projectName, "name", "n", "", "Name of airflow project") - cmd.Flags().StringVarP(&airflowVersion, "airflow-version", "v", "", "Version of airflow you want to deploy") + cmd.Flags().StringVarP(&airflowVersion, "airflow-version", "v", "", "Version of airflow you want to deploy. Only used if --use-astronomer-certified is true.") cmd.Flags().BoolVarP(&useAstronomerCertified, "use-astronomer-certified", "", false, "If specified, initializes a project using Astronomer Certified Airflow image instead of Astro Runtime.") + if appConfig != nil && appConfig.Flags.AstroRuntimeEnabled { + cmd.Flags().StringVarP(&runtimeVersion, "runtime-version", "r", "", "Version of astro-runtime you want to deploy. Only used if --use-astronomer-certified is false.") + } + return cmd } @@ -312,8 +317,18 @@ func airflowInit(cmd *cobra.Command, _ []string, out io.Writer) error { useAstronomerCertified = true } + // Validate runtimeVersion and airflowVersion + if airflowVersion != "" && runtimeVersion != "" { + return errInvalidBothAirflowAndRuntimeVersions + } + if useAstronomerCertified && runtimeVersion != "" { + fmt.Println("You provided a runtime version with the --use-astronomer-certified flag. Thus, this command will ignore the --runtime-version value you provided.") + runtimeVersion = "" + } + + // if runtime version is set, use runtime version. If user specifies --use-astronomer-certified, or use runtime but didn't provide version, get default image tag httpClient := airflowversions.NewClient(httputil.NewHTTPClient(), useAstronomerCertified) - defaultImageTag, err := prepareDefaultAirflowImageTag(airflowVersion, httpClient, houstonClient, out) + defaultImageTag, err := prepareDefaultAirflowImageTag(airflowVersion, runtimeVersion, httpClient, houstonClient, out) if err != nil { return err } diff --git a/cmd/utils.go b/cmd/utils.go index ef57005dc..bef07cfa8 100644 --- a/cmd/utils.go +++ b/cmd/utils.go @@ -13,7 +13,7 @@ import ( var errAirflowVersionNotSupported = errors.New("the --airflow-version flag is not supported if you're not authenticated to Astronomer. Please authenticate and try again") -func prepareDefaultAirflowImageTag(airflowVersion string, httpClient *airflowversions.Client, houstonClient houston.ClientInterface, out io.Writer) (string, error) { +func prepareDefaultAirflowImageTag(airflowVersion, userRuntimeVersion string, httpClient *airflowversions.Client, houstonClient houston.ClientInterface, out io.Writer) (string, error) { deploymentConfig, err := houstonClient.GetDeploymentConfig() if err == nil { acceptableAirflowVersions := deploymentConfig.AirflowVersions @@ -29,7 +29,7 @@ func prepareDefaultAirflowImageTag(airflowVersion string, httpClient *airflowver } } - defaultImageTag, _ := airflowversions.GetDefaultImageTag(httpClient, airflowVersion) + defaultImageTag, _ := airflowversions.GetDefaultImageTag(httpClient, airflowVersion, userRuntimeVersion) if defaultImageTag == "" { if useAstronomerCertified { diff --git a/cmd/utils_test.go b/cmd/utils_test.go index ee4d65225..5e1230810 100644 --- a/cmd/utils_test.go +++ b/cmd/utils_test.go @@ -67,7 +67,7 @@ func Test_prepareDefaultAirflowImageTag(t *testing.T) { output := new(bytes.Buffer) - myTests := []struct { + airflowTests := []struct { airflowVersion string expectedImageTag string expectedError string @@ -78,8 +78,122 @@ func Test_prepareDefaultAirflowImageTag(t *testing.T) { {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 { - defaultTag, err := prepareDefaultAirflowImageTag(tt.airflowVersion, httpClient, api, output) + for _, tt := range airflowTests { + 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_prepareDefaultImageTagRuntime(t *testing.T) { + runtimeResponse := ` +{ + "runtimeVersions":{ + "3.0.0":{ + "metadata":{ + "airflowVersion":"2.1.1", + "channel":"stable", + "releaseDate":"2021-08-12" + }, + "migrations":{ + "airflowDatabase":false + } + }, + "3.0.1":{ + "metadata":{ + "airflowVersion":"2.1.1", + "channel":"stable", + "releaseDate":"2021-08-31" + }, + "migrations":{ + "airflowDatabase":false + } + }, + "4.0.0":{ + "metadata":{ + "airflowVersion":"2.2.0", + "channel":"stable", + "releaseDate":"2021-10-12" + }, + "migrations":{ + "airflowDatabase":true + } + }, + "4.0.1":{ + "metadata":{ + "airflowVersion":"2.2.0", + "channel":"stable", + "releaseDate":"2021-10-25" + }, + "migrations":{ + "airflowDatabase":false + } + }, + "4.1.0":{ + "metadata":{ + "airflowVersion":"2.2.4", + "channel":"stable", + "releaseDate":"2022-02-22" + }, + "migrations":{ + "airflowDatabase":true + } + }, + "4.2.0":{ + "metadata":{ + "airflowVersion":"2.2.4", + "channel":"stable", + "releaseDate":"2022-03-08" + }, + "migrations":{ + "airflowDatabase":false + } + }, + "5.0.0":{ + "metadata":{ + "airflowVersion":"2.3.0", + "channel":"alpha", + "releaseDate":"2022-04-21" + }, + "migrations":{ + "airflowDatabase":true + } + } + }, + "schemaVersion":"1.3" +} +` + // prepare fake response from houston + api := new(mocks.ClientInterface) + api.On("GetDeploymentConfig").Return(mockDeploymentConfig, nil) + + output := new(bytes.Buffer) + + clientRuntime := testUtil.NewTestClient(func(req *http.Request) *http.Response { + return &http.Response{ + StatusCode: http.StatusOK, + Body: ioutil.NopCloser(bytes.NewBufferString(runtimeResponse)), + Header: make(http.Header), + } + }) + runtimeHTTPClient := airflowversions.NewClient(clientRuntime, false) + + runtimeTests := []struct { + runtimeVersion string + expectedImageTag string + expectedError string + }{ + {runtimeVersion: "3.0.1", expectedImageTag: "3.0.1", expectedError: ""}, + {runtimeVersion: "", expectedImageTag: "4.2.0", expectedError: ""}, + {runtimeVersion: "9.9.9", expectedImageTag: "4.2.0", expectedError: ""}, + } + + for _, tt := range runtimeTests { + defaultTag, err := prepareDefaultAirflowImageTag("", tt.runtimeVersion, runtimeHTTPClient, api, output) if tt.expectedError != "" { assert.EqualError(t, err, tt.expectedError) } else { @@ -113,7 +227,7 @@ func Test_fallbackDefaultAirflowImageTag(t *testing.T) { output := new(bytes.Buffer) - defaultTag, err := prepareDefaultAirflowImageTag("", httpClient, api, output) + defaultTag, err := prepareDefaultAirflowImageTag("", "", httpClient, api, output) assert.NoError(t, err) assert.Equal(t, "2.0.0-buster-onbuild", defaultTag) }) @@ -139,7 +253,7 @@ func Test_fallbackDefaultAirflowImageTag(t *testing.T) { output := new(bytes.Buffer) - defaultTag, err := prepareDefaultAirflowImageTag("", httpClient, api, output) + defaultTag, err := prepareDefaultAirflowImageTag("", "", httpClient, api, output) assert.NoError(t, err) assert.Equal(t, "3.0.0", defaultTag) }) @@ -183,7 +297,7 @@ func Test_prepareDefaultAirflowImageTagHoustonBadRequest(t *testing.T) { output := new(bytes.Buffer) - defaultTag, err := prepareDefaultAirflowImageTag("2.0.2", httpClient, api, output) + defaultTag, err := prepareDefaultAirflowImageTag("2.0.2", "", httpClient, api, output) assert.Equal(t, mockErrorResponse.Error(), err.Error()) assert.Equal(t, "", defaultTag) } @@ -226,7 +340,7 @@ func Test_prepareDefaultAirflowImageTagHoustonUnauthedRequest(t *testing.T) { output := new(bytes.Buffer) - defaultTag, err := prepareDefaultAirflowImageTag("2.0.2", httpClient, api, output) + defaultTag, err := prepareDefaultAirflowImageTag("2.0.2", "", httpClient, api, output) assert.Equal(t, "the --airflow-version flag is not supported if you're not authenticated to Astronomer. Please authenticate and try again", err.Error()) assert.Equal(t, "", defaultTag) }