feat: User Defined Stages in Spark CI/CD Pipeline

examples/spark-data-lake/README.md

@@ -70,6 +70,35 @@ pip install -r requirements.txt
 }
 ```
 
+Alternatively, if further customization is necessary, the following allows multiple different environments to be created:
+
+**Stage names must be unique**
+
+```json
+{
+  "environments": [
+    {
+      "stageName": "<STAGE_NAME_1>",
+      "account": "<STAGE_ACCOUNT_ID>",
+      "region": "<REGION>",
+      "triggerIntegTest": "<OPTIONAL_BOOLEAN_CAN_BE_OMMITTED>"
+    },
+    {
+      "stageName": "<STAGE_NAME_2>",
+      "account": "<STAGE_ACCOUNT_ID>",
+      "region": "<REGION>",
+      "triggerIntegTest": "<OPTIONAL_BOOLEAN_CAN_BE_OMMITTED>"
+    },
+    {
+      "stageName": "<STAGE_NAME_3>",
+      "account": "<STAGE_ACCOUNT_ID>",
+      "region": "<REGION>",
+      "triggerIntegTest": "<OPTIONAL_BOOLEAN_CAN_BE_OMMITTED>"
+    }
+  ]
+}
+```
+
 5. Create a connection, this will server to link your code repository to Amazon Code Pipeline. You can follow the instruction in the [AWS documentation](https://docs.aws.amazon.com/dtconsole/latest/userguide/connections.html)
 to create a connection.
 

framework/src/processing/README.md

@@ -206,6 +206,37 @@ You need to also provide the accounts information in the cdk.json in the form of
 }
 ```
 
+## User Defined Stages
+
+To define multiple stages (which can also be deployed in different AWS accounts by following the bootstrap command in the previous section), configure the `cdk.json` file with the following:
+
+**Stage names must be unique**
+
+```json
+{
+  "environments": [
+    {
+      "stageName": "<STAGE_NAME_1>",
+      "account": "<STAGE_ACCOUNT_ID>",
+      "region": "<REGION>",
+      "triggerIntegTest": "<OPTIONAL_BOOLEAN_CAN_BE_OMMITTED>"
+    },
+    {
+      "stageName": "<STAGE_NAME_2>",
+      "account": "<STAGE_ACCOUNT_ID>",
+      "region": "<REGION>",
+      "triggerIntegTest": "<OPTIONAL_BOOLEAN_CAN_BE_OMMITTED>"
+    },
+    {
+      "stageName": "<STAGE_NAME_3>",
+      "account": "<STAGE_ACCOUNT_ID>",
+      "region": "<REGION>",
+      "triggerIntegTest": "<OPTIONAL_BOOLEAN_CAN_BE_OMMITTED>"
+    }
+  ]
+}
+```
+
 ## Defining a CDK Stack for the Spark application
 
 The `SparkCICDPipeline` construct deploys an application stack, which contains your business logic, into staging and production environments.

framework/src/processing/lib/cicd-pipeline/spark-emr-cicd-pipeline.ts

@@ -18,6 +18,18 @@ import {
 } from '../../../utils';
 import { DEFAULT_SPARK_IMAGE, SparkImage } from '../emr-releases';
 
