Skip to content

Commit

Permalink
add flag --runtime-version to dev init command (#533)
Browse files Browse the repository at this point in the history 
* add flag --runtime-version to dev init command

* add warning message if user provided invalid runtime v in dev init

* fix review
  • Loading branch information
@AmFlint @neel-astro
AmFlint authored and neel-astro committed Apr 29, 2022
1 parent 63f8f69 commit 9655db1
Show file tree
Hide file tree
Showing 5 changed files with 160 additions and 19 deletions.
15 changes: 12 additions & 3 deletions airflow_versions/airflow_versions.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
7 changes: 5 additions & 2 deletions airflow_versions/airflow_versions_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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)
}
25 changes: 20 additions & 5 deletions cmd/airflow.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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
}
Expand Down
4 changes: 2 additions & 2 deletions cmd/utils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand Down
128 changes: 121 additions & 7 deletions cmd/utils_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ func Test_prepareDefaultAirflowImageTag(t *testing.T) {

output := new(bytes.Buffer)

myTests := []struct {
airflowTests := []struct {
airflowVersion string
expectedImageTag string
expectedError string
Expand All @@ -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 {
Expand Down Expand Up @@ -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)
})
Expand All @@ -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)
})
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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)
}

0 comments on commit 9655db1

Please sign in to comment.