Only use flags for retrieving env connections (#1423)

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
 }

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)
 	})

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)

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)
 	})
 }

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 != ""
+}