[SDK] Fail when using any integrations that aren't connected yet

sdk/aqueduct/artifacts/preview.py

@@ -37,6 +37,13 @@ def preview_artifacts(
     Returns a list of artifacts, each corresponding to one of the provided `target_artifact_ids`, in
     the same order.
     """
+    global_engine_config: Optional[EngineConfig] = None
+    if globals.__GLOBAL_CONFIG__.engine is not None:
+        global_engine_config = generate_engine_config(
+            globals.__GLOBAL_API_CLIENT__.list_integrations(),
+            globals.__GLOBAL_CONFIG__.engine,
+        )
+
     subgraph = apply_deltas_to_dag(
         dag,
         deltas=[
@@ -50,13 +57,6 @@ def preview_artifacts(
         ],
         make_copy=True,
     )
-
-    global_engine_config: Optional[EngineConfig] = None
-    if globals.__GLOBAL_CONFIG__.engine is not None:
-        global_engine_config = generate_engine_config(
-            globals.__GLOBAL_API_CLIENT__.list_integrations(),
-            globals.__GLOBAL_CONFIG__.engine,
-        )
     subgraph.set_engine_config(global_engine_config=global_engine_config)
 
     engine_statuses = globals.__GLOBAL_API_CLIENT__.get_dynamic_engine_status_by_dag(dag=subgraph)

sdk/aqueduct/error.py

@@ -135,3 +135,13 @@ class ClientValidationError(Error):
 # Exception raised when requirements.txt file is missing
 class RequirementsMissingError(Error):
     pass
+
+
+# Exception raised when the user attempts to use an integration that has failed to connected.
+class IntegrationFailedToConnect(Error):
+    pass
+
+
+# Exception raised when the user attempts to use an integration that is still connecting.
+class IntegrationConnectionInProgress(Error):
+    pass

sdk/aqueduct/integrations/dynamic_k8s_integration.py

@@ -3,6 +3,7 @@
 from aqueduct.constants.enums import K8sClusterActionType, K8sClusterStatusType
 from aqueduct.error import InvalidIntegrationException, InvalidUserArgumentException
 from aqueduct.integrations.connect_config import DynamicK8sConfig
+from aqueduct.integrations.validation import validate_is_connected
 from aqueduct.models.integration import Integration, IntegrationInfo
 from aqueduct.models.response_models import DynamicEngineStatusResponse
 
@@ -41,6 +42,7 @@ class DynamicK8sIntegration(Integration):
     def __init__(self, metadata: IntegrationInfo):
         self._metadata = metadata
 
+    @validate_is_connected()
     def status(self) -> str:
         """Get the current status of the dynamic Kubernetes cluster."""
         engine_statuses = globals.__GLOBAL_API_CLIENT__.get_dynamic_engine_status(
@@ -51,6 +53,7 @@ def status(self) -> str:
 
         return engine_statuses[self._metadata.name].status.value
 
+    @validate_is_connected()
     def create(
         self, config_delta: Union[Dict[str, Union[int, str]], DynamicK8sConfig] = {}
     ) -> None:
@@ -93,6 +96,7 @@ def create(
             config_delta=config_delta,
         )
 
+    @validate_is_connected()
     def update(self, config_delta: Union[Dict[str, Union[int, str]], DynamicK8sConfig]) -> None:
         """Update the dynamic Kubernetes cluster. This can only be done when the cluster is in
             Active status.
@@ -135,6 +139,7 @@ def update(self, config_delta: Union[Dict[str, Union[int, str]], DynamicK8sConfi
             config_delta=config_delta,
         )
 
+    @validate_is_connected()
     def delete(self, force: bool = False) -> None:
         """Deletes the dynamic Kubernetes cluster if it is running, ignoring the keepalive period.
 

sdk/aqueduct/integrations/google_sheets_integration.py

@@ -3,6 +3,7 @@
 from aqueduct.artifacts.base_artifact import BaseArtifact
 from aqueduct.artifacts.table_artifact import TableArtifact
 from aqueduct.constants.enums import ArtifactType, GoogleSheetsSaveMode
+from aqueduct.integrations.validation import validate_is_connected
 from aqueduct.models.artifact import ArtifactMetadata
 from aqueduct.models.dag import DAG
 from aqueduct.models.integration import Integration, IntegrationInfo
@@ -29,6 +30,7 @@ def __init__(self, dag: DAG, metadata: IntegrationInfo):
         self._dag = dag
         self._metadata = metadata
 
+    @validate_is_connected()
     def spreadsheet(
         self,
         spreadsheet_id: str,
@@ -96,6 +98,7 @@ def spreadsheet(
             artifact_id=output_artifact_id,
         )
 
+    @validate_is_connected()
     def save(
         self,
         artifact: BaseArtifact,

sdk/aqueduct/integrations/mongodb_integration.py

@@ -7,6 +7,7 @@
 from aqueduct.constants.enums import ArtifactType, ExecutionMode, LoadUpdateMode
 from aqueduct.integrations.parameters import _validate_parameters
 from aqueduct.integrations.save import _save_artifact
+from aqueduct.integrations.validation import validate_is_connected
 from aqueduct.models.artifact import ArtifactMetadata
 from aqueduct.models.dag import DAG
 from aqueduct.models.integration import Integration, IntegrationInfo
@@ -33,6 +34,7 @@ def __init__(self, dag: DAG, metadata: IntegrationInfo, collection_name: str) ->
         self._dag = dag
         self._collection_name = collection_name
 
+    @validate_is_connected()
     def find(
         self,
         *args: List[Any],
@@ -156,6 +158,7 @@ def find(
             # We are in lazy mode.
             return TableArtifact(self._dag, output_artf_id)
 
+    @validate_is_connected()
     def save(self, artifact: BaseArtifact, update_mode: LoadUpdateMode) -> None:
         """Registers a save operator of the given artifact, to be executed when it's computed in a published flow.
 
@@ -188,6 +191,7 @@ def __init__(self, dag: DAG, metadata: IntegrationInfo):
         self._dag = dag
         self._metadata = metadata
 
+    @validate_is_connected()
     def collection(self, name: str) -> MongoDBCollectionIntegration:
         """Returns a specific collection object to call `.find()` method.
 
@@ -203,6 +207,7 @@ def describe(self) -> None:
         print("==================== MongoDB Integration  =============================")
         self._metadata.describe()
 
+    @validate_is_connected()
     def save(self, artifact: BaseArtifact, collection: str, update_mode: LoadUpdateMode) -> None:
         """Registers a save operator of the given artifact, to be executed when it's computed in a published flow.
 

sdk/aqueduct/integrations/s3_integration.py

@@ -4,6 +4,7 @@
 from aqueduct.artifacts import preview as artifact_utils
 from aqueduct.artifacts.base_artifact import BaseArtifact
 from aqueduct.constants.enums import ArtifactType, ExecutionMode, S3TableFormat
+from aqueduct.integrations.validation import validate_is_connected
 from aqueduct.models.artifact import ArtifactMetadata
 from aqueduct.models.dag import DAG
 from aqueduct.models.integration import Integration, IntegrationInfo
@@ -51,6 +52,7 @@ def __init__(self, dag: DAG, metadata: IntegrationInfo):
         self._dag = dag
         self._metadata = metadata
 
+    @validate_is_connected()
     def file(
         self,
         filepaths: Union[List[str], str],
@@ -171,6 +173,7 @@ def _is_multi_file_search() -> bool:
             # We are in lazy mode.
             return to_artifact_class(self._dag, output_artifact_id, artifact_type)
 
+    @validate_is_connected()
     def save(self, artifact: BaseArtifact, filepath: str, format: Optional[str] = None) -> None:
         """Registers a save operator of the given artifact, to be executed when it's computed in a published flow.
 

sdk/aqueduct/integrations/salesforce_integration.py

@@ -4,6 +4,7 @@
 from aqueduct.artifacts.base_artifact import BaseArtifact
 from aqueduct.artifacts.table_artifact import TableArtifact
 from aqueduct.constants.enums import ArtifactType, SalesforceExtractType
+from aqueduct.integrations.validation import validate_is_connected
 from aqueduct.models.artifact import ArtifactMetadata
 from aqueduct.models.dag import DAG
 from aqueduct.models.integration import Integration, IntegrationInfo
@@ -30,6 +31,7 @@ def __init__(self, dag: DAG, metadata: IntegrationInfo):
         self._dag = dag
         self._metadata = metadata
 
+    @validate_is_connected()
     def search(
         self,
         search_query: str,
@@ -69,6 +71,7 @@ def search(
             artifact_id=output_artifact_id,
         )
 
+    @validate_is_connected()
     def query(
         self,
         query: str,
@@ -102,6 +105,7 @@ def query(
             artifact_id=output_artifact_id,
         )
 
+    @validate_is_connected()
     def save(self, artifact: BaseArtifact, object: str) -> None:
         """Registers a save operator of the given artifact, to be executed when it's computed in a published flow.
 

sdk/aqueduct/integrations/sql_integration.py

@@ -8,6 +8,7 @@
 from aqueduct.error import InvalidUserActionException, InvalidUserArgumentException
 from aqueduct.integrations.parameters import _validate_builtin_expansions, _validate_parameters
 from aqueduct.integrations.save import _save_artifact
+from aqueduct.integrations.validation import validate_is_connected
 from aqueduct.models.artifact import ArtifactMetadata
 from aqueduct.models.dag import DAG
 from aqueduct.models.integration import Integration, IntegrationInfo
@@ -46,6 +47,7 @@ def __init__(self, dag: DAG, metadata: IntegrationInfo):
         self._dag = dag
         self._metadata = metadata
 
+    @validate_is_connected()
     def list_tables(self) -> pd.DataFrame:
         """
         Lists the tables available in the RelationalDB integration.
@@ -80,6 +82,7 @@ def list_tables(self) -> pd.DataFrame:
         sql_artifact = self.sql(query=list_tables_query)
         return sql_artifact.get()
 
+    @validate_is_connected()
     def table(self, name: str) -> pd.DataFrame:
         """
         Retrieves a table from a RealtionalDB integration.
@@ -94,6 +97,7 @@ def table(self, name: str) -> pd.DataFrame:
         sql_artifact = self.sql(query=GET_TABLE_QUERY % name)
         return sql_artifact.get()
 
+    @validate_is_connected()
     def sql(
         self,
         query: Union[str, List[str], RelationalDBExtractParams],
@@ -233,6 +237,7 @@ def sql(
             # We are in lazy mode.
             return TableArtifact(self._dag, sql_output_artifact_id)
 
+    @validate_is_connected()
     def save(self, artifact: BaseArtifact, table_name: str, update_mode: LoadUpdateMode) -> None:
         """Registers a save operator of the given artifact, to be executed when it's computed in a published flow.
 
@@ -269,6 +274,11 @@ def describe(self) -> None:
         print("==================== SQL Integration =============================")
         print("Integration Information:")
         self._metadata.describe()
-        print("Integration Table List Preview:")
-        print(self.list_tables()["name"].head().to_string())
-        print("(only first 5 tables are shown)")
+
+        # Only list the tables if the integration is connected.
+        try:
+            print("Integration Table List Preview:")
+            print(self.list_tables()["name"].head().to_string())
+            print("(only first 5 tables are shown)")
+        except:
+            pass

sdk/aqueduct/integrations/validation.py

@@ -0,0 +1,19 @@
+from typing import Any, Callable
+
+from aqueduct.utils.integration_validation import validate_integration_is_connected
+
+AnyFunc = Callable[..., Any]
+
+
+def validate_is_connected() -> Callable[[AnyFunc], AnyFunc]:
+    """This decorator, when used on an Integration class method, ensures that the integration is
+    connected before allowing the method to be called."""
+
+    def decorator(method: AnyFunc) -> Callable[[AnyFunc], AnyFunc]:
+        def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            validate_integration_is_connected(self.name(), self._metadata.exec_state)
+            return method(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

sdk/aqueduct/models/integration.py

@@ -4,6 +4,7 @@
 from typing import Any, Dict, Optional
 
 from aqueduct.constants.enums import ExecutionStatus, ServiceType
+from aqueduct.error import IntegrationConnectionInProgress, IntegrationFailedToConnect
 from aqueduct.models.execution_state import ExecutionState
 from aqueduct.models.utils import human_readable_timestamp
 from pydantic import BaseModel

sdk/aqueduct/utils/integration_validation.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+from aqueduct.constants.enums import ExecutionStatus
+from aqueduct.error import IntegrationConnectionInProgress, IntegrationFailedToConnect
+from aqueduct.models.execution_state import ExecutionState
+
+
+def validate_integration_is_connected(name: str, exec_state: Optional[ExecutionState]) -> None:
+    """Method used to determine if this integration was successfully connected to or not.
+    If not successfully connected (or pending), we will raise an Exception.
+    """
+    # TODO(ENG-2813): Remove the assumption that a missing `exec_state` means success.
+    if exec_state is None or exec_state.status == ExecutionStatus.SUCCEEDED:
+        return
+
+    if exec_state.status == ExecutionStatus.FAILED:
+        assert exec_state.error is not None
+        raise IntegrationFailedToConnect(
+            "Cannot use integration %s because it has not been successfully connected to: "
+            "%s\n%s\n\n Please see the /integrations page on the UI for more details."
+            % (
+                name,
+                exec_state.error.tip,
+                exec_state.error.context,
+            )
+        )
+    else:
+        # The assumption is that we are in the running state here.
+        raise IntegrationConnectionInProgress(
+            "Cannot use integration %s because it is still in the process of connecting."
+            "Please see the /integrations page on the UI for more details." % name
+        )

sdk/aqueduct/utils/utils.py

@@ -14,6 +14,7 @@
 from aqueduct.models.dag import Schedule
 from aqueduct.models.integration import IntegrationInfo
 from aqueduct.models.operators import ParamSpec
+from aqueduct.utils.integration_validation import validate_integration_is_connected
 from croniter import croniter
 
 from ..models.execution_state import Logs
@@ -160,6 +161,8 @@ def generate_engine_config(
         )
 
     integration = integrations[integration_name]
+    validate_integration_is_connected(integration_name, integration.exec_state)
+
     if integration.service == ServiceType.AIRFLOW:
         return EngineConfig(
             type=RuntimeType.AIRFLOW,

src/golang/lib/lambda/utils.go

@@ -78,7 +78,11 @@ func CreateLambdaFunction(ctx context.Context, functionsToShip []LambdaFunctionT
 	defer close(pushImageChannel)
 	AddFunctionTypeToChannel(functionsToShip[:], pullImageChannel)
 
-	for i := 0; i < MaxConcurrentDownload; i++ {
+	numPullWorkers := MaxConcurrentDownload
+	if numPullWorkers > len(functionsToShip) {
+		numPullWorkers = len(functionsToShip)
+	}
+	for i := 0; i < numPullWorkers; i++ {
 		errGroup.Go(func() error {
 			for {
 				select {
@@ -105,7 +109,11 @@ func CreateLambdaFunction(ctx context.Context, functionsToShip []LambdaFunctionT
 	AddFunctionTypeToChannel(functionsToShip[:], incompleteWorkChannel)
 
 	// Receive the downloaded docker images from push channels and create lambda functions on a concurrency of "MaxConcurrentUpload".
-	for i := 0; i < MaxConcurrentUpload; i++ {
+	numPushWorkers := MaxConcurrentUpload
+	if numPushWorkers > len(functionsToShip) {
+		numPushWorkers = len(functionsToShip)
+	}
+	for i := 0; i < numPushWorkers; i++ {
 		errGroup.Go(func() error {
 			for {
 				select {
@@ -245,6 +253,8 @@ func PushImageToPrivateECR(functionType LambdaFunctionType, roleArn string) erro
 	repositoryUri := fmt.Sprintf("%s:%s", *result.Repository.RepositoryUri, lib.ServerVersionNumber)
 
 	cmd := exec.Command("docker", "tag", versionedLambdaImageUri, repositoryUri)
+	cmd.Stdout = stdout
+	cmd.Stderr = stderr
 	err = cmd.Run()
 	if err != nil {
 		log.Info(stdout.String())
@@ -253,6 +263,8 @@ func PushImageToPrivateECR(functionType LambdaFunctionType, roleArn string) erro
 	}
 
 	cmd = exec.Command("docker", "push", repositoryUri)
+	cmd.Stdout = stdout
+	cmd.Stderr = stderr
 	err = cmd.Run()
 	if err != nil {
 		log.Info(stdout.String())