+const MISSING_ENVIRONMENTS_ERROR = 'MissingEnvironmentsError';
+const DUPLICATE_STAGE_NAME_ERROR = 'DuplicateStageNameError';
+
+/**
+ * User defined CI/CD environment stages
+ */
+interface CICDEnvironment {
+  stageName: string;
+  account: string;
+  region: string;
+  triggerIntegTest?: boolean;
+}
 
 /**
  * A CICD Pipeline to test and deploy a Spark application on Amazon EMR in cross-account environments using CDK Pipelines.
@@ -211,36 +223,63 @@ export class SparkEmrCICDPipeline extends TrackedConstruct {
       },
     });
 
-    // Create the Staging stage of the CICD
-    const staging = new ApplicationStage(this, 'Staging', {
-      env: this.getAccountFromContext('staging'),
+    try {
+      const environments = this.getUserDefinedEnvironmentsFromContext();
+
+      for (const e of environments) {
+        this.integrationTestStage = this.attachStageToPipeline(e.stageName.toUpperCase(), {
+          account: e.account,
+          region: e.region,
+        }, e.triggerIntegTest || false, buildStage, props);
+      }
+    } catch (e) {
+      const error = e as Error;
+      if (error.name === DUPLICATE_STAGE_NAME_ERROR) {
+        throw e;
+      }
+
+      this.integrationTestStage = this.attachStageToPipeline('Staging', this.getAccountFromContext('staging'), true, buildStage, props);
+      this.attachStageToPipeline('Prod', this.getAccountFromContext('prod'), false, buildStage, props);
+    }
+  }
+
+  /**
+   * Attaches the given stage to the pipeline
+   * @param stageName
+   * @param resourceEnvironment
+   * @param attachIntegTest
+   * @param buildStage
+   * @param props
+   * @returns {CodeBuildStep|undefined} if integration step is configured, this returns the corresponding `CodeBuildStep` for the test
+   */
+  private attachStageToPipeline(stageName: string, resourceEnvironment: ResourceEnvironment
+    , attachIntegTest: boolean, buildStage: CodeBuildStep
+    , props: SparkEmrCICDPipelineProps): CodeBuildStep|undefined {
+    const applicationStage = new ApplicationStage(this, stageName, {
+      env: resourceEnvironment,
       applicationStackFactory: props.applicationStackFactory,
-      outputsEnv: props.integTestEnv,
-      stage: CICDStage.STAGING,
+      outputsEnv: (attachIntegTest && props.integTestScript) ? props.integTestEnv : undefined,
+      stage: CICDStage.of(stageName.toUpperCase()),
     });
-    const stagingDeployment = this.pipeline.addStage(staging);
+    const stageDeployment = this.pipeline.addStage(applicationStage);
+
+    let integrationTestStage:CodeBuildStep|undefined = undefined;
 
-    if (props.integTestScript) {
+    if (attachIntegTest && props.integTestScript) {
       // Extract the path and script name from the integration tests script path
       const [integPath, integScript] = SparkEmrCICDPipeline.extractPath(props.integTestScript);
 
-      this.integrationTestStage = new CodeBuildStep('IntegrationTests', {
+      integrationTestStage = new CodeBuildStep(`${stageName}IntegrationTests`, {
         input: buildStage.addOutputDirectory(integPath),
         commands: [`chmod +x ${integScript} && ./${integScript}`],
-        envFromCfnOutputs: staging.stackOutputsEnv,
+        envFromCfnOutputs: applicationStage.stackOutputsEnv,
         rolePolicyStatements: props.integTestPermissions,
       });
       // Add a post step to run the integration tests
-      stagingDeployment.addPost(this.integrationTestStage);
+      stageDeployment.addPost(integrationTestStage);
     }
 
-    // Create the Production stage of the CICD
-    this.pipeline.addStage(new ApplicationStage(this, 'Production', {
-      env: this.getAccountFromContext('prod'),
-      applicationStackFactory: props.applicationStackFactory,
-      stage: CICDStage.PROD,
-    }));
-
+    return integrationTestStage;
   }
 
   /**
@@ -251,4 +290,34 @@ export class SparkEmrCICDPipeline extends TrackedConstruct {
     if (!account) throw new Error(`Missing context variable ${name}`);
     return account;
   }
+
+  /**
+   * Retrieves the list of user defined environments from the context
+   * @returns {CICDEnvironment[]} list of user defined environments
+   */
+  private getUserDefinedEnvironmentsFromContext(): CICDEnvironment[] {
+    const environments = this.node.tryGetContext('environments') as CICDEnvironment[];
+
+    if (!environments) {
+      const missingContextError = new Error('Missing context variable environments');
+      missingContextError.name = MISSING_ENVIRONMENTS_ERROR;
+      throw missingContextError;
+    } else {
+      //check for duplicates
+
+      const stageNameTracker = [];
+
+      for (let e of environments) {
+        if (stageNameTracker.indexOf(e.stageName) != -1) {
+          const duplicateStageError = new Error('Duplicate stage name found');
+          duplicateStageError.name = DUPLICATE_STAGE_NAME_ERROR;
+          throw duplicateStageError;
+        }
+
+        stageNameTracker.push(e.stageName);
+      }
+    }
+
+    return environments;
+  }
 }

framework/src/utils/lib/application-stage.ts

@@ -10,9 +10,31 @@ import { ApplicationStackFactory } from './application-stack-factory';
 /**
  * The list of CICD Stages used in CICD Pipelines.
  */
-export enum CICDStage {
-  STAGING = 'staging',
-  PROD = 'prod',
+export class CICDStage {
+
+  /**
+   * Prod stage
+   */
+  public static readonly PROD = CICDStage.of('PROD');
+
+  /**
+   * Staging stage
+   */
+  public static readonly STAGING = CICDStage.of('STAGING');
+
+  /**
+   * Custom stage
+   * @param stage the stage inside the pipeline
+   * @returns
+   */
+  public static of(stage: string) {
+    return new CICDStage(stage);
+  }
+
+  /**
+   * @param stage the stage inside the pipeline
+   */
+  private constructor(public readonly stage: string) {}
 }
 
 /**

framework/test/e2e/spark-cicd-pipeline.e2e.test.ts

@@ -4,7 +4,7 @@
 /**
  * E2E test for SparkCICDPipeline
  *
- * @group e2e/processing/spark-cicd
+ * @group e2e/processing/default-spark-cicd
  */
 
 import { RemovalPolicy, CfnOutput, Stack, StackProps, App } from 'aws-cdk-lib';
@@ -18,7 +18,7 @@ jest.setTimeout(9000000);
 
 // GIVEN
 const app = new App();
-const cicdStack = new Stack(app, 'CICDStack', {
+const cicdStack = new Stack(app, 'DefaultCICDStack', {
   env: {
     region: 'eu-west-1',
   },
@@ -27,8 +27,8 @@ const testStack = new TestStack('SparkCICDPipelineTestStack', app, cicdStack);
 const { stack } = testStack;
 
 stack.node.setContext('@data-solutions-framework-on-aws/removeDataOnDestroy', true);
-stack.node.setContext('staging', { accountId: '123456789012', region: 'eu-west-1' });
-stack.node.setContext('prod', { accountId: '123456789012', region: 'eu-west-1' });
+stack.node.setContext('staging', { accountId: stack.account, region: stack.region });
+stack.node.setContext('prod', { accountId: stack.account, region: stack.region });
 
 interface MyApplicationStackProps extends StackProps {
   readonly prodBoolean: Boolean;