From fe503bca21142ba7ad3e220b65e62dd48f910579 Mon Sep 17 00:00:00 2001 From: Jeremy Beard Date: Mon, 30 Oct 2023 12:52:27 -0400 Subject: [PATCH] Only use flags for retrieving env connections (#1423) --- cloud/environment/environment.go | 27 +++++---- cloud/environment/environment_test.go | 39 +++--------- cmd/airflow.go | 16 +++-- cmd/airflow_test.go | 87 +++++---------------------- cmd/cloud/setup.go | 12 +++- cmd/cloud/setup_test.go | 36 +++++++++++ 6 files changed, 95 insertions(+), 122 deletions(-) diff --git a/cloud/environment/environment.go b/cloud/environment/environment.go index 6848c8ca5..766630592 100644 --- a/cloud/environment/environment.go +++ b/cloud/environment/environment.go @@ -2,25 +2,31 @@ package environment import ( http_context "context" + "errors" astrocore "github.com/astronomer/astro-cli/astro-client-core" "github.com/astronomer/astro-cli/config" ) -func ListConnections(workspaceID, deploymentID string, coreClient astrocore.CoreClient) map[string]astrocore.EnvironmentObjectConnection { - envObjs := listEnvironmentObjects(workspaceID, deploymentID, astrocore.ListEnvironmentObjectsParamsObjectTypeCONNECTION, coreClient) +var ErrorEntityIDNotSpecified = errors.New("workspace or deployment ID must be specified") + +func ListConnections(workspaceID, deploymentID string, coreClient astrocore.CoreClient) (map[string]astrocore.EnvironmentObjectConnection, error) { + envObjs, err := listEnvironmentObjects(workspaceID, deploymentID, astrocore.ListEnvironmentObjectsParamsObjectTypeCONNECTION, coreClient) + if err != nil { + return nil, err + } connections := make(map[string]astrocore.EnvironmentObjectConnection) for _, envObj := range envObjs { connections[envObj.ObjectKey] = *envObj.Connection } - return connections + return connections, nil } -func listEnvironmentObjects(workspaceID, deploymentID string, objectType astrocore.ListEnvironmentObjectsParamsObjectType, coreClient astrocore.CoreClient) []astrocore.EnvironmentObject { +func listEnvironmentObjects(workspaceID, deploymentID string, objectType astrocore.ListEnvironmentObjectsParamsObjectType, coreClient astrocore.CoreClient) ([]astrocore.EnvironmentObject, error) { c, err := config.GetCurrentContext() if err != nil { - return []astrocore.EnvironmentObject{} + return nil, err } showSecrets := true resolvedLinked := true @@ -40,24 +46,21 @@ func listEnvironmentObjects(workspaceID, deploymentID string, objectType astroco case workspaceID != "": // or, if the workspace is specified during the command, use that listParams.WorkspaceId = &workspaceID - case c.Workspace != "": - // or, if the workspace is specified as part of the context, use that - listParams.WorkspaceId = &c.Workspace default: // otherwise, we don't have an entity to list for, so we return an empty list - return []astrocore.EnvironmentObject{} + return nil, ErrorEntityIDNotSpecified } resp, err := coreClient.ListEnvironmentObjectsWithResponse(http_context.Background(), c.Organization, listParams) if err != nil { - return []astrocore.EnvironmentObject{} + return nil, err } err = astrocore.NormalizeAPIError(resp.HTTPResponse, resp.Body) if err != nil { - return []astrocore.EnvironmentObject{} + return nil, err } envObjsPaginated := *resp.JSON200 envObjs := envObjsPaginated.EnvironmentObjects - return envObjs + return envObjs, nil } diff --git a/cloud/environment/environment_test.go b/cloud/environment/environment_test.go index a77e0d53c..38b35decb 100644 --- a/cloud/environment/environment_test.go +++ b/cloud/environment/environment_test.go @@ -41,7 +41,8 @@ func TestListConnections(t *testing.T) { }}, }, nil).Once() - conns := ListConnections("", deploymentID, mockClient) + conns, err := ListConnections("", deploymentID, mockClient) + assert.NoError(t, err) assert.Len(t, conns, 1) assert.Equal(t, "postgres", conns["conn1"].Type) @@ -66,32 +67,8 @@ func TestListConnections(t *testing.T) { }}, }, nil).Once() - conns := ListConnections(workspaceID, "", mockClient) - assert.Len(t, conns, 1) - assert.Equal(t, "postgres", conns["conn1"].Type) - - mockClient.AssertExpectations(t) - }) - - t.Run("List connections with context workspace ID", func(t *testing.T) { - workspaceID := context.Workspace - listParams := &astrocore.ListEnvironmentObjectsParams{ - WorkspaceId: &workspaceID, - ObjectType: &objectType, - ShowSecrets: &showSecrets, - ResolveLinked: &resolvedLinked, - Limit: &limit, - } - - mockClient := new(astrocore_mocks.ClientWithResponsesInterface) - mockClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, organization, listParams).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{StatusCode: 200}, - JSON200: &astrocore.EnvironmentObjectsPaginated{EnvironmentObjects: []astrocore.EnvironmentObject{ - {ObjectKey: "conn1", Connection: &astrocore.EnvironmentObjectConnection{Type: "postgres"}}, - }}, - }, nil).Once() - - conns := ListConnections("", "", mockClient) + conns, err := ListConnections(workspaceID, "", mockClient) + assert.NoError(t, err) assert.Len(t, conns, 1) assert.Equal(t, "postgres", conns["conn1"].Type) @@ -106,8 +83,8 @@ func TestListConnections(t *testing.T) { mockClient := new(astrocore_mocks.ClientWithResponsesInterface) - conns := ListConnections("", "", mockClient) - assert.Len(t, conns, 0) + _, err := ListConnections("", "", mockClient) + assert.Error(t, err) mockClient.AssertExpectations(t) }) @@ -124,8 +101,8 @@ func TestListConnections(t *testing.T) { mockClient := new(astrocore_mocks.ClientWithResponsesInterface) mockClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, organization, listParams).Return(nil, assert.AnError).Once() - conns := ListConnections("", "", mockClient) - assert.Len(t, conns, 0) + _, err := ListConnections(context.Workspace, "", mockClient) + assert.Error(t, err) mockClient.AssertExpectations(t) }) diff --git a/cmd/airflow.go b/cmd/airflow.go index 8eff27284..f487bde09 100644 --- a/cmd/airflow.go +++ b/cmd/airflow.go @@ -639,8 +639,12 @@ func airflowStart(cmd *cobra.Command, args []string, astroCoreClient astrocore.C } var envConns map[string]astrocore.EnvironmentObjectConnection - if !config.CFG.DisableEnvObjects.GetBool() { - envConns = environment.ListConnections(workspaceID, deploymentID, astroCoreClient) + if !config.CFG.DisableEnvObjects.GetBool() && (workspaceID != "" || deploymentID != "") { + var err error + envConns, err = environment.ListConnections(workspaceID, deploymentID, astroCoreClient) + if err != nil { + return err + } } containerHandler, err := containerHandlerInit(config.WorkingPath, envFile, dockerfile, "") @@ -759,8 +763,12 @@ func airflowRestart(cmd *cobra.Command, args []string, astroCoreClient astrocore noBrowser = true var envConns map[string]astrocore.EnvironmentObjectConnection - if !config.CFG.DisableEnvObjects.GetBool() { - envConns = environment.ListConnections(workspaceID, deploymentID, astroCoreClient) + if !config.CFG.DisableEnvObjects.GetBool() && (workspaceID != "" || deploymentID != "") { + var err error + envConns, err = environment.ListConnections(workspaceID, deploymentID, astroCoreClient) + if err != nil { + return err + } } return containerHandler.Start(customImageName, settingsFile, composeFile, noCache, noBrowser, waitTime, envConns) diff --git a/cmd/airflow_test.go b/cmd/airflow_test.go index 2dd9930f7..469658f47 100644 --- a/cmd/airflow_test.go +++ b/cmd/airflow_test.go @@ -410,36 +410,21 @@ func TestAirflowStart(t *testing.T) { cmd := newAirflowStartCmd(nil) args := []string{"test-env-file"} - envObj := astrocore.EnvironmentObject{ - ObjectKey: "test-object-key", - Connection: &astrocore.EnvironmentObjectConnection{ - Type: "test-conn-type", - }, - } - mockCoreClient := new(coreMocks.ClientWithResponsesInterface) - mockCoreClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{ - StatusCode: 200, - }, - JSON200: &astrocore.EnvironmentObjectsPaginated{ - EnvironmentObjects: []astrocore.EnvironmentObject{envObj}, - }, - }, nil).Once() - mockContainerHandler := new(mocks.ContainerHandler) containerHandlerInit = func(airflowHome, envFile, dockerfile, imageName string) (airflow.ContainerHandler, error) { - mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", false, false, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection{envObj.ObjectKey: *envObj.Connection}).Return(nil).Once() + mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", false, false, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection(nil)).Return(nil).Once() return mockContainerHandler, nil } - err := airflowStart(cmd, args, mockCoreClient) + err := airflowStart(cmd, args, nil) assert.NoError(t, err) mockContainerHandler.AssertExpectations(t) - mockCoreClient.AssertExpectations(t) }) - t.Run("success with environment objects disabled", func(t *testing.T) { + t.Run("success with deployment id flag set but environment objects disabled", func(t *testing.T) { cmd := newAirflowStartCmd(nil) + deploymentID = "test-deployment-id" + cmd.Flag("deployment-id").Value.Set(deploymentID) args := []string{"test-env-file"} config.CFG.DisableEnvObjects.SetHomeString("true") defer config.CFG.DisableEnvObjects.SetHomeString("false") @@ -534,19 +519,13 @@ func TestAirflowStart(t *testing.T) { cmd := newAirflowStartCmd(nil) args := []string{"test-env-file"} - mockCoreClient := new(coreMocks.ClientWithResponsesInterface) - mockCoreClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{StatusCode: 200}, - JSON200: &astrocore.EnvironmentObjectsPaginated{EnvironmentObjects: []astrocore.EnvironmentObject{}}, - }, nil).Once() - mockContainerHandler := new(mocks.ContainerHandler) containerHandlerInit = func(airflowHome, envFile, dockerfile, imageName string) (airflow.ContainerHandler, error) { - mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", false, false, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection{}).Return(errMock).Once() + mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", false, false, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection(nil)).Return(errMock).Once() return mockContainerHandler, nil } - err := airflowStart(cmd, args, mockCoreClient) + err := airflowStart(cmd, args, nil) assert.ErrorIs(t, err, errMock) mockContainerHandler.AssertExpectations(t) }) @@ -555,17 +534,11 @@ func TestAirflowStart(t *testing.T) { cmd := newAirflowStartCmd(nil) args := []string{} - mockCoreClient := new(coreMocks.ClientWithResponsesInterface) - mockCoreClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{StatusCode: 200}, - JSON200: &astrocore.EnvironmentObjectsPaginated{EnvironmentObjects: []astrocore.EnvironmentObject{}}, - }, nil).Once() - containerHandlerInit = func(airflowHome, envFile, dockerfile, imageName string) (airflow.ContainerHandler, error) { return nil, errMock } - err := airflowStart(cmd, args, mockCoreClient) + err := airflowStart(cmd, args, nil) assert.ErrorIs(t, err, errMock) }) } @@ -924,28 +897,14 @@ func TestAirflowRestart(t *testing.T) { cmd.Flag("no-cache").Value.Set("true") args := []string{"test-env-file"} - envObj := astrocore.EnvironmentObject{ - ObjectKey: "test-object-key", - Connection: &astrocore.EnvironmentObjectConnection{Type: "test-conn-type"}, - } - mockCoreClient := new(coreMocks.ClientWithResponsesInterface) - mockCoreClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{ - StatusCode: 200, - }, - JSON200: &astrocore.EnvironmentObjectsPaginated{ - EnvironmentObjects: []astrocore.EnvironmentObject{envObj}, - }, - }, nil).Once() - mockContainerHandler := new(mocks.ContainerHandler) containerHandlerInit = func(airflowHome, envFile, dockerfile, imageName string) (airflow.ContainerHandler, error) { mockContainerHandler.On("Stop", true).Return(nil).Once() - mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", true, true, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection{envObj.ObjectKey: *envObj.Connection}).Return(nil).Once() + mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", true, true, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection(nil)).Return(nil).Once() return mockContainerHandler, nil } - err := airflowRestart(cmd, args, mockCoreClient) + err := airflowRestart(cmd, args, nil) assert.NoError(t, err) mockContainerHandler.AssertExpectations(t) }) @@ -1025,19 +984,13 @@ func TestAirflowRestart(t *testing.T) { cmd.Flag("no-cache").Value.Set("true") args := []string{"test-env-file"} - mockCoreClient := new(coreMocks.ClientWithResponsesInterface) - mockCoreClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{StatusCode: 200}, - JSON200: &astrocore.EnvironmentObjectsPaginated{EnvironmentObjects: []astrocore.EnvironmentObject{}}, - }, nil).Once() - mockContainerHandler := new(mocks.ContainerHandler) containerHandlerInit = func(airflowHome, envFile, dockerfile, imageName string) (airflow.ContainerHandler, error) { mockContainerHandler.On("Stop", true).Return(errMock).Once() return mockContainerHandler, nil } - err := airflowRestart(cmd, args, mockCoreClient) + err := airflowRestart(cmd, args, nil) assert.ErrorIs(t, err, errMock) mockContainerHandler.AssertExpectations(t) }) @@ -1047,20 +1000,14 @@ func TestAirflowRestart(t *testing.T) { cmd.Flag("no-cache").Value.Set("true") args := []string{"test-env-file"} - mockCoreClient := new(coreMocks.ClientWithResponsesInterface) - mockCoreClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{StatusCode: 200}, - JSON200: &astrocore.EnvironmentObjectsPaginated{EnvironmentObjects: []astrocore.EnvironmentObject{}}, - }, nil).Once() - mockContainerHandler := new(mocks.ContainerHandler) containerHandlerInit = func(airflowHome, envFile, dockerfile, imageName string) (airflow.ContainerHandler, error) { mockContainerHandler.On("Stop", true).Return(nil).Once() - mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", true, true, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection{}).Return(errMock).Once() + mockContainerHandler.On("Start", "", "airflow_settings.yaml", "", true, true, 1*time.Minute, map[string]astrocore.EnvironmentObjectConnection(nil)).Return(errMock).Once() return mockContainerHandler, nil } - err := airflowRestart(cmd, args, mockCoreClient) + err := airflowRestart(cmd, args, nil) assert.ErrorIs(t, err, errMock) mockContainerHandler.AssertExpectations(t) }) @@ -1070,17 +1017,11 @@ func TestAirflowRestart(t *testing.T) { cmd.Flag("no-cache").Value.Set("true") args := []string{"test-env-file"} - mockCoreClient := new(coreMocks.ClientWithResponsesInterface) - mockCoreClient.On("ListEnvironmentObjectsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&astrocore.ListEnvironmentObjectsResponse{ - HTTPResponse: &http.Response{StatusCode: 200}, - JSON200: &astrocore.EnvironmentObjectsPaginated{EnvironmentObjects: []astrocore.EnvironmentObject{}}, - }, nil).Once() - containerHandlerInit = func(airflowHome, envFile, dockerfile, imageName string) (airflow.ContainerHandler, error) { return nil, errMock } - err := airflowRestart(cmd, args, mockCoreClient) + err := airflowRestart(cmd, args, nil) assert.ErrorIs(t, err, errMock) }) } diff --git a/cmd/cloud/setup.go b/cmd/cloud/setup.go index 8806e6ee9..31d1c5ba0 100644 --- a/cmd/cloud/setup.go +++ b/cmd/cloud/setup.go @@ -69,10 +69,12 @@ func Setup(cmd *cobra.Command, client astro.Client, coreClient astrocore.CoreCli return nil } - // If the user is using dev commands no need to go through auth setup. - if cmd.CalledAs() == "dev" && cmd.Parent().Use == topLvlCmd { + // If the user is using dev commands no need to go through auth setup, + // unless the workspace or deployment ID flag is set. + if cmd.CalledAs() == "dev" && cmd.Parent().Use == topLvlCmd && !workspaceOrDeploymentIDFlagSet(cmd) { return nil } + // If the user is using flow commands no need to go through auth setup. if cmd.CalledAs() == "flow" && cmd.Parent().Use == topLvlCmd { return nil @@ -442,3 +444,9 @@ func checkAPIToken(isDeploymentFile bool, coreClient astrocore.CoreClient) (bool } return true, nil } + +func workspaceOrDeploymentIDFlagSet(cmd *cobra.Command) bool { + wsID, _ := cmd.Flags().GetString("workspace-id") + depID, _ := cmd.Flags().GetString("deployment-id") + return wsID != "" || depID != "" +} diff --git a/cmd/cloud/setup_test.go b/cmd/cloud/setup_test.go index 7d7c5c22d..f1e880296 100644 --- a/cmd/cloud/setup_test.go +++ b/cmd/cloud/setup_test.go @@ -50,6 +50,42 @@ func TestSetup(t *testing.T) { assert.NoError(t, err) }) + t.Run("dev cmd with workspace flag set", func(t *testing.T) { + testUtil.SetupOSArgsForGinkgo() + cmd := &cobra.Command{Use: "dev"} + cmd.Flags().StringVarP(&workspaceID, "workspace-id", "w", "test-workspace-id", "") + cmd, err := cmd.ExecuteC() + assert.NoError(t, err) + + rootCmd := &cobra.Command{Use: "astro"} + rootCmd.AddCommand(cmd) + + authLogin = func(domain, token string, client astro.Client, coreClient astrocore.CoreClient, out io.Writer, shouldDisplayLoginLink bool) error { + return nil + } + + err = Setup(cmd, nil, nil) + assert.NoError(t, err) + }) + + t.Run("dev cmd with deployment flag set", func(t *testing.T) { + testUtil.SetupOSArgsForGinkgo() + cmd := &cobra.Command{Use: "dev"} + cmd.Flags().StringVarP(&workspaceID, "deployment-id", "w", "test-deployment-id", "") + cmd, err := cmd.ExecuteC() + assert.NoError(t, err) + + rootCmd := &cobra.Command{Use: "astro"} + rootCmd.AddCommand(cmd) + + authLogin = func(domain, token string, client astro.Client, coreClient astrocore.CoreClient, out io.Writer, shouldDisplayLoginLink bool) error { + return nil + } + + err = Setup(cmd, nil, nil) + assert.NoError(t, err) + }) + t.Run("flow cmd", func(t *testing.T) { testUtil.SetupOSArgsForGinkgo() cmd := &cobra.Command{Use: "flow"}