This solution helps you deploy ETL processes and data storage resources to create an Insurance Lake using Amazon S3 buckets for storage, AWS Glue for data transformation, and AWS CDK Pipelines. It is originally based on the AWS blog Deploy data lake ETL jobs using CDK Pipelines, and complements the InsuranceLake Infrastructure project
# !/usr/bin/env python3
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import os
import aws_cdk as cdk
from cdk_nag import AwsSolutionsChecks, NagSuppressions

from lib.pipeline_stack import PipelineStack
from lib.code_commit_stack import CodeCommitStack
from lib.configuration import (
    ACCOUNT_ID, CODECOMMIT_MIRROR_REPOSITORY_NAME, DEPLOYMENT, DEV, TEST, PROD, REGION, CODE_BRANCH,
    get_logical_id_prefix, get_all_configurations
)
from lib.tagging import tag

app = cdk.App()

# Enable CDK Nag for the Mirror repository, Pipeline, and related stacks
# Environment stacks must be enabled on the Stage resource
cdk.Aspects.of(app).add(AwsSolutionsChecks())

raw_mappings = get_all_configurations()

deployment_account = raw_mappings[DEPLOYMENT][ACCOUNT_ID]
deployment_region = raw_mappings[DEPLOYMENT][REGION]
deployment_aws_env = {
    'account': deployment_account,
    'region': deployment_region,
}
logical_id_prefix = get_logical_id_prefix()

if raw_mappings[DEPLOYMENT][CODECOMMIT_MIRROR_REPOSITORY_NAME] != '':
    mirror_repository_stack = CodeCommitStack(
        app,
        f'{DEPLOYMENT}-{logical_id_prefix}EtlMirrorRepository',
        description='InsuranceLake stack for ETL repository mirror (uksb-1tu7mtee2)',
        target_environment=DEPLOYMENT,
        env=deployment_aws_env,
    )
    tag(mirror_repository_stack, DEPLOYMENT)

if os.environ.get('ENV', DEV) == DEV:
    target_environment = DEV
    dev_account = raw_mappings[DEV][ACCOUNT_ID]
    dev_region = raw_mappings[DEV][REGION]
    dev_aws_env = {
        'account': dev_account,
        'region': dev_region,
    }
    dev_pipeline_stack = PipelineStack(
        app,
        f'{target_environment}-{logical_id_prefix}EtlPipeline',
        description=f'InsuranceLake stack for ETL pipeline - {DEV} environment (uksb-1tu7mtee2)',
        target_environment=DEV,
        target_branch=raw_mappings[DEV][CODE_BRANCH],
        target_aws_env=dev_aws_env,
        env=deployment_aws_env,
    )
    tag(dev_pipeline_stack, DEPLOYMENT)

# TODO: Check for cross-region or cross-account deployment and explicitly create a replication bucket with naming convention, encryption, access logs
if os.environ.get('ENV', TEST) == TEST:
    target_environment = TEST
    test_account = raw_mappings[TEST][ACCOUNT_ID]
    test_region = raw_mappings[TEST][REGION]
    test_aws_env = {
        'account': test_account,
        'region': test_region,
    }
    test_pipeline_stack = PipelineStack(
        app,
        f'{target_environment}-{logical_id_prefix}EtlPipeline',
        description=f'InsuranceLake stack for ETL pipeline - {TEST} environment (uksb-1tu7mtee2)',
        target_environment=TEST,
        target_branch=raw_mappings[TEST][CODE_BRANCH],
        target_aws_env=test_aws_env,
        env=deployment_aws_env,
    )
    tag(test_pipeline_stack, DEPLOYMENT)

if os.environ.get('ENV', PROD) == PROD:
    target_environment = PROD
    prod_account = raw_mappings[PROD][ACCOUNT_ID]
    prod_region = raw_mappings[PROD][REGION]
    prod_aws_env = {
        'account': prod_account,
        'region': prod_region,
    }
    prod_pipeline_stack = PipelineStack(
        app,
        f'{target_environment}-{logical_id_prefix}EtlPipeline',
        description=f'InsuranceLake stack for ETL pipeline - {PROD} environment (uksb-1tu7mtee2)',
        target_environment=PROD,
        target_branch=raw_mappings[PROD][CODE_BRANCH],
        target_aws_env=prod_aws_env,
        env=deployment_aws_env,
    )
    tag(prod_pipeline_stack, DEPLOYMENT)

    # TODO: Modify replication bucket to have access logs and key rotation
    # Apply tagging to cross-region support stacks
    for stack in app.node.children:
        # All other stacks in the app are custom constructs
        if type(stack) == cdk.Stack:
            # Use the deployment environment for tagging because there
            # is no way to determine 1:1 which pipeline created the stack
            tag(stack, DEPLOYMENT)

            NagSuppressions.add_resource_suppressions(stack, [
                {
                    'id': 'AwsSolutions-S1',
                    'reason': 'Cross-region support stack and bucket are auto-created by Codepipeline'
                },
                {
                    'id': 'AwsSolutions-KMS5',
                    'reason': 'Cross-region support stack and bucket are auto-created by Codepipeline'
                },
            ], apply_to_children=True)

app.synth()
app.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
from constructs import Construct
import aws_cdk as cdk
import aws_cdk.aws_s3 as s3
import aws_cdk.aws_athena as athena
from cdk_nag import NagSuppressions

from .stack_import_helper import ImportedBuckets
from .configuration import (
    DEV, PROD, TEST, get_logical_id_prefix, get_resource_name_prefix, get_environment_configuration,
)

class AthenaHelperStack(cdk.Stack):
    def __init__(
            self, 
            scope: Construct, 
            construct_id: str, 
            target_environment: str, 
            glue_scripts_temp_bucket: s3.Bucket,
            **kwargs
        ):
        """CloudFormation stack to create DynamoDB Tables

        Parameters
        ----------
        scope
            Parent of this stack, usually an App or a Stage, but could be any construct
        construct_id
            The construct ID of this stack; if stackName is not explicitly defined,
            this ID (and any parent IDs) will be used to determine the physical ID of the stack
        target_environment
            The target environment for stacks in the deploy stage
        kwargs: optional
            Optional keyword arguments to pass up to parent Stack class
        """
        super().__init__(scope, construct_id, **kwargs)

        self.mappings = get_environment_configuration(target_environment)
        logical_id_prefix = get_logical_id_prefix()
        resource_name_prefix = get_resource_name_prefix().replace('-', '_')

        if target_environment == PROD or target_environment == TEST:
            recursive_delete = False
        else:
            recursive_delete = True

        buckets = ImportedBuckets(self, logical_id_suffix='AthenaHelperStack')

        athena_workgroup = athena.CfnWorkGroup(
            self,
            f'{target_environment}{logical_id_prefix}AthenaWorkgroup',
            name=resource_name_prefix,
            description='InsuranceLake ETL Maintainer helper workgroup ',
            recursive_delete_option=recursive_delete,
            work_group_configuration=athena.CfnWorkGroup.WorkGroupConfigurationProperty(
                publish_cloud_watch_metrics_enabled=True,
                result_configuration=athena.CfnWorkGroup.ResultConfigurationProperty(
                    output_location=f's3://{glue_scripts_temp_bucket.bucket_name}/query-results/',
                    # Glue Scripts Temp Bucket from Glue Stack has server-side encryption enabled
                    encryption_configuration=athena.CfnWorkGroup.EncryptionConfigurationProperty(
                        encryption_option='SSE_KMS',
                        kms_key=buckets.s3_kms_key.key_arn,
                    )
                )
            )
        )
        NagSuppressions.add_resource_suppressions(athena_workgroup, [
            {
                'id': 'AwsSolutions-ATH1',
                'reason': 'Glue Scripts Temp Bucket from Glue Stack has server-side encryption enabled; query results encryption would be redundant'
            },
        ])
lib/athena_helper_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import aws_cdk as cdk
from constructs import Construct
import aws_cdk.aws_iam as iam
import aws_cdk.aws_codecommit as CodeCommit

from .configuration import (
    CODECOMMIT_MIRROR_REPOSITORY_NAME, CODECOMMIT_MIRROR_REPOSITORY_NAME,
	get_logical_id_prefix, get_resource_name_prefix, get_all_configurations
)


class CodeCommitStack(cdk.Stack):

    def __init__(
        self, scope: Construct, construct_id: str,
        target_environment: str,
        **kwargs
    ):
        """CloudFormation stack to create CodeCommit mirror repository, if needed.

        Parameters
        ----------
        scope
            Parent of this stack, usually an App or a Stage, but could be any construct
        construct_id
            The construct ID of this stack; if stackName is not explicitly defined,
            this ID (and any parent IDs) will be used to determine the physical ID of the stack
        target_environment
            The target environment for the CodeCommit repository
        kwargs: optional
            Optional keyword arguments to pass up to parent Stack class
        """
        super().__init__(scope, construct_id, **kwargs)

        self.mappings = get_all_configurations()
        self.create_mirror_repository(target_environment)

    def create_mirror_repository(self, target_environment):
        """Creates CodeCommit repository to mirror source code from an unsupported version
        control system (e.g. Bitbucket, Gitlab) and integrate with CodePipeline.

        Parameters
        ----------
        target_environment
            The target environment for the CodeCommit repository
        """
        logical_id_prefix = get_logical_id_prefix()
        resource_name_prefix = get_resource_name_prefix()

        repo = CodeCommit.Repository(
            self, 
            f'{target_environment}{logical_id_prefix}EtlMirrorRepository',
            description='InsuranceLake ETL source code repository mirror for CodePipeline integration',
            repository_name=self.mappings[target_environment][CODECOMMIT_MIRROR_REPOSITORY_NAME],
        )

        git_mirror_user = iam.User(
            self,
            f'{logical_id_prefix}EtlGitExternalMirrorUser',
            user_name=f'{resource_name_prefix}-etl-git-mirror',
        )

        git_mirror_user.attach_inline_policy(
            iam.Policy(
                self,
                f'{logical_id_prefix}EtlCodecommitPushPullPolicy',
                statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        actions=[
                            'codecommit:GitPull',
                            'codecommit:GitPush',
                        ],
                        resources=[
                            repo.repository_arn
                        ],
                    )
                ]
            )
        )

        cdk.CfnOutput(
            self,
            f'{target_environment}{logical_id_prefix}CodeCommitMirrorRepositoryName',
            value=repo.repository_name,
            export_name=f'{target_environment}EtlMirrorRepository'
        )

        cdk.CfnOutput(
            self,
            f'{target_environment}{logical_id_prefix}CodeCommitMirrorUser',
            value=git_mirror_user.user_name,
            export_name=f'{target_environment}EtlMirrorUser'
        )
lib/code_commit_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import re
import boto3

# Environments (targeted at accounts)
DEPLOYMENT = 'Deploy'
DEV = 'Dev'
TEST = 'Test'
PROD = 'Prod'

# The following constants are used to map to parameter/secret paths
ENVIRONMENT = 'environment'

# Manual Inputs
GITHUB_REPOSITORY_OWNER_NAME = 'github_repository_owner_name'
GITHUB_REPOSITORY_NAME = 'github_repository_name'
CODECOMMIT_REPOSITORY_NAME = 'codecommit_repository_name'
CODECOMMIT_MIRROR_REPOSITORY_NAME = 'codecommit_mirror_repository_name'
ACCOUNT_ID = 'account_id'
REGION = 'region'
VPC_CIDR = 'vpc_cidr'
LOGICAL_ID_PREFIX = 'logical_id_prefix'
RESOURCE_NAME_PREFIX = 'resource_name_prefix'
CODE_BRANCH = 'code_branch'
LINEAGE='lineage'

# Secrets Manager Inputs
GITHUB_TOKEN = 'github_token'

# Used in Automated Outputs
VPC_ID = 'vpc_id'
AVAILABILITY_ZONE_1 = 'availability_zone_1'
AVAILABILITY_ZONE_2 = 'availability_zone_2'
AVAILABILITY_ZONE_3 = 'availability_zone_3'
SUBNET_ID_1 = 'subnet_id_1'
SUBNET_ID_2 = 'subnet_id_2'
SUBNET_ID_3 = 'subnet_id_3'
ROUTE_TABLE_1 = 'route_table_1'
ROUTE_TABLE_2 = 'route_table_2'
ROUTE_TABLE_3 = 'route_table_3'
SHARED_SECURITY_GROUP_ID = 'shared_security_group_id'
S3_KMS_KEY = 's3_kms_key'
S3_ACCESS_LOG_BUCKET = 's3_access_log_bucket'
S3_RAW_BUCKET = 's3_raw_bucket'
S3_CONFORMED_BUCKET = 's3_conformed_bucket'
S3_PURPOSE_BUILT_BUCKET = 's3_purpose_built_bucket'
CROSS_ACCOUNT_DYNAMODB_ROLE = 'cross_account_dynamodb_role'
STATE_MACHINE = 'sfn_state_machine'
NOTIFICATION_TOPIC = 'sns_topic'

MAX_S3_BUCKET_NAME_LENGTH = 63

def get_local_configuration(environment: str, local_mapping: dict = None) -> dict:
    """Provides manually configured variables that are validated for quality and safety

    Parameters
    ----------
    environment
        The environment used to retrieve corresponding configuration
    local_mapping: optional
        Optional override the embedded local_mapping; used for testing

    Raises
    ------
    AttributeError
        If the resource_name_prefix does not conform or if the requested
        environment does not exist

    Returns
    -------
    dict
        Configuration for the requested environment
    """
    active_account_id = boto3.client("sts").get_caller_identity()["Account"]

    if local_mapping is None:
        local_mapping = {
            DEPLOYMENT: {
                ACCOUNT_ID: active_account_id,
                REGION: 'us-east-2',

                # If you use GitHub / GitHub Enterprise, this will be the organization name
                GITHUB_REPOSITORY_OWNER_NAME: '',

                # Use your forked Github repo here!
                # Leave empty if you do not use Github
                GITHUB_REPOSITORY_NAME: '',

                # Use only if your repository is already in CodecCommit, otherwise leave empty!
                # Use your CodeCommit repo name here
                CODECOMMIT_REPOSITORY_NAME: '',

                # Use only if you do NOT use Github or CodeCommit and need to mirror your repository
                # Name your CodeCommit mirror repo here (recommend matching your external repo)
                # Leave empty if you use Github or your repository is in CodeCommit already
                CODECOMMIT_MIRROR_REPOSITORY_NAME: 'aws-insurancelake-etl',

                # This is used in the Logical Id of CloudFormation resources.
                # We recommend Capital case for consistency, e.g. DataLakeCdkBlog
                LOGICAL_ID_PREFIX: 'InsuranceLake',

                # Important: This is used as a prefix for resources that must be **globally** unique!
                # Resource names may only contain alphanumeric characters, hyphens, and cannot contain trailing hyphens.
                # S3 bucket names from this application must be under the 63 character bucket name limit
                RESOURCE_NAME_PREFIX: 'insurancelake',
            },
            DEV: {
                ACCOUNT_ID: active_account_id,
                REGION: 'us-east-2',
                LINEAGE: True,
                # VPC_CIDR: '10.20.0.0/24',
                CODE_BRANCH: 'develop',
            },
            TEST: {
                ACCOUNT_ID: active_account_id,
                REGION: 'us-east-2',
                LINEAGE: True,
                # VPC_CIDR: '10.10.0.0/24',
                CODE_BRANCH: 'test',
            },
            PROD: {
                ACCOUNT_ID: active_account_id,
                REGION: 'us-east-2',
                LINEAGE: True,
                # VPC_CIDR: '10.0.0.0/24',
                CODE_BRANCH: 'master',
            }
        }

    resource_prefix = local_mapping[DEPLOYMENT][RESOURCE_NAME_PREFIX]
    if (
        not re.fullmatch('^[a-z0-9-]+', resource_prefix)
        or '-' in resource_prefix[-1:] or '-' in resource_prefix[1]
    ):
        raise AttributeError('Resource names may only contain lowercase alphanumeric and hyphens '
                        'and cannot contain leading or trailing hyphens')

    for each_env in local_mapping:
        # NOTE: Resource with longest bucket name is from the infra
        #       code base, but we will assume the user wants to have
        #       a consistent resource prefix across all stacks
        longest_bucket_name = \
            f'{each_env}-{resource_prefix}-{local_mapping[each_env][ACCOUNT_ID]}-{local_mapping[each_env][REGION]}-access-logs'
        if len(longest_bucket_name) > MAX_S3_BUCKET_NAME_LENGTH:
            raise AttributeError('Resource name prefix is too long; at least one S3 bucket name '
                        f'would exceed maximum allowed length of {MAX_S3_BUCKET_NAME_LENGTH} '
                        f'characters, e.g. {longest_bucket_name}')

    if environment not in local_mapping:
        raise AttributeError(f'The requested environment: {environment} does not exist in local mappings')

    return local_mapping[environment]


def get_environment_configuration(environment: str, local_mapping: dict = None) -> dict:
    """Provides all configuration values for the given target environment

    Parameters
    ----------
    environment
        The environment used to retrieve corresponding configuration
    local_mapping: optional
        Optionally override the embedded local_mapping; used for testing

    Returns
    -------
    dict
        Combined configuration and Cloudformation output names for target environment
    """
    cloudformation_output_mapping = {
        ENVIRONMENT: f'{environment}',
        VPC_ID: f'{environment}VpcId',
        AVAILABILITY_ZONE_1: f'{environment}AvailabilityZone1',
        AVAILABILITY_ZONE_2: f'{environment}AvailabilityZone2',
        AVAILABILITY_ZONE_3: f'{environment}AvailabilityZone3',
        SUBNET_ID_1: f'{environment}SubnetId1',
        SUBNET_ID_2: f'{environment}SubnetId2',
        SUBNET_ID_3: f'{environment}SubnetId3',
        ROUTE_TABLE_1: f'{environment}RouteTable1',
        ROUTE_TABLE_2: f'{environment}RouteTable2',
        ROUTE_TABLE_3: f'{environment}RouteTable3',
        SHARED_SECURITY_GROUP_ID: f'{environment}SharedSecurityGroupId',
        S3_KMS_KEY: f'{environment}S3KmsKeyArn',
        S3_ACCESS_LOG_BUCKET: f'{environment}S3AccessLogBucket',
        S3_RAW_BUCKET: f'{environment}CollectBucketName',
        S3_CONFORMED_BUCKET: f'{environment}CleanseBucketName',
        S3_PURPOSE_BUILT_BUCKET: f'{environment}ConsumeBucketName',
        CROSS_ACCOUNT_DYNAMODB_ROLE: f'{environment}CrossAccountDynamoDbRoleArn',
        STATE_MACHINE: f'{environment}StepFunctionsStateMachineName',
        NOTIFICATION_TOPIC: f'{environment}EtlNotificationSnsTopicName',
    }

    return {**cloudformation_output_mapping, **get_local_configuration(environment, local_mapping = local_mapping)}


def get_all_configurations() -> dict:
    """Returns a dict mapping of configurations for all environments.
    These keys correspond to static values, CloudFormation outputs, and Secrets Manager
    (passwords only) records.

    Returns
    -------
    dict
        Combined configuration and Cloudformation output names for all environments
    """
    return {
        DEPLOYMENT: {
            ENVIRONMENT: DEPLOYMENT,
            GITHUB_TOKEN: '/InsuranceLake/GitHubToken',
            **get_local_configuration(DEPLOYMENT),
        },
        DEV: get_environment_configuration(DEV),
        TEST: get_environment_configuration(TEST),
        PROD: get_environment_configuration(PROD),
    }


def get_logical_id_prefix() -> str:
    """Returns the logical id prefix to apply to all CloudFormation resources

    Returns
    -------
    str
        Logical ID prefix from deployment configuration
    """
    return get_local_configuration(DEPLOYMENT)[LOGICAL_ID_PREFIX]


def get_resource_name_prefix() -> str:
    """Returns the resource name prefix to apply to all resources names

    Returns
    -------
    str
        Resource name prefix from deployment configuration
    """
    return get_local_configuration(DEPLOYMENT)[RESOURCE_NAME_PREFIX]
lib/configuration.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
from constructs import Construct
import aws_cdk as cdk
import aws_cdk.aws_dynamodb as dynamodb

from .stack_import_helper import ImportedBuckets
from .configuration import (
    DEV, PROD, TEST, LINEAGE,
    get_logical_id_prefix, get_resource_name_prefix, get_environment_configuration,
)

class DynamoDbStack(cdk.Stack):
    def __init__(self, scope: Construct, construct_id: str, target_environment: str, **kwargs):
        """CloudFormation stack to create DynamoDB Tables.

        Parameters
        ----------
        scope
            Parent of this stack, usually an App or a Stage, but could be any construct
        construct_id
            The construct ID of this stack; if stackName is not explicitly defined,
            this ID (and any parent IDs) will be used to determine the physical ID of the stack
        target_environment
            The target environment for stacks in the deploy stage
        kwargs: optional
            Optional keyword arguments to pass up to parent Stack class
        """
        super().__init__(scope, construct_id, **kwargs)

        self.mappings = get_environment_configuration(target_environment)
        logical_id_prefix = get_logical_id_prefix()
        resource_name_prefix = get_resource_name_prefix().replace('-', '_')

        buckets = ImportedBuckets(self, logical_id_suffix='DynamoDbStack')

        # DynamoDB stack will allow tables to be destroyed (along with the stack) in the Dev environment only
        self.removal_policy = cdk.RemovalPolicy.DESTROY
        if (target_environment == PROD or target_environment == TEST):
            self.removal_policy = cdk.RemovalPolicy.RETAIN

        # DynamoDB table to store audit log and step function state
        self.job_audit_table = dynamodb.Table(
            self,
            f'{target_environment}{logical_id_prefix}EtlAuditTable',
            table_name=f'{target_environment.lower()}-{resource_name_prefix}-etl-job-audit',
            partition_key=dynamodb.Attribute(name='execution_id', type=dynamodb.AttributeType.STRING),
            billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
            encryption=dynamodb.TableEncryption.DEFAULT,
            point_in_time_recovery=True,
            removal_policy=self.removal_policy,
        )

        # DynamoDB table to store raw data to hash value mapping
        # raw data stored in 'raw_data' column
        self.hash_values_table = dynamodb.Table(
            self,
            f'{target_environment}{logical_id_prefix}HashValuesTable',
            table_name=f'{target_environment.lower()}-{resource_name_prefix}-etl-hash-values',
            partition_key=dynamodb.Attribute(name='hash_key', type=dynamodb.AttributeType.STRING),
            billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
            encryption=dynamodb.TableEncryption.CUSTOMER_MANAGED,
            encryption_key=buckets.s3_kms_key,
            point_in_time_recovery=True,
            removal_policy=self.removal_policy,
        )

        # DynamoDB table to store lookup values for lookup transform
        # Lookup values stored in 'lookup_data' column
        self.value_lookup_table = dynamodb.Table(
            self,
            f'{target_environment}{logical_id_prefix}ValueLookupTable',
            table_name=f'{target_environment.lower()}-{resource_name_prefix}-etl-value-lookup',
            partition_key=dynamodb.Attribute(name='source_system', type=dynamodb.AttributeType.STRING),
            sort_key=dynamodb.Attribute(name='column_name', type=dynamodb.AttributeType.STRING),
            billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
            encryption=dynamodb.TableEncryption.DEFAULT,
            point_in_time_recovery=True,
            removal_policy=self.removal_policy,
        )

        # DynamoDB table to store lookup values for multi-value
        # lookup transform
        # Composite hash of all keys stored in 'lookup_item'
        # Lookup return values stored in user defined columns
        self.multi_lookup_table = dynamodb.Table(
            self,
            f'{target_environment}{logical_id_prefix}MultiValueLookupTable',
            table_name=f'{target_environment.lower()}-{resource_name_prefix}-etl-multi-lookup',
            partition_key=dynamodb.Attribute(name='lookup_group', type=dynamodb.AttributeType.STRING),
            sort_key=dynamodb.Attribute(name='lookup_item', type=dynamodb.AttributeType.STRING),
            billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
            encryption=dynamodb.TableEncryption.DEFAULT,
            point_in_time_recovery=True,
            removal_policy=self.removal_policy,
        )

        # DynamoDB table to store Glue Data Quality results
        self.dq_results_table = dynamodb.Table(
            self,
            f'{target_environment}{logical_id_prefix}DqResultsTable',
            table_name=f'{target_environment.lower()}-{resource_name_prefix}-etl-dq-results',
            partition_key=dynamodb.Attribute(name='execution_id', type=dynamodb.AttributeType.STRING),
            sort_key=dynamodb.Attribute(name='job_id_action_rule', type=dynamodb.AttributeType.STRING),
            billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
            encryption=dynamodb.TableEncryption.DEFAULT,
            point_in_time_recovery=True,
            removal_policy=self.removal_policy,
        )

        if self.mappings[LINEAGE]:
            # DynamoDB table to store custom data lineage logging per Step Function exection
            self.data_lineage_table = dynamodb.Table(
                self,
                f'{target_environment}{logical_id_prefix}DataLineageTable',
                table_name=f'{target_environment.lower()}-{resource_name_prefix}-etl-data-lineage',
                partition_key=dynamodb.Attribute(name='step_function_execution_id', type=dynamodb.AttributeType.STRING),
                sort_key=dynamodb.Attribute(name='job_id_operation_seq', type=dynamodb.AttributeType.STRING),
                billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
                encryption=dynamodb.TableEncryption.DEFAULT,
                point_in_time_recovery=True,
                removal_policy=self.removal_policy,
            )
        else:
            self.data_lineage_table = None
lib/dynamodb_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import json
import boto3
import botocore
import os
import logging
import os.path
from datetime import datetime
import dateutil.tz

# Logger initiation
logger = logging.getLogger()
logger.setLevel(logging.INFO)

def lambda_handler(event: dict, context: dict) -> dict:
    """Lambda function's entry point. This function receives a success/fail
    event from Step Functions State machine, transforms error message, and
    updates DynamoDB table.

    Parameters
    ----------
    event
        State Machine event dictionary
    context
        Lambda context dictionary

    Returns
    -------
    dict
        Lambda result dictionary
    """
    dynamo_client = boto3.resource('dynamodb')
    audit_table_name = os.environ['DYNAMODB_TABLE_NAME']
    print(event)
    print(event['Input']['execution_id'])
    execution_id = event['Input']['execution_id']

    # Capturing the current time in UTC
    now = datetime.now(tz=dateutil.tz.gettz('UTC'))
    audit_message_timestamp = now.strftime('%Y-%m-%d %H:%M:%S.%f')
    return_message = ''

    if "JobRunState" in event['Input']['taskresult']:
        print(f"JobRunState exists: {event['Input']['taskresult']['JobRunState']}")
        status = event['Input']['taskresult']['JobRunState']

        try:
            # Update audit table
            table = dynamo_client.Table(audit_table_name)
            table.update_item(
                Key={
                    'execution_id': execution_id
                },
                UpdateExpression="set job_last_updated_timestamp=:lut,job_latest_status=:sts",
                ExpressionAttributeValues={
                    ':sts': status,
                    ':lut': audit_message_timestamp,
                },
                ReturnValues="UPDATED_NEW"
            )
            return_message = f'status {status}'
        except botocore.exceptions.ClientError as error:
            logger.error(f'DynamoDB update_item failed: {error}')
            raise error
    else:
        print("JobRunState does not exist, assuming FAILED")
        status = "FAILED"
        try:
            error_event = json.loads(event['Input']['taskresult']['Cause'])
            error_message = error_event.get('ErrorMessage', error_event['JobRunState'])
        except json.JSONDecodeError:
            error_message = event['Input']['taskresult']['Cause']
        print(error_message)

        try:
            # Update table
            table = dynamo_client.Table(audit_table_name)
            table.update_item(
                Key={
                    'execution_id': execution_id
                },
                UpdateExpression="set joblast_updated_timestamp=:lut,job_latest_status=:sts,error_message=:emsg",
                ExpressionAttributeValues={
                    ':sts': status,
                    ':lut': audit_message_timestamp,
                    ':emsg': error_message,
                },
                ReturnValues="UPDATED_NEW"
            )
            return_message = f"status {status} {error_message}"
        except botocore.exceptions.ClientError as error:
            logger.error(f'DynamoDB update_item failed: {error}')
            raise error

    return {
        'statusCode': 200,
        'body': f'DynamoDB status updated with {return_message}'
    }
lib/etl_job_auditor/lambda_handler.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import boto3
import sys
import time
import os
import re

from pyspark.context import SparkContext

from awsglue.utils import getResolvedOptions
from awsglue.context import GlueContext
from awsglue.job import Job

# For running in local Glue container
sys.path.append(os.path.dirname(__file__) + '/lib')
from glue_catalog_helpers import create_database
from datalineage import DataLineageGenerator

expected_arguments = [
    'JOB_NAME',
    'environment',
    'TempDir',
    'txn_bucket',
    'txn_sql_prefix_path',
    'source_bucket',
    'target_bucket',
    'state_machine_name',
    'execution_id',
    'database_name_prefix',
    'table_name',
    'base_file_name',
    'p_year',
    'p_month',
    'p_day',
]

# Handle optional arguments
for arg in sys.argv:
    if '--data_lineage_table' in arg:
        expected_arguments.append('data_lineage_table')


def athena_execute_query(database: str, query: str, result_bucket: str, max_attempts: int = 15) -> str:
    """Function to execute query using Athena boto client, loop until
    result is returned, and return status.

    Parameters
    ----------
    aws_account_id
        AWS Account ID
    database
        Athena database in which to run the query
    query
        Single or multi-line query to run
    max_attempts
        Number of loops (1s apart) to attempt to get query status

    Returns
    -------
    str
        Status of query result execution
    """
    athena = boto3.client('athena')

    query_response = athena.start_query_execution(
            QueryExecutionContext = {'Database': database},
            QueryString = query,
            ResultConfiguration = {'OutputLocation': result_bucket + '/'}
        )
    print(f'Executed query response: {query_response}')

    # Valid "state" Values: QUEUED | RUNNING | SUCCEEDED | FAILED | CANCELLED
    attempts = max_attempts
    status = 'QUEUED'
    while status in [ 'QUEUED', 'RUNNING' ]:
        if attempts == 0:
            raise RuntimeError('athena_execute_query() exceeded max_attempts waiting for query')

        query_details = athena.get_query_execution(QueryExecutionId = query_response['QueryExecutionId'])
        print(f'Get query execution response: {query_details}')
        status = query_details['QueryExecution']['Status']['State']
        if status in [ 'FAILED', 'CANCELED' ]:
            # Raise errors that are not raised from StartQueryExecution
            raise RuntimeError("athena_execute_query() failed with query engine error: "
                              f"{query_details['QueryExecution']['Status'].get('StateChangeReason', '')}")
        if status != 'SUCCEEDED':
            time.sleep(1)
        attempts -= 1

    return status


def main():
    args = getResolvedOptions(sys.argv, expected_arguments)

    sc = SparkContext()
    glueContext = GlueContext(sc)
    spark = glueContext.spark_session
    job = Job(glueContext)
    job.init(args['JOB_NAME'], args)

    hadoop_conf = sc._jsc.hadoopConfiguration()
    hadoop_conf.set("fs.s3.impl", "org.apache.hadoop.fs.s3a.S3AFileSystem")

    # Argument passed as database_name_prefix will have mixed case and match the S3 file system
    # Glue Catalog Databases will be lower case only
    source_database = args['database_name_prefix'].lower()
    target_database = f"{args['database_name_prefix'].lower()}_consume"
    target_table = args['table_name']
    substitution_data = args.copy()
    substitution_data.update({
        'source_database': source_database,
        'target_database': target_database,
    })
    sql_prefix = args['txn_sql_prefix_path']

    # TODO: Check if Lakeformation permissions break Spark SQL query and find workaround
    spark_sql_key = f"{sql_prefix[1:]}spark-{args['database_name_prefix']}-{args['table_name']}.sql"
    print(f'Using Spark SQL transformation from: {spark_sql_key}')
    try:
        spark_sql_data = sc.textFile(f"{args['txn_bucket']}/{spark_sql_key}")
        spark_sql = '\n'.join(spark_sql_data.collect())
    except Exception as e:
        print(f'No Spark SQL transformation exists or error reading: {e.java_exception.getMessage()}, skipping')
        spark_sql = None

    # Spark SQL is used to populate the target database, stored in S3
    if spark_sql is not None:
        re_create_table = re.compile(r'\s*CREATE[\s\w]*\s+TABLE\s+["`\']?([\w]+)["`\']?\s+AS(.*)', re.IGNORECASE | re.DOTALL)
        match_create_table = re_create_table.match(spark_sql)
        if match_create_table:
            print(f'Table name override detected: {match_create_table.group(1)}')
            target_table = match_create_table.group(1)
            spark_sql = match_create_table.group(2)

        substitution_data.update({ 'target_table': target_table })
        spark_sql = spark_sql.strip()
        # Make sure there is something to execute
        if not spark_sql:
            raise RuntimeError('Spark SQL file found but appears to be empty')

        print(f'Executing Spark SQL: {spark_sql.format(**substitution_data)}')
        # Spark SQL should query entire table, because the existing consume table will be overwritten
        df = spark.sql(spark_sql.format(**substitution_data))
        print(f'Retrieved (from Cleanse) dataframe row count: {df.count()}'
              f'  column count: {len(df.columns)}'
              f'  number of partitions: {df.rdd.getNumPartitions()}')

        df.cache()
        lineage = DataLineageGenerator(args)
        dataset_name = args['database_name_prefix'] + '/' + args['table_name']
        lineage.update_lineage(df, dataset_name, 'numericaudit')
        lineage.update_lineage(df, dataset_name, 'sparksql', transform=[ spark_sql ])

        storage_location = f"{args['target_bucket']}/{args['database_name_prefix']}/{target_table}"

        create_database(target_database)
        df.write.partitionBy('year', 'month', 'day').saveAsTable(
            f'{target_database}.{target_table}',
            path=storage_location,
            format='parquet',
            mode='overwrite',
        )
        df.unpersist()

    # Athena SQL is used to create views
    athena_sql_key = f"{sql_prefix[1:]}athena-{args['database_name_prefix']}-{args['table_name']}.sql"
    print(f'Using Athena SQL transformation from: {athena_sql_key}')
    try:
        athena_sql_data = sc.textFile(f"{args['txn_bucket']}/{athena_sql_key}")
        # Format with line breaks and preserve whitespace between lines so logging is easier to read
        athena_sql = '\n'.join(athena_sql_data.collect())
    except Exception as e:
        print(f'No Athena SQL transformation exists or error reading: {e.java_exception.getMessage()}, skipping')
        athena_sql = None

    if athena_sql:
        for sql in athena_sql.split(';'):
            sql = sql.strip()
            # Check if there is actual SQL content since split will return an element after the final ;
            if sql:
                print(f'Executing Athena SQL: {sql.format(**substitution_data)}')
                status = athena_execute_query(target_database, sql.format(**substitution_data), args['TempDir'])
                print(f'Athena query execution status: {status}')

    job.commit()


if __name__ == '__main__':
    main()
lib/glue_scripts/etl_cleanse_to_consume.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import sys
import csv
import json
import os
from functools import reduce

from pyspark.context import SparkContext
from pyspark.sql.functions import trim
from pyspark.sql.utils import IllegalArgumentException

from awsglue.utils import getResolvedOptions
from awsglue.context import GlueContext
from awsglue.job import Job
from awsglue.dynamicframe import DynamicFrame

# For running in local Glue container
sys.path.append(os.path.dirname(__file__) + '/lib')

from glue_catalog_helpers import upsert_catalog_table, clean_column_names, generate_spec, put_s3_object
from custom_mapping import custommapping
from datatransform_typeconversion import *
from datatransform_dataprotection import *
from datatransform_regex import *
from datatransform_lookup import *
from datatransform_premium import *
from datalineage import DataLineageGenerator

try:
    # Try/Except block for running in AWS-provided Glue container until Glue DQ support is added
    from dataquality_check import check_dataquality_warn, check_dataquality_quarantine, check_dataquality_halt
except ModuleNotFoundError:
    pass

expected_arguments = [
    'JOB_NAME',
    'environment',
    'TempDir',
    'txn_bucket',
    'txn_spec_prefix_path',
    'source_bucket',
    'target_bucket',
    'hash_value_table',
    'value_lookup_table',
    'multi_lookup_table',
    'dq_results_table',
    'state_machine_name',
    'execution_id',
    'source_key',
    'target_database_name',
    'table_name',
    'base_file_name',
    'p_year',
    'p_month',
    'p_day',
]

# Handle optional arguments
for arg in sys.argv:
    if '--data_lineage_table' in arg:
        expected_arguments.append('data_lineage_table')


def main():
    args = getResolvedOptions(sys.argv, expected_arguments)

    sc = SparkContext()
    glueContext = GlueContext(sc)
    spark = glueContext.spark_session
    job = Job(glueContext)
    job.init(args['JOB_NAME'], args)

    filename, ext = os.path.splitext(args['base_file_name'])
    source_path = args['source_bucket'] + '/' + args['source_key'] + '/' + args['base_file_name']
    print(f'Source path: {source_path}')

    # Set default schema change detection based on InsuranceLake environment
    if args['environment'] == 'Dev':
        allow_schema_change = 'permissive'
    elif args['environment'] == 'Test':
        allow_schema_change = 'reorder'
    else:
        allow_schema_change = 'strict'

    # Read input/transformation spec file
    txn_spec_prefix = args['txn_spec_prefix_path']
    txn_spec_key = txn_spec_prefix[1:] + args['target_database_name'] + '-' + args['table_name'] + '.json'
    print(f'Using input/transformation specification from: {txn_spec_key}')
    try:
        spec_data = sc.textFile(f"{args['txn_bucket']}/{txn_spec_key}")
        spec_json_data = json.loads('\n'.join(spec_data.collect()))
    except Exception as e:
        print(f'No input/transformation spec file exists or error reading: {e.java_exception.getMessage()}')
        spec_json_data = None

    # Read schema mapping file
    txn_spec_prefix = args['txn_spec_prefix_path']
    txn_map_key = txn_spec_prefix[1:] + args['target_database_name'] + '-' + args['table_name'] + '.csv'
    print(f'Using schema mapping from: {txn_map_key}')
    try:
        mapping_file = sc.textFile(f"{args['txn_bucket']}/{txn_map_key}").collect()
        reader = csv.DictReader(mapping_file)
        # Implement case insensitive field name matching
        reader.fieldnames = [ field_name.lower() for field_name in reader.fieldnames ]
        mapping_data = [ field_map for field_map in reader ]
    except Exception as e:
        print(f'No transformation mapping file exists or error reading: {e.java_exception.getMessage()}, skipping')
        mapping_data = None

    # Read collected S3 object
    input_spec = spec_json_data.get('input_spec', {}) if spec_json_data else {}
    if 'fixed' in input_spec:
        # mapping_data should be SourceName (ignored for fixed format), DestinationName, (Column) Width
        initial_df = spark.read.text(source_path)
        initial_df = initial_df.select(
            [
                trim(initial_df.value.substr(
                    # Calculate start index based on all prior column widths
                    reduce(lambda a, b: a + int(b['width']), mapping_data[:index], 0) + 1,
                    int(field_data['width'])
                )).alias(field_data['destname'])
                    for index, field_data in enumerate(mapping_data)
                        if field_data['destname'].lower() != 'null'
            ]
        )
        print('Performed fixed width file load and field mapping')

    elif ext.lower() in [ '.xlsx', '.xls' ]:
        if 'excel' in input_spec:
            sheet_names = input_spec['excel']['sheet_names']
            data_address = input_spec['excel']['data_address']
            # Retrieve and remove the password so it is not printed
            workbook_password = input_spec['excel'].pop('password', None)
            print(f"Using Excel input specification: {input_spec['excel']}")
            header = input_spec['excel'].get('header', True)
        else:
            print(f'No Excel input specification, using defaults')
            sheet_names = [ '0' ]
            data_address = 'A1'
            header = True
            workbook_password = None

        for sheet_name in sheet_names:
            try:
                initial_df = spark.read.format('excel') \
                    .option('header', header) \
                    .option('inferSchema', 'true') \
                    .option('mode', 'PERMISSIVE') \
                    .option("usePlainNumberFormat", 'true') \
                    .option('dataAddress', f"'{sheet_name}'!{data_address}") \
                    .option('workbookPassword', workbook_password) \
                    .load(source_path)
                print(f"Found table at '{sheet_name}'!{data_address}")
                break
            except IllegalArgumentException:
                pass

        if 'initial_df' not in locals():
            raise RuntimeError(f'None of sheet names {sheet_names} found in Excel workbook')

    elif ext.lower() == '.parquet':
        # TODO: Support partitioned Parquet folder structures (read correctly and repartition later)
        initial_df = spark.read.format('parquet').load(source_path)

    else:
        # Comma delimited is default, with support for tab delimited if specified
        delimiter = ','
        header = True
        if 'csv' in input_spec:
            header = input_spec['csv'].get('header', True)
        if 'tsv' in input_spec:
            delimiter = '\t'
            header = input_spec['tsv'].get('header', True)

        initial_df = spark.read.format('csv') \
            .option('header', header) \
            .option('delimiter', delimiter) \
            .option('inferSchema', 'true') \
            .option('mode', 'PERMISSIVE') \
            .load(source_path)

    initial_df.cache()
    lineage = DataLineageGenerator(args)
    lineage.update_lineage(initial_df, args['source_key'], 'read')
    lineage.update_lineage(initial_df, args['source_key'], 'numericaudit')

    if initial_df.count() == 0:
        raise RuntimeError('No rows of data in source file; aborting')

    # Skip fixed format because mapping is already done
    if 'fixed' not in input_spec:
        # Perform schema mapping if mapping data exists
        if mapping_data:
            initial_df = custommapping(initial_df, mapping_data, args, lineage)
            print('Performed field mapping')
        else:
            generated_map_path = args['TempDir'] + '/' + args['target_database_name'] + '-' + args['table_name'] + '.csv'
            print(f'No mapping found, generated recommended mapping to: {generated_map_path}')
            initial_df, generated_mapping_data = clean_column_names(initial_df)
            put_s3_object(generated_map_path, generated_mapping_data)

    totransform_df = initial_df
    totransform_df.cache()

    # Perform transforms
    if spec_json_data and 'transform_spec' in spec_json_data:
        transform_spec = spec_json_data['transform_spec']
        print(f'Using transformation specification: {transform_spec}')

        for transform in transform_spec.keys():
            if 'transform_' + transform not in globals():
                print(f'Transform function transform_{transform} called for in {txn_map_key} not implemented')
                continue

            transform_func = globals()['transform_' + transform]
            totransform_df = transform_func(
                totransform_df,
                transform_spec[transform],
                args,
                lineage,
                sc
            )
            print(f'Performed {transform} transform')
    else:
        generated_spec_path = args['TempDir'] + '/' + args['target_database_name'] + '-' + args['table_name'] + '.json'
        print(f'No transformation specification found, generating recommended spec to: {generated_spec_path}')
        generated_spec_data = generate_spec(totransform_df, ext)
        put_s3_object(generated_spec_path, json.dumps(generated_spec_data, indent=4))

    fields_to_add = {
        # Add State Machine execution ID to DataFrame for lineage tracking
        'execution_id':  args['execution_id'],
        # Add partition columns for daily data load partitioning strategy
        'year': args['p_year'],
        'month': args['p_month'],
        'day': args['p_day'],
    }
    transformed_df = transform_literal(totransform_df, fields_to_add, args, lineage)
    transformed_df.cache()
    print('Added partition columns and State Machine execution_id column')

    # Record final DataFrame schema
    transformed_schema = list(transformed_df.schema)
    print(str(transformed_schema))

    # Run Data Quality rules and take appropriate action
    dq_rules_key = 'etl/dq-rules/dq-' + args['target_database_name'] + '-' + args['table_name'] + '.json'
    print(f'Using data quality rules from: {dq_rules_key}')
    try:
        rules_data = sc.textFile(f"{args['txn_bucket']}/{dq_rules_key}")
        rules_json_data = json.loads('\n'.join(rules_data.collect()))
    except Exception as e:
        print(f'No data quality rules file exists or error reading: {e.java_exception.getMessage()}, skipping')
        rules_json_data = None

    # NOTE: Only quarantine rules will change filtered Dataframe
    filtered_df = transformed_df
    if rules_json_data:
        if 'dataquality_check' not in sys.modules:
            print('Skipping data quality rules because awsgluedq library is not available')
        else:
            dq_rules = rules_json_data['dq_rules']
            print(f'Using data quality rules: {dq_rules}')
            fordq_dyf = DynamicFrame.fromDF(transformed_df, glueContext, f"{args['execution_id']}-fordq_dyf")
            if 'warn_rules' in dq_rules:
                print(f'Glue Data Quality Warn-on-Failure results:')
                check_dataquality_warn(fordq_dyf, dq_rules['warn_rules'], args, sc)
            if 'quarantine_rules' in dq_rules:
                print(f'Glue Data Quality Quarantine-on-Failure results:')
                filtered_df.unpersist()
                filtered_df = check_dataquality_quarantine(fordq_dyf, dq_rules['quarantine_rules'], args, sc)
                lineage.update_lineage(filtered_df, args['source_key'], 'dataqualityquarantine', transform=dq_rules['quarantine_rules'])
            if 'halt_rules' in dq_rules:
                # NOTE: Original DynamicFrame is used to process halt rules, so even rows filtered above can trigger abort
                print(f'Glue Data Quality Halt-on-Failure results:')
                check_dataquality_halt(fordq_dyf, dq_rules['halt_rules'], args, sc)

    # Post transform and DQ filter numeric audit
    filtered_df.cache()
    lineage.update_lineage(filtered_df, args['source_key'], 'numericaudit')

    # Build Glue Catalog for data
    storage_location = args['target_bucket'] + '/' + args['target_database_name'] + '/' + args['table_name']
    allow_schema_change = input_spec.get('allow_schema_change', allow_schema_change)
    upsert_catalog_table(
            filtered_df,
            args['target_database_name'],
            args['table_name'],
            ['year', 'month', 'day'],
            storage_location,
            allow_schema_change,
        )

    # Set dynamic overwrite to preserve historical partitions, but overwrite current partition
    spark.conf.set('spark.sql.sources.partitionOverwriteMode', 'dynamic')
    spark.conf.set('hive.exec.dynamic.partition', 'true')
    spark.conf.set('hive.exec.dynamic.partition.mode', 'nonstrict')

    # Write the data with new schema and partitioning to the target location in hadoop parquet
    filtered_df.write.partitionBy('year', 'month', 'day').save(storage_location, format='parquet', mode='overwrite')

    # Do not use RECOVER PARTITIONS because we already updated the Glue Catalog if it was needed
    spark.sql(f"ALTER TABLE {args['target_database_name']}.{args['table_name']} ADD IF NOT EXISTS PARTITION"
        f"(year='{args['p_year']}', month='{args['p_month']}', day='{args['p_day']}')")

    job.commit()


if __name__ == '__main__':
    main()
lib/glue_scripts/etl_collect_to_cleanse.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
from awsglue.dynamicframe import DynamicFrame
from awsglue.transforms import ApplyMapping
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.functions import col
from rapidfuzz import fuzz
from rapidfuzz import process as fuzz_process
from rapidfuzz.utils import default_process

def custommapping(df: DataFrame, field_mapping_list: list, args: list, lineage) -> DataFrame:
    """Apply a custom field mapping to the data schema in a DataFrame
    Uses RapidFuzz for fuzzy matching: https://maxbachmann.github.io/RapidFuzz/

    Parameters
    ----------
    df
        Spark DataFrame on which to apply schema mapping
    field_mapping_list
        List of dictionary objects with the form:
            'sourcename': 'source_field_name',
            'destname': 'destination_field_name', use Null to drop a column
            'threshold': (optional) fuzzy match minimum confidence level
            'scorer': (optional, must be specified with threshhold) fuzzy match scoring algorithm
    args
        Glue job arguments, from which source_key and execution_id are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        Spark DataFrame with the custom mapping applied
    """
    # Direct field mapping rows have 2 values only
    field_map = { map_row['sourcename']: map_row['destname']
        for map_row in field_mapping_list if not map_row.get('threshold', None) }
    select_list = [ col('`' + field.name + '`').alias(field_map[field.name])
        for field in df.schema
            if field.name in field_map and field_map[field.name].lower() != 'null' ]

    # Only attempt to fuzzy match fields not directly mapped
    unmapped_fields = [ field.name for field in df.schema if field.name not in field_map ]

    # Perform fuzzy match with rapidfuzz, specified sort, and column alias
    if unmapped_fields:
        for map_row in field_mapping_list:
            if map_row.get('threshold', None):
                match, score, index = fuzz_process.extractOne(
                    map_row['sourcename'],
                    unmapped_fields,
                    processor=default_process,
                    scorer=getattr(fuzz, map_row['scorer'])
                )
                if score >= int(map_row['threshold']):
                    select_list.append(col('`' + match + '`').alias(map_row['destname']))
                    unmapped_fields.remove(match)
                    # Add match to mutable dictionary object to use for lineage
                    map_row['match'] = match

                print(f"Fuzzy matched {map_row['sourcename']} with column {match} and score {score}")

    if unmapped_fields:
        print(f'Discarded unmapped fields {unmapped_fields}')

    lineage.update_lineage(df, args['source_key'], 'mapping', map=field_mapping_list)
    return df.select(select_list)


def custommapping_with_glue(dyf: DynamicFrame, field_mapping_list: list, args: list, lineage) -> DynamicFrame:
    """Apply a custom field mapping to the data schema in a Glue DynamicFrame
    """
    # NOTE: Fuzzy matching mappings will be skipped
    prepared_map = [ (map_row['sourcename'], map_row['destname'])
        for map_row in field_mapping_list
        if not map_row['threshold'] and map_row['destname'].lower() != 'null' ]

    # NOTE: ApplyMapping.apply "optimizes" DecimalTypes, reorders fields, and re-samples columns
    #   resulting in Null/Void column types
    # NOTE: Discarded unmapped fields will not be logged
    lineage.update_lineage(dyf, args['source_key'], 'mapping', map=field_mapping_list)
    return ApplyMapping.apply(
        info='Field name mapping transform',
        frame=dyf,
        mappings=prepared_map,
        transformation_ctx=f"{args['execution_id']}-custommapping",
    )
lib/glue_scripts/lib/custom_mapping.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
"""
InsuranceLake implementation of Data Lineage generation class. This class handles different operations
on a dataset and registers them in a DynamoDB. There is an incremental sequence number added 
to the records in the order the operations occur.

"""
import uuid
import json
import boto3
from datetime import datetime
import dateutil.tz
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.functions import sum, input_file_name

class DataLineageGenerator:
    """The InsuranceLake custom data lineage class

    Methods: generateuuid
             insert_table
             update_lineage
    """
    def __init__(self, args: dict, uniid: str = None):
        """
        Parameters
        ----------
        args
            Glue job parameters; required values: JOB_NAME, JOB_RUN_ID, state_machine_name, execution_id, data_lineage_table
        uniid: optional
            Optional UUID to use for lineage events; creates a new UUID if not supplied
        """
        if 'data_lineage_table' in args:
            dynamodb = boto3.resource('dynamodb')

            self.table = dynamodb.Table(args['data_lineage_table'])
            self.job_id = args['JOB_RUN_ID']
            self.job_name  = args['JOB_NAME']
            self.step_function_name = args['state_machine_name']
            self.step_function_execution_id = args['execution_id'] 
            self.uniid = uniid if uniid is not None else self.generateuuid()
            self.count = 1
        else:
            # Allow the calling code to toggle lineage on/off by specifying a table (or not)
            self.table = None

    def generateuuid(self) -> str:
        """Generates a new UUID
        """
        return(str(uuid.uuid4()))

    def insert_table(self, dataset: str, operation: str, transform_spec: any):
        """Inserts the lineage event to DynamoDB data lineage table

        Parameters
        ----------
        dataset
            Name of the dataset on which the operation has been performed
        operation
            Information about the type of operation; e.g. read, write, numericaudit, or other transform name
        transform_spec
            Detailed information about the operation
        """
        now = datetime.now(tz=dateutil.tz.gettz('UTC'))

        self.table.put_item(
            Item = {
                'step_function_execution_id' : self.step_function_execution_id,
                'job_id_operation_seq': f'{self.job_id}-{self.count}',
                'operation_seq' : self.count,
                'job_id': self.job_id,
                'job_name': self.job_name,
                'dataset' : dataset,
                'uuid' : self.uniid,
                'time_of_operation' : str(now.strftime("%Y-%m-%d %H:%M:%S")),
                'dataset_operation' : operation,
                'step_function_name' : self.step_function_name,
                'info' : json.dumps(transform_spec)
            }
        )
        self.count += 1

    def update_lineage(self, df: DataFrame, dataset: str, operation: str, **kwargs):
        """Prepares the event data to be recorded in DynamoDB data lineage table

        Parameters
        ----------
        df
            Spark DataFrame to use for lineage data
        dataset
            Name of the dataset on which the operation has been performed
        operation
            Information about the type of operation; e.g. read, write, numericaudit, or other transform name
        kwargs: optional
            List of optional keywords used for specific operations; most operations expect 'transform'
        """
        if not self.table:
            # No table is defined so no data lineage is possible
            return

        dataset_counts = {
            'row': df.count(),
            'column': len(df.columns),
        }
        try:
            dataset_filename = df.select(input_file_name()).limit(1).collect()[0][0]
        except:
            # input_file_name() does not support more than one source, some source types do not capture input_file_name
            dataset_filename = 'N/A'

        # Read / Write
        if operation == 'read':
            lineage_info = {
                'datasetname': dataset_filename,
                'dataset_counts': dataset_counts,
                'schema': str(df.dtypes),
            }
            self.insert_table(dataset, operation, lineage_info)

        elif operation == 'write':
            lineage_info = {
                'datasetname': dataset_filename,
                'format': kwargs['format'],
                'dataset_counts': dataset_counts,
                'schema': str(df.dtypes),
            }
            self.insert_table(dataset, operation, lineage_info)

        # Numeric Audit
        elif operation == 'numericaudit': 
            lineage_info = {
                'dataset_counts': dataset_counts,
                'field_total': {
                    # Sum any double or decimal columns; convert to string explicitly to not rely on the JSON parser
                    field[0]: str(df.select(sum(f'`{field[0]}`')).collect()[0][0])
                        for field in df.dtypes if field[1] == 'double' or field[1][:7] == 'decimal'
                }
            }
            self.insert_table(dataset, operation, lineage_info)

        # Mapping
        elif operation == 'mapping':
            for map_record in kwargs['map']:
                if map_record['destname'].lower() == 'null':
                    self.insert_table(dataset, 'dropped', { 'dropped': map_record['sourcename'] })
                elif map_record.get('threshold', None):
                    self.insert_table(dataset, 'fuzzymapping', map_record)
                else:
                    self.insert_table(dataset, operation, map_record)

        # All other transforms (that do not require special handling)
        else:
            for spec in kwargs['transform']:
                self.insert_table(dataset, operation, { operation: spec })
lib/glue_scripts/lib/datalineage.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
from typing import Callable
from awsglue.context import GlueContext
from awsglue.dynamicframe import DynamicFrame
from awsglue.transforms import SelectFromCollection
from pyspark.sql.dataframe import DataFrame
from pyspark.context import SparkContext
from pyspark.sql.functions import col, current_timestamp
from awsgluedq.transforms import EvaluateDataQuality
from glue_catalog_helpers import upsert_catalog_table


def create_add_dynamodb_columns_func(rule_action: str, args: dict) -> Callable:
    """Create and return DynamicFrame map function to add columns to DQ results
    Use second order function because DynamicFrame.map/RDD.map have no access to global scope
    """
    def _function(rec):
        rec['execution_id'] = args['execution_id']
        rec['job_id_action_rule'] = args['JOB_RUN_ID'] + '-' + rule_action + '-' + rec['Rule']
        rec['rule_action'] = rule_action
        rec['job_id'] = args['JOB_RUN_ID']
        rec['job_name'] = args['JOB_NAME']
        rec['step_function_name'] = args['state_machine_name']
        # DynamoDB Athena connector doesn't handle Null values well
        rec['FailureReason'] = rec['FailureReason'] or ''
        return rec

    return _function

def write_results_to_dynamodb(results_dyf: DynamicFrame, rule_action: str, args: dict, sc: SparkContext):
    """Prepare and write Glue Data Quality results to DynamoDB table

    Parameters
    ----------
    dyf
        Glue DynamicFrame on which to check data quality
    action_name
        abort/warn/quarantine action name to associate with rule results
    args
        Glue job arguments, from which dq_results_table is used
    sc
        Spark Context class from the calling job
    """
    glueContext = GlueContext(sc)

    add_dynamodb_columns = create_add_dynamodb_columns_func(rule_action, args)
    results_with_keys_dyf = results_dyf.map(f = add_dynamodb_columns)

    glueContext.write_dynamic_frame_from_options(
        frame = results_with_keys_dyf,
        connection_type = 'dynamodb',
        connection_options = {
            'dynamodb.output.tableName': args['dq_results_table'],
        }
    )


def check_dataquality_halt(dyf: DynamicFrame, dq_rules: list, args: dict, sc: SparkContext):
    """Evalute Glue Data Quality rules and raise runtime error if any fail

    Parameters
    ----------
    dyf
        Glue DynamicFrame on which to check data quality
    dq_rules
        List of strings containing Glue Data Quality rules (single quotes around field names supported)
    args
        Glue job arguments, from which source_key is used
    sc
        Spark Context class from the calling job
    """
    rule_string = 'Rules = [ ' + ','.join(dq_rules).replace("'", '"') + ' ]'
    dq_results = EvaluateDataQuality.apply(
        frame=dyf,
        ruleset=rule_string,
        publishing_options={
            'dataQualityEvaluationContext': args['source_key'].replace('/', '-') + '-halt-on-failure',
            'enableDataQualityCloudWatchMetrics': True,
            'enableDataQualityResultsPublishing': True,
        },
    )
    write_results_to_dynamodb(dq_results, 'halt-on-failure', args, sc)
    dq_results_df = dq_results.toDF()
    dq_results_df.show(truncate=False)
    failed_df = dq_results_df.filter(col('Outcome').contains('Failed'))
    if failed_df.count() > 0:
        raise RuntimeError('Data quality check failed')


def check_dataquality_warn(dyf: DynamicFrame, dq_rules: list, args: dict, sc: SparkContext):
    """Evalute Glue Data Quality rules and log a warning message if any fail

    Parameters
    ----------
    dyf
        Glue DynamicFrame on which to check data quality
    dq_rules
        List of strings containing Glue Data Quality rules (single quotes around field names supported)
    args
        Glue job arguments, from which source_key is used
    sc
        Spark Context class from the calling job
    """
    rule_string = 'Rules = [ ' + ','.join(dq_rules).replace("'", '"') + ' ]'
    dq_results = EvaluateDataQuality.apply(
        frame=dyf,
        ruleset=rule_string,
        publishing_options={
            'dataQualityEvaluationContext': args['source_key'].replace('/', '-') + '-warn-on-failure',
            'enableDataQualityCloudWatchMetrics': True,
            'enableDataQualityResultsPublishing': True,
        },
    )
    write_results_to_dynamodb(dq_results, 'warn-on-failure', args, sc)
    dq_results_df = dq_results.toDF()
    dq_results_df.show(truncate=False)
    warn_df = dq_results_df.filter(col('Outcome').contains('Failed'))
    if warn_df.count() > 0:
        print('Data quality checks failed: warning only')


def check_dataquality_quarantine(dyf: DynamicFrame, dq_rules: list, args: dict, sc: SparkContext) -> DataFrame:
    """Evalute Glue Data Quality rules and filter/quarantine rows that fail
    Reference: https://docs.aws.amazon.com/glue/latest/ug/data-quality-notebooks.html

    Parameters
    ----------
    dyf
        Glue DynamicFrame on which to check data quality
    dq_rules
        List of strings containing Glue Data Quality rules (single quotes around field names supported)
    args
        Glue job arguments, from which source_key is used
    sc
        Spark Context class from the calling job

    Returns
    -------
    DataFrame
        DataFrame (not DynamicFrame) with failed rows filtered (i.e. only passed rows)
    """
    glueContext = GlueContext(sc)
    spark = glueContext.spark_session

    rule_string = 'Rules = [ ' + ','.join(dq_rules).replace("'", '"') + ' ]'
    dq_multiframe = EvaluateDataQuality().process_rows(
        frame=dyf,
        ruleset=rule_string,
        publishing_options={
            'dataQualityEvaluationContext': args['source_key'].replace('/', '-') + '-quarantine-on-failure',
            'enableDataQualityCloudWatchMetrics': True,
            'enableDataQualityResultsPublishing': True,
        },
        additional_options={'performanceTuning.caching': 'CACHE_NOTHING'},
    )

    dq_results = SelectFromCollection.apply(dfc=dq_multiframe, key='ruleOutcomes')
    write_results_to_dynamodb(dq_results, 'quarantine-on-failure', args, sc)
    dq_results_df = dq_results.toDF()
    dq_results_df.show(truncate=False)

    dq_rowleveloutcomes = SelectFromCollection.apply(dfc=dq_multiframe, key='rowLevelOutcomes')
    dq_rowleveloutcomes_df = dq_rowleveloutcomes.toDF()

    # Passed records will be returned as DataFrame
    passed_df = dq_rowleveloutcomes_df.filter(col('DataQualityEvaluationResult').contains('Passed'))
    # Remove DQ outcome columns because there are only passed rows in this DataFrame
    passed_df = passed_df.drop('DataQualityRulesPass', 'DataQualityRulesFail', 'DataQualityRulesSkip', 'DataQualityEvaluationResult')

    # Failed records will be quarantined
    failed_df = dq_rowleveloutcomes_df.filter(col('DataQualityEvaluationResult').contains('Failed'))
    if failed_df.count() > 0:
        failed_df = failed_df.withColumn('quarantine_timestamp', current_timestamp())

        storage_location = args['target_bucket'] + '/quarantine/' + args['target_database_name'] + '/' + args['table_name']
        upsert_catalog_table(
            failed_df,
            args['target_database_name'],
            args['table_name'] + '_quarantine',
            ['year', 'month', 'day'],
            storage_location,
            # Allow schema changes, because we don't want a failure when saving quarantined rows
            allow_schema_change='permissive',
        )
        failed_df.write.partitionBy('year', 'month', 'day').format('parquet').save(storage_location, mode='overwrite')
        spark.sql(f"ALTER TABLE {args['target_database_name']}.{args['table_name'] + '_quarantine'} RECOVER PARTITIONS")

        if passed_df.count() == 0:
            raise RuntimeError('Data quality check quarantined all rows')

        print(f'Data quality checks failed: {failed_df.count()} record(s) quarantined to {storage_location}')

    # Return passed rows as new DataFrame
    return passed_df
lib/glue_scripts/lib/dataquality_check.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import hashlib
import pyspark
from pyspark.context import SparkContext
from pyspark.sql.functions import udf, lit
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.types import StructType, StructField, StringType
from awsglue.context import GlueContext
from awsglue.dynamicframe import DynamicFrame

def transform_hash(df: DataFrame, hash_fields: list, args: dict, lineage, *extra) -> DataFrame:
    """Hash specified column values using SHA256 and Spark UDF

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply field hashing
    hash_fields
        List of field names to hash
    args
        Glue job arguments, from which source_key is used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with hashing applied
    """
    tohash = udf(lambda x: hashlib.sha256(str(x).encode('utf-8')).hexdigest())

    cols_map = {}
    for hash_field in hash_fields:
        cols_map.update({ hash_field: tohash(hash_field) })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'hashing', transform=hash_fields)
    return df.withColumns(cols_map)


def transform_redact(df: DataFrame, redact_fields: dict, args: dict, lineage, *extra) -> DataFrame:
    """Redact specified column values using supplied redaction string

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply field redaction
    redact_fields
        Dictionary of fields to redact in the form:
            field: redact_string
    args
        Glue job arguments, from which source_key is used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with redaction applied
    """
    cols_map = {}
    for redact_field, redact_string in redact_fields.items():
        cols_map.update({ redact_field: lit(redact_string) })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'redact', transform=redact_fields)
    return df


def transform_tokenize(df_with_values: DataFrame, tokenize_fields: list, args: dict, lineage, sc: SparkContext, *extra) -> DataFrame:
    """Replace specified column values with hash and store original value in separate table
    Uses SHA256 and Spark UDF similar to transform_hash() but with original value stored

    Parameters
    ----------
    df_with_values
        pySpark DataFrame on which to apply field hashing
    tokenize_fields
        List of fields to tokenize
    args
        Glue job arguments, from which source_key, execution_id, and hash_value_table are used
    lineage:
        Initialized lineage class object from the calling job
    sc
        Spark Context class from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with hashing applied
    """
    # TODO: Reduce number of dataframes used with select() and column aliases
    # TODO: Support hash with salt or composite hash value (improves uniqueness of token data)

    tohash = udf(lambda x: hashlib.sha256(str(x).encode('utf-8')).hexdigest())

    glueContext = GlueContext(sc)
    spark = glueContext.spark_session

    # Columns needed for the token store table
    token_store_columns = StructType([
        StructField('raw_data', StringType(), False),
        StructField('hash_key', StringType(), False),
    ])

    # Create an empty DataFrame for the token store
    empty_rdd = sc.emptyRDD()
    token_store_df = spark.createDataFrame(
        data = empty_rdd,
        schema = token_store_columns
    )

    # Using a different variable name here ensures that a pointer to
    # the DataFrame with original values persists and can be referenced.
    # Pointers to the DataFrame with hashes will be repeatedly 
    # de-referenced in the loop as the variable is replaced with a new
    # DataFrame after each operation.
    df_with_hashes = df_with_values

    for tokenize_field in tokenize_fields:

        # Hash specified field and store value in a new column 'hash_key'
        df_with_hashes = df_with_hashes.withColumn('hash_key', tohash(tokenize_field))

        # Select the tokenized column and the hash_key column into a
        # temp DataFrame (to merge with the token store table)
        df_token_store_to_merge = df_with_hashes.select(tokenize_field, 'hash_key')

        # In the temp DataFrame, rename the tokenized column to 'raw_data'
        df_token_store_to_merge.withColumnRenamed(tokenize_field, 'raw_data')

        # Merge these new original values with the existing token_store DataFrame
        token_store_df = token_store_df.union(df_token_store_to_merge)

        # We are done with this temp DataFrame
        df_token_store_to_merge.unpersist()

        # Drop the original values from the DataFrame that is being transformed
        df_with_hashes = df_with_hashes.drop(tokenize_field)

        # Rename 'hash_key' to the tokenized column name
        # Drop and Rename is faster than tokenizing twice
        df_with_hashes = df_with_hashes.withColumnRenamed('hash_key', tokenize_field)

    # Remove all duplicate values to reduce write operations and because
    # hash_key (which will be duplicated too) is a partition key in the
    # DynamoDB table
    token_store_df = token_store_df.dropDuplicates(['raw_data'])

    token_store_dyf = DynamicFrame.fromDF(token_store_df, glueContext, f"{args['execution_id']}-token_store_to_write")

    # Store the new original values
    glueContext.write_dynamic_frame_from_options(
        frame = token_store_dyf,
        connection_type = 'dynamodb',
        connection_options = {
            'dynamodb.output.tableName': args['hash_value_table'],
        },
        transformation_ctx = f"{args['execution_id']}-store_token_table"
    )

    lineage.update_lineage(df_with_hashes, args['source_key'], 'tokenize', transform=tokenize_fields)

    # Return the transformed DataFrame, no other dataframes are needed
    return df_with_hashes
lib/glue_scripts/lib/datatransform_dataprotection.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pyspark
import json
import boto3
from boto3.dynamodb.conditions import Key
from pyspark.context import SparkContext
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.functions import coalesce, lit, broadcast, concat_ws
from awsglue.context import GlueContext

def get_lookup_data(table_name: str, key_values: dict, lookup_data_attribute_name: str) -> list:
    """Gets a single data item from a DynamoDB table based on key values
    key_values must contain all keys needed to represent a primary key

    Parameters
    ----------
    table_name
        DynamoDB table from which to get item
    key_values
        Dictionary of key values to identify item
    lookup_data_attribute_name
        Attribute name to use for lookup_data (will be json parsed)

    Returns
    -------
    list
        List of lookup keys and values suitable for creating a DataFrame
    """
    dynamodb = boto3.resource('dynamodb')
    lookup_table = dynamodb.Table(table_name)

    response = lookup_table.get_item(
            Key = key_values
        )

    if 'Item' not in response:
        raise RuntimeError(f'No lookup data found for key values {key_values} on table {table_name}')

    lookup_json = json.loads(response['Item'][lookup_data_attribute_name])
    # Format the data so it can be used to create a dataframe
    return [ ( key, lookup_json[key] ) for key in lookup_json.keys() ]

def transform_lookup(df_for_lookup: DataFrame, lookup_spec: list, args: dict, lineage, sc: SparkContext, *extra) -> DataFrame:
    """Replace specified column values with values looked up from an external table

    Parameters
    ----------
    df_for_lookup
        pySpark DataFrame on which to perform the lookups and replacement
    lookup_spec
        List of dictionary objects in the form:
            field: 'NewFieldName' (must be a new field if source is specified)
            source: 'OriginalFieldName' (optional, field will be replaced if ommitted)
            lookup: 'LookupData' (DynamoDB column_name value to use to get looked up values')
            nomatch: 'N/A' (optional, no matches will be empty/null)
    args
        Glue job arguments, from which source_key, value_lookup_table, and target_database_name are used
    lineage
        Initialized lineage class object from the calling job
    sc
        Spark Context class from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with lookup applied
    """
    glueContext = GlueContext(sc)
    spark = glueContext.spark_session

    # Pointers to the DataFrame df_for_lookup will be repeatedly dereferenced in the loop (and
    # freed from memory when Spark is able) as the variable is replaced with a new DataFrame 
    # with each operation.
    for spec in lookup_spec:

        # Get lookup data from DynamoDB
        lookup_data = get_lookup_data(
                args['value_lookup_table'],
                {
                    'source_system': args['target_database_name'],
                    'column_name': spec['lookup'],
                },
                'lookup_data'
            )

        # spec will indicate a source field if the looked up value should be added as a new column
        sourcefield = spec.get('source', spec['field'])

        # Create a DataFrame with lookup values to join to
        df_for_join = spark.createDataFrame(lookup_data, schema=['orig_value', 'new_value'])

        # Left join so that no source data is lost
        # After join, drop the now unneeded orig_value column
        df_for_lookup = df_for_lookup.join(
                broadcast(df_for_join),
                df_for_lookup[sourcefield] == df_for_join['orig_value'],
                'left'
            ). \
            drop('orig_value')
        df_for_join.unpersist()

        if 'source' not in spec:
            # Drop the original value, because the spec indicates replacing the existing column
            df_for_lookup = df_for_lookup.drop(spec['field'])

        # Rename the new_value column
        df_for_lookup = df_for_lookup.withColumnRenamed('new_value', spec['field'])

        if 'nomatch' in spec:
            # Fill default literal value for unmatched
            df_for_lookup = df_for_lookup.fillna(spec['nomatch'], spec['field'])

        lineage.update_lineage(df_for_lookup, args['source_key'], 'lookup', transform=spec)

    return df_for_lookup


def get_multilookup_data(table_name: str, lookup_group: str, lookup_data_attributes: list, limit: int = None) -> list:
    """Gets data items for multilookup from a DynamoDB table based on lookup_group

    Parameters
    ----------
    table_name
        DynamoDB table from which to get item
    lookup_group
        Filter string to query DynamoDB table lookup results
    lookup_data_attributes
        List of attribute names to retrieve; lookup_item is always retrieved
    limit: optional
        Optional pagination size primarily used for testing

    Returns
    -------
    list
        List of lookup items from DynamoDB suitable for creating a DataFrame
    """
    dynamodb = boto3.resource('dynamodb')
    table = dynamodb.Table(table_name)
    query_parameters = {
        'KeyConditionExpression': ( Key('lookup_group').eq(lookup_group) ),
        'ProjectionExpression': ','.join(['lookup_item'] + lookup_data_attributes)
    }
    if limit:
        query_parameters.update({ 'Limit': limit})

    items = []
    response = table.query(**query_parameters)
    while True:
        items += response['Items']
        if 'LastEvaluatedKey' in response:
            # Results are paginated due to size over 1 MB
            response = table.query(ExclusiveStartKey=response['LastEvaluatedKey'], **query_parameters)
        else:
            break

    if not items:
        raise RuntimeError(f'No lookup data found for lookup_group {lookup_group}')
    return items

def transform_multilookup(df_for_lookup: DataFrame, lookup_spec: list, args: dict, lineage, sc: SparkContext, *extra) -> DataFrame:
    """Add specified multiple values looked up from an external table as new columns

    Parameters
    ----------
    df_for_lookup
        pySpark DataFrame on which to perform the lookups and replacement
    lookup_spec
        List of dictionary objects in the form:
            field: 'NewFieldName' (must be a new field if source is specified)
            lookup: 'LookupData' (DynamoDB column_name value for filter to get looked up values')
            nomatch: 'N/A' (optional, no matches will be empty/null)
    args
        Glue job arguments, from which source_key, multi_lookup_table, and target_database_name are used
    lineage
        Initialized lineage class object from the calling job
    sc
        Spark Context class from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with multi lookup applied
    """
    glueContext = GlueContext(sc)
    spark = glueContext.spark_session

    # Pointers to the DataFrame df_for_lookup will be repeatedly dereferenced in the loop (and
    # freed from memory when Spark is able) as the variable is replaced with a new DataFrame
    # with each operation.
    for spec in lookup_spec:

        lookup_items = get_multilookup_data(
            table_name=args['multi_lookup_table'],
            lookup_group=spec['lookup_group'],
            lookup_data_attributes=spec['return_attributes'],
        )

        # Create a DataFrame with lookup values to join to
        df_for_join = spark.createDataFrame(lookup_items)

        # Left join so that no source data is lost
        # After join, drop the now unneeded lookup_item column
        df_for_lookup = df_for_lookup.join(
                broadcast(df_for_join),
                concat_ws('-', *spec['match_columns']) == df_for_join['lookup_item'],
                'left'
            ). \
             drop('lookup_item')

        df_for_join.unpersist()

        if 'nomatch' in spec:
            # Fill default literal value for unmatched
            fill_map = { column_name: spec['nomatch'] for column_name in spec['return_attributes'] }
            df_for_lookup = df_for_lookup.fillna(fill_map)

        lineage.update_lineage(df_for_lookup, args['source_key'], 'multilookup', transform=spec)

    return df_for_lookup


def transform_merge(df: DataFrame, merge_spec: list, args: dict, lineage, *extra) -> DataFrame:
    """Merge columns using coalesce

    Parameters
    ----------
    merge_spec
        List of dictionary objects in the form:
            field: 'NewFieldName'
            source_list: [ 'first', 'second', 'third' ] list of fields in order of priority to coalesce
            default: 'N/A' optional lowest priority default value to use in the coalesce (instead of null)
    """
    cols_map = {}
    for spec in merge_spec:
        column_list = spec['source_list'].copy()
        if 'default' in spec:
            # Append the default value as a literal to the list of fields
            column_list.append(lit(spec['default']))

        cols_map.update({ spec['field']: coalesce(*column_list) })
        lineage.update_lineage(df, args['source_key'], 'merge', transform=spec)

    return df.withColumns(cols_map)
lib/glue_scripts/lib/datatransform_lookup.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
from decimal import Decimal
import datetime
from dateutil import relativedelta, rrule
import calendar
from functools import reduce
from operator import add
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.functions import udf, expr, months_between, explode_outer, coalesce, lit, col
from pyspark.sql.types import DecimalType, DateType, ArrayType, IntegerType

def add_columns(*source_columns):
    """Add values from an arbitrary number of Spark columns together and return the result (as a column)
    Uses only Spark native functions for performance (contributed by Gonzalo Herreros <gonher@amazon.com>)
    """
    return reduce(add, [ coalesce(col(c), lit(0)) for c in source_columns ])


def policy_month_list(policy_effective_date: datetime.date, policy_expiration_date: datetime.date):
    """Return list of months from the first month of the policy to the last
    """
    first_day_of_month = datetime.date(policy_effective_date.year, policy_effective_date.month, 1)
    return list(rrule.rrule(freq=rrule.MONTHLY, dtstart=first_day_of_month, until=policy_expiration_date))


def months_between_normalized(expiration_date: datetime.date, effective_date: datetime.date) -> any:
    """Calculate the number of months between 2 dates with floor-style rounding to nearest whole month
    Designed to work as regular Python function and Spark UDF
    """
    if None in [ expiration_date, effective_date ]:
        return None

    # Columns of DateType are automatically converted to Python datetime.date
    # Achieve a floor-style rounding of number of months by subtracting 1 month from the end date
    expiration_prior_month = expiration_date - relativedelta.relativedelta(months=1)
    if expiration_prior_month < effective_date:
        # Except when policy length is one month or less
        expiration_prior_month = expiration_date

    # Using a list of months will ensure an integer number of months
    return len(policy_month_list(effective_date, expiration_prior_month))


def transform_enddate(df: DataFrame, enddate_spec: list, args: dict, lineage, *extra) -> DataFrame:
    """Add a number of months to a specified date to get an ending/expiration date

    Parameters
    ----------
    policymonths_spec
        A list of dictionary objects in the form:
            field: "NewFieldName"
            start_date: field name containing the start date or policy effective date
            num_months: field name containing the number of months to add to the start date
                Column must be integer values
    """
    cols_map = {}
    for spec in enddate_spec:
        # "Column is not iterable" workaround:
        # https://sparkbyexamples.com/pyspark/pyspark-typeerror-column-is-not-iterable/
        cols_map.update({spec['field']: expr(
            'add_months(' + spec['start_date'] + ',' + spec['num_months'] + ')'
        )})

    lineage.update_lineage(df, args['source_key'], 'enddate', transform=enddate_spec)
    return df.withColumns(cols_map)


def transform_policymonths(df: DataFrame, policymonths_spec: list, args: dict, lineage, *extra) -> DataFrame:
    """Calculate number of months between policy start/end dates

    Parameters
    ----------
    policymonths_spec
        A list of dictionary objects in the form:
            field: "NewFieldName"
            policy_effective_date: field name containing policy effective date
            policy_expiration_date: field name containing policy expiration date
            normalized: (optional) boolean if number of months should be floor-rounded
    """
    cols_map = {}
    for spec in policymonths_spec:
        dates = [ spec['policy_expiration_date'], spec['policy_effective_date'] ]

        if 'normalized' in spec and spec['normalized']:
            months_function = udf(months_between_normalized, IntegerType())
            cols_map.update({ spec['field']: months_function(*dates) })
        else:
            cols_map.update({ spec['field']: months_between(*dates).cast('Decimal(16,2)') })

    lineage.update_lineage(df, args['source_key'], 'policymonths', transform=policymonths_spec)
    return df.withColumns(cols_map)


@udf(ArrayType(DateType()))
def add_policy_months_list(policy_effective_date, policy_expiration_date):
    """Return an array column containing a list of active policy months for a policy
    Intended to be used with Spark explode
    """
    if None not in [ policy_effective_date, policy_expiration_date ]:
        # ArrayType is equivalent to Python list
        return policy_month_list(policy_effective_date, policy_expiration_date)

@udf(DateType())
def last_day_of_month(any_date_in_month):
    """Given a date, return the last day of the month for that date
    """
    if any_date_in_month:
        return datetime.date(any_date_in_month.year, any_date_in_month.month, 
            calendar.monthrange(any_date_in_month.year, any_date_in_month.month)[1]
        )

def transform_expandpolicymonths(df: DataFrame, expandpremium_spec: dict, args: dict, lineage, *extra) -> DataFrame:
    """Expand dataset to one row for each month the policy is active with a calculated earned premium

    Parameters
    ----------
    earnedpremium_spec
        A dictionary object in the form:
            uniqueid: (optional) Name of field to add with generated unique identifier (uuid)
			policy_month_start_field: Name of field to store the policy month
			policy_month_end_field: Name of field to store the policy month
            policy_effective_date: field name containing policy effective date
            policy_expiration_date: field name containing policy expiration date
    """
    # If specified, add a unique ID for each policy before exploding the rows
    if 'uniqueid' in expandpremium_spec:
        df = df.withColumn(expandpremium_spec['uniqueid'], expr('uuid()'))

    # Create array column with list of policy months (first day of the month)
    df = df.withColumn(expandpremium_spec['policy_month_start_field'], 
            add_policy_months_list(
                expandpremium_spec['policy_effective_date'],
                expandpremium_spec['policy_expiration_date']
        ))

    # Explode array column into multiple rows
    df = df.withColumn(expandpremium_spec['policy_month_start_field'], 
        explode_outer(expandpremium_spec['policy_month_start_field']))

    # Add column for last day of month, so we have both start and end date of period
    df = df.withColumn(expandpremium_spec['policy_month_end_field'],
            last_day_of_month(expandpremium_spec['policy_month_start_field']))

    lineage.update_lineage(df, args['source_key'], 'expandpolicymonths', transform=expandpremium_spec)
    return df


@udf(DecimalType(16,2))
def earnedpremium_straightline(total_premium, effective_date, expiration_date, period_start_date, period_end_date):
    """Calculate Earned Premium for a given policy and period dates using a straighline method
    """
    if None in [ total_premium, effective_date, expiration_date, period_start_date, period_end_date ]:
        return None

    month_list = policy_month_list(effective_date, expiration_date)
    # Check for empty list due to effective_date >= expiration_date (i.e. bad data)
    # Use last month in normalized month list instead of expiration_date to provide consistent results
    if not month_list or \
            period_end_date < effective_date or period_end_date > month_list[-1].date():
        # Last month of the policy will have $0, first month will have a full earned premium
        # expiration_date must be on or earlier than effective_date (prevents division by 0)
        return None

    return total_premium / months_between_normalized(expiration_date, effective_date)


@udf(DecimalType(16,2))
def earnedpremium_byday(total_premium, effective_date, expiration_date, period_start_date, period_end_date):
    """Calculate Earned Premium for a given policy and period dates using number of days in period
    """
    if None in [ total_premium, effective_date, expiration_date, period_start_date, period_end_date ]:
        return None

    if period_end_date < effective_date or period_start_date > expiration_date or \
            (expiration_date - effective_date).days < 0:
        # Specified period is outside the policy dates
        # expiration_date must be on or earlier than effective_date (prevents division by 0)
        return None

    start_date = effective_date if period_start_date < effective_date else period_start_date
    end_date = expiration_date if period_end_date > expiration_date else period_end_date
    return total_premium * \
        Decimal(
            ((end_date - start_date).days + 1)
            /
            ((expiration_date - effective_date).days + 1)
        )

def transform_earnedpremium(df: DataFrame, earnedpremium_spec: list, args: dict, lineage, *extra) -> DataFrame:
    """Calculate monthly earned premium

    Parameters
    ----------
    earnedpremium_spec
        A list of dictionary objects in the form:
            field: "NewFieldName"
            written_premium_list: [ "first", "second", "third" ] list of fields containing written premium to be summed
            period_start_date: field name containing period start date
            period_end_date: field name containing period end date
            policy_effective_date: field name containing policy effective date
            policy_expiration_date: field name containing policy expiration date
            byday: (optional) boolean to determine method of calculation (by month vs. by day)
    """
    cols_map = {}
    for spec in earnedpremium_spec:
        if 'byday' in spec and spec['byday']:
            cols_map.update({spec['field']: earnedpremium_byday(
                add_columns(*spec['written_premium_list']),
                spec['policy_effective_date'],
                spec['policy_expiration_date'],
                spec['period_start_date'],
                spec['period_end_date'],
            )})
        else:
            cols_map.update({spec['field']: earnedpremium_straightline(
                add_columns(*spec['written_premium_list']),
                spec['policy_effective_date'],
                spec['policy_expiration_date'],
                spec['period_start_date'],
                spec['period_end_date'],
            )})

    lineage.update_lineage(df, args['source_key'], 'earnedpremium', transform=earnedpremium_spec)
    return df.withColumns(cols_map)


def transform_addcolumns(df: DataFrame, addcolumns_spec: dict, args: dict, lineage, *extra) -> DataFrame:
    """Add two or more columns together in a new column

    Parameters
    ----------
    addcolumns_spec
        A list of dictionary objects in the form:
            field: "NewFieldName"
            source_columns: [ "first", "second", "third" ] list of fields to add together
    """
    cols_map = {}
    for spec in addcolumns_spec:
        cols_map.update({ spec['field']: add_columns(*spec['source_columns']) })

    lineage.update_lineage(df, args['source_key'], 'addcolumns', transform=addcolumns_spec)
    return df.withColumns(cols_map)


def transform_flipsign(df: DataFrame, field_list: list, args: dict, lineage, *extra) -> DataFrame:
    """Flip the sign of a numeric column in a Spark DataFrame, optionally in a new column

    Parameters
    ----------
    field_list
        List of dictionary objects in the form:
            field: "NeworCurrentFieldName"
            source: "SourceFieldName" (optional, if omitted, field will be changed in place)
    """
    cols_map = {}
    for spec in field_list:
        sourcefield = spec.get('source', spec['field'])
        cols_map.update({ spec['field']: - df[sourcefield] })

    lineage.update_lineage(df, args['source_key'], 'flipsign', transform=field_list)
    return df.withColumns(cols_map)
lib/glue_scripts/lib/datatransform_premium.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pyspark
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.types import StringType
from pyspark.sql.functions import lit, regexp_extract, udf
import re

def transform_filename(df: DataFrame, filename_patterns: list, args: dict, lineage, *extra):
    """Add column to DataFrame based on regexp pattern on the filename argument to the Glue job

    Parameters
    ----------
    df
        pySpark DataFrame on which to add the column
    filename_patterns
        List of patterns and fieldnames to add:
            pattern: 'regex-with-grouping'
            field: 'FieldName'
            required: boolean
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with filename transforms applied
    """
    cols_map = {}
    for spec in filename_patterns:
        filename_re = re.compile(spec['pattern'])
        match = filename_re.search(args['base_file_name'])
        if match:
            cols_map.update({ spec['field']: lit(match.group(1)) })
        else:
            if spec['required']:
                # Pattern is required, so stop workflow
                raise RuntimeError(f"Filename '{args['base_file_name']}' failed to match required pattern {spec['pattern']} for field '{spec['field']}'")
            else:
                # Pattern is not required, so add empty column as placeholder
                cols_map.update({ spec['field']: lit(None).cast(StringType()) })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'filename', transform=filename_patterns)
    return df.withColumns(cols_map)


def transform_columnfromcolumn(df: DataFrame, cfc_spec: list, args: dict, lineage, *extra):
    """Add column to DataFrame based on regexp pattern on another column

    Parameters
    ----------
    cfc_spec
        List of patterns and fieldnames to add:
            field: 'FieldName'
            source: 'FieldName' (optional, field will be replaced if ommitted)
            pattern: 'regex-with-grouping'
    """
    cols_map = {}
    for spec in cfc_spec:
        # spec will indicate a source field if the matched value should be added as a new column
        sourcefield = spec.get('source', spec['field'])

        cols_map.update({
            spec['field']:
            regexp_extract(sourcefield, spec['pattern'], 1)
        })

    lineage.update_lineage(df, args['source_key'], 'columnfromcolumn', transform=cfc_spec)
    return df.withColumns(cols_map)


def transform_literal(df: DataFrame, spec: dict, args: dict, lineage, *extra):
    """Add column to DataFrame with static/literal value supplied in specification

    Parameters
    ----------
    spec
        Dictionary of field names and values
    """
    cols_map = {}
    for field_name, field_value in spec.items():
        cols_map.update({ field_name: lit(field_value) })

    lineage.update_lineage(df, args['source_key'], 'literal', transform=spec)
    return df.withColumns(cols_map)


@udf(StringType())
def format_column(formatstring, *args):
    """Format list of column values using a Python format string; convert None values to
    empty strings and attempt to clean up dangling whitespace separators
    """
    if formatstring:
        checked_args = [ '' if value is None else value for value in args ]
        return formatstring.format(*checked_args).strip()

def transform_combinecolumns(df: DataFrame, combine_spec: list, args: dict, lineage, *extra):
    """Add column to DataFrame using format string and source columns

    Parameters
    ----------
    combine_spec
        List of dictionary objects in the form:
            field: 'FieldName'
            source_columns: [ 'FieldName', 'FieldName2' ] list of columns
            format: '{}-{}.{}' Python style format string with substitutions for source columns
    """
    cols_map = {}
    for spec in combine_spec:
        cols_map.update({
            spec['field']:
            format_column(lit(spec['format']), *spec['source_columns'])
        })

    lineage.update_lineage(df, args['source_key'], 'columnfromcolumn', transform=combine_spec)
    return df.withColumns(cols_map)


def transform_filterrows(df: DataFrame, filter_spec: list, args: dict, lineage, *extra):
    """Filter out (remove) rows of data that fail the provided filter operation
    NOTE: Use only when certain rows can be systematically discarded; otherwise use data quality quarantine rules

    Parameters
    ----------
    filter_spec
        List of dictionary objects in the form:
            condition: '`field` IS NOT NULL and `field2` == 'Test'
    """
    for spec in filter_spec:
        df = df.filter(spec['condition'])
    lineage.update_lineage(df, args['source_key'], 'filterrows', transform=filter_spec)
    return df
lib/glue_scripts/lib/datatransform_regex.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pyspark
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.functions import col, to_date, to_timestamp, regexp_extract, concat_ws, regexp_replace, initcap

def transform_date(df: DataFrame, date_formats: list, args: dict, lineage, *extra):
    """Convert specified date fields to ISO format based on known input format

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply date format conversion
    date_formats
        List of fieldnames and expected date format to convert _from_ in the form:
            field: 'FieldName'
            format: 'ExpectedDateFormat'
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with date format conversion applied
    """
    cols_map = {}
    for conversion in date_formats:
        cols_map.update({
            conversion['field']:
            to_date( col(conversion['field']), conversion['format'] )
        })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'dateconversion', transform=date_formats)
    return df.withColumns(cols_map)


def transform_timestamp(df: DataFrame, timestamp_formats: list, args: dict, lineage, *extra):
    """Convert specified date/time fields to ISO format based on known input format

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply timestamp format conversion
    timestamp_formats
        List of fieldnames and expected timestamp format to convert _from_ in the form:
            field: 'FieldName'
            format: 'ExpectedTimestampFormat'
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with timestamp format conversion applied
    """
    cols_map = {}
    for conversion in timestamp_formats:
        cols_map.update({
            conversion['field']:
            to_timestamp( col(conversion['field']), conversion['format'] )
        })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'timestampconversion', transform=timestamp_formats)
    return df.withColumns(cols_map)


def transform_decimal(df: DataFrame, decimal_formats: list, args: dict, lineage, *extra):
    """Convert specified numeric field (usually Float or Double) fields to Decimal (fixed
    precision) type

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply decimal conversion
    decimal_formats
        List of fieldnames and decimal precision, scale in the form:
            field: 'FieldName'
            format: 'precision,scale'
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with decimal conversion applied
    """
    cols_map = {}
    for conversion in decimal_formats:
        cols_map.update({
            conversion['field']:
            col(conversion['field']).cast(f"Decimal({conversion['format']})")
        })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'decimalconversion', transform=decimal_formats)
    return df.withColumns(cols_map)


def transform_implieddecimal(df: DataFrame, decimal_formats: list, args: dict, lineage, *extra):
    """Convert specified numeric field (usually Float or Double) fields to Decimal (fixed
    precision) type with implied decimal point support (i.e. last 2 digits are to the right of
    decimal, unless specified otherwise)

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply decimal conversion
    decimal_formats
        List of fieldnames and decimal precision, scale in the form:
            field: 'FieldName'
            format: 'precision,scale'
            num_implied: (optional) number of implied decimal points, default 2
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with decimal conversion applied
    """
    implied_decimal_pattern = r'([+-]?\d+)(\d{{{num_implied}}})$'

    cols_map = {}
    for conversion in decimal_formats:
        pattern = implied_decimal_pattern.format(num_implied = conversion.get('num_implied', 2))
        cols_map.update({
            conversion['field']:
            concat_ws(
                '.',
                regexp_extract(conversion['field'], pattern, 1),
                regexp_extract(conversion['field'], pattern, 2)
            ) \
            .cast(f"Decimal({conversion['format']})")
        })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'implieddecimalconversion', transform=decimal_formats)
    return df.withColumns(cols_map)


def transform_currency(df: DataFrame, currency_formats: list, args: dict, lineage, *extra):
    """Convert specified numeric field with currnecy formatting to Decimal (fixed precision)

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply currency conversion
    currency_formats
        List of fieldnames and decimal precision, scale in the form:
            field: 'FieldName'
            format: 'precision,scale' (optional, defaults to 16,2)
            euro: boolean (if true, expects European 5.000.000,12 format, default false)
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with currency conversion applied
    """
    cols_map = {}
    for conversion in currency_formats:
        decimal_format = conversion.get('format', '16,2')
        if conversion.get('euro', False):
            # Decimal = ,  Thousands = .
            new_column = regexp_replace(regexp_replace(conversion['field'], r'[^\-\d,-]+', ''), r',', '.')
        else:
            # Decimal = .  Thousands = ,
            new_column = regexp_replace(conversion['field'], r'[^\-\d\.]+', '')
        cols_map.update({
            conversion['field']: new_column.cast(f'Decimal({decimal_format})')
        })

    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'currencyconversion', transform=currency_formats)
    return df.withColumns(cols_map)


def transform_titlecase(df: DataFrame, titlecase_fields: list, args: dict, lineage, *extra):
    """Convert specified string field to title case

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply title case conversion
    titlecase_fields
        Simple list of fieldnames
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with title case conversion applied
    """
    cols_map = {
        field: initcap(col(field))
            for field in titlecase_fields
    }
    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'titlecaseconversion', transform=titlecase_fields)
    return df.withColumns(cols_map)


def transform_bigint(df: DataFrame, bigint_fields: list, args: dict, lineage, *extra):
    """Convert specified numeric field to bigint

    Parameters
    ----------
    df
        pySpark DataFrame on which to apply bigint conversion
    bigint_fields
        Simple list of fieldnames to covert
    args
        Glue job arguments, from which source_key and base_file_name are used
    lineage
        Initialized lineage class object from the calling job

    Returns
    -------
    DataFrame
        pySpark DataFrame with bigint conversion applied
    """
    cols_map = {
        field: col(field).cast('long')
            for field in bigint_fields
    }
    # No need to unpersist as there is only one reference to the dataframe and it is returned
    lineage.update_lineage(df, args['source_key'], 'bigintconversion', transform=bigint_fields)
    return df.withColumns(cols_map)
lib/glue_scripts/lib/datatransform_typeconversion.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import boto3
import botocore
import io
import csv
from urllib.parse import urlparse
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.types import DoubleType, NullType
from pyspark.sql.functions import col
from awsglue.dynamicframe import DynamicFrame

def table_exists(target_database: str, table_name: str) -> bool:
    """Function to check if table exists returns true/false
    All other exceptions from get_table are raised normally

    Parameters
    ----------
    target_database
        Glue Catalog database in which table exists
    table_name
        Glue Catalog table to check for existance

    Returns
    -------
    bool
        whether table exists or not
    """
    try:
        glue_client = boto3.client('glue')
        glue_client.get_table(DatabaseName=target_database, Name=table_name)
        return True
    except glue_client.exceptions.EntityNotFoundException:
        return False


def create_database(database_name: str):
    """Function to create catalog database if does not exists

    Parameters
    ----------
    database_name
        Glue Catalog database name for creating or confirming existance
    """
    response = None
    glue_client = boto3.client('glue')

    try:
        response = glue_client.get_database(
            Name=database_name
        )
    except botocore.exceptions.ClientError as error:
        if error.response['Error']['Code'] == 'EntityNotFoundException':
            print(f'The requested database: {database_name} was not found, creating it')
        else:
            raise RuntimeError(f"Error getting database {database_name}: {error.response['Error']['Message']}")

    if response is None:
        # Database does not exist, so we should create it
        response = glue_client.create_database(
            DatabaseInput={
                'Name': database_name
            }
        )
        print(f'create_database response: {response}')


def check_schema_change(existing_schema: list, new_schema: list, allow_schema_change: str) -> bool:
    """Function to check if schema change is allowed based on allow_schema_change setting

    Parameters
    ----------
    existing_schema
        Schema that already exists in the data lake
    new_schema
        Incoming data file schema that may be different from the existing schema
    allow_schema_changes
        strict - no changes
        reorder - allow reordering of columns
        evolve - allow reorder, adding, changing certain data types
        permissive - allow everything (may break Athena queries)

    Returns
    -------
    bool
        New schema passes the rules for the provided schema change setting

    """
    if allow_schema_change == 'permissive':
        # Permissive
        print('Permissive schema change: allowed')
        return True

    if allow_schema_change == 'strict':
        # Strict
        print('Strict schema change: not allowed')
        return existing_schema == new_schema

    existing_schema_set = set([ field_def['Name'] for field_def in existing_schema ])
    new_schema_set = set([ field_def['Name'] for field_def in new_schema ])
    if allow_schema_change == 'reorder':
        # Reorder (duplicate fields in new_schema will fail)
        print('Reorder schema change: checking')
        return existing_schema_set == new_schema_set \
            and len(new_schema) == len(new_schema_set)

    if allow_schema_change == 'evolve':
        pass
        # TODO: Finish implementation
        # field additions: new_schema_set - existing_schema_set
        # field removals: existing_schema_set - new_schema_set

    raise RuntimeError(f'Unsupported value for allow_schema_change {allow_schema_change}')


def upsert_catalog_table(
        df: any, 
        target_database: str, 
        table_name: str, 
        partition_keys: list, 
        storage_location: str, 
        allow_schema_change: str = 'strict',
    ):
    """Function to upsert Glue Data Catalog table

    Parameters
    ----------
    df
        Spark Dataframe or Glue DynamicFrame from which to gather schema
    target_database
        Glue Catalog database name to use
    table_name
        Glue Catalog table name to use
    partition_keys
        list of key names to use as partition keys
    storage_location
        S3 path to data stored in parquet files
    allow_schema_change: optional
        Schema evolution setting to be used by check_schema_change()
    """
    create_database(target_database)

    df_schema = df.toDF().dtypes if type(df) == DynamicFrame else df.dtypes
    schema = [
        # Glue Catalog will transparently convert all column names to lowercase
        # so do it explicitly so that schema change detection works
        { 'Name': field[0].lower(), 'Type': field[1] }
        for field in df_schema
            # Skip explicitly defining partition keys in the catalog
            # schema if they already exist in the dataframe schema
            if field[0] not in partition_keys
    ]

    input_format = 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat'
    output_format = 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat'
    serde_info = {
        'Parameters': {
            'serialization.format': '1'
        },
        'SerializationLibrary': 'org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe'
    }
    storage_descriptor = {
        'Columns': schema,
        'InputFormat': input_format,
        'OutputFormat': output_format,
        'SerdeInfo': serde_info,
        'Location': storage_location
    }
    partition_key = [
        { 'Name': key_name, 'Type': 'string' } for key_name in partition_keys
    ]
    table_input = {
        'Name': table_name,
        'StorageDescriptor': storage_descriptor,
        'Parameters': {
            'classification': 'PARQUET',
            'SourceType': 's3',
            'SourceTableName': table_name,
            'TableVersion': '0'
        },
        'TableType': 'EXTERNAL_TABLE',
        'PartitionKeys': partition_key
    }

    print(f'Upserting Glue Catalog schema: {schema}')
    try:
        glue_client = boto3.client('glue')
        if not table_exists(target_database, table_name):
            print(f'Target Table name: {table_name} does not exist')
            glue_client.create_table(DatabaseName=target_database, TableInput=table_input)
            return

        # Compare new schema to existing Glue catalog
        table_response = glue_client.get_table(DatabaseName=target_database, Name=table_name)
        schema_equal = ( table_response['Table']['StorageDescriptor']['Columns'] == schema )
        if schema_equal:
            print('No schema changes detected')
            return

        if check_schema_change(table_response['Table']['StorageDescriptor']['Columns'], schema, allow_schema_change):
            print(f'Trying to update Target Table: {table_name}')
            glue_client.update_table(DatabaseName=target_database, TableInput=table_input)
            return

        # Abort if there are changes because it breaks Athena queries to have partitions 
        # with different schema. Even though we are updating the same Glue catalog, the
        # existing partitions will retain the original schema. and be in conflict.
        print(f'Schema change not permissible by {allow_schema_change} alllow_schema_changes '
            'setting detected between existing partitions and new partitions; aborting')
        print(f"Existing schema: {table_response['Table']['StorageDescriptor']['Columns']}")
        print(f'New schema: {schema}')
        raise RuntimeError(f'Nonpermissible schema change detected with existing Glue catalog {target_database}.{table_name}')

    except botocore.exceptions.ClientError as error:
        print(f'Glue job client process failed: {error}')
        raise error


def clean_column_names(df: DataFrame) -> tuple:
    """Rename DataFrame column names to conform to Parquet/Athena naming requirements

    Parameters
    ----------
    df
        Spark DataFrame to use for field names and initial data types

    Returns
    -------
    DataFrame, field_map
        DataFrame with cleaned column names, and dictionary of mapping operations performed
    """
    cols = []
    field_map_rows = [[ 'SourceName', 'DestName' ]]
    for field in df.schema:
        column = col('`' + field.name + '`')

        # Remove leading/trailing whitespace, and trim length to max allowed
        new_name = field.name.strip()[:255].lower()

        for char in ',;{}()\n\t=':
            # Remove bad Parquet column naming characters
            new_name = new_name.replace(char, '')
        for char in ' .':
            # Replace characters for Parquet rules/conventions
            new_name = new_name.replace(char, '_')

        # Consolidate excessive character replacement
        new_name = new_name.replace('_-_', '-').replace('__', '_').replace('__', '_')

        if field.name != new_name:
            column = column.alias(new_name)
        cols.append(column)

        # Fields appear in the recommended map regardless of whether they were aliased
        field_map_rows.append([ field.name, new_name ])

    return df.select(cols), field_map_rows


def generate_spec(df: DataFrame, input_file_extension: str) -> dict:
    """Generate a suggested transform spec file (primarily type conversions)

    Parameters
    ----------
    df
        Spark DataFrame to use for field names and initial data types
    input_file_extension
        File extension of input file to determine if input spec is needed

    Returns
    -------
    dict
        Dictionary of recommended transform specifications
    """
    input_spec = {}
    if input_file_extension.lower() in [ '.xlsx', '.xls' ]:
        input_spec.update({
            'excel': { 'sheet_names': [ '0' ], 'data_address': 'A1', 'header': True }
        })

    transform_spec = { 'date': [], 'decimal': [] }
    for field in df.schema:
        if field.dataType == DoubleType():
            transform_spec['decimal'].append({ 'field': field.name, 'format': '16,2' })

        if 'date' in field.name.lower():
            transform_spec['date'].append({ 'field': field.name, 'format': 'MM/dd/yy' })

    return { 'input_spec': input_spec, 'transform_spec': transform_spec }


def put_s3_object(uri: str, data: any):
    """Save data to S3 object bypassing Spark write operation

    Parameters
    ----------
    uri
        Full HDFS S3 path URI to create
    data
        list or string of data to write to S3 (list will be converted to CSV)
    """
    if type(data) is list:
        string_output = io.StringIO()
        writer = csv.writer(string_output)
        writer.writerows(data)
        data = string_output.getvalue()
    parsed_uri = urlparse(uri)
    s3 = boto3.client('s3')
    s3.put_object(Body=str(data), Bucket=parsed_uri.netloc, Key=parsed_uri.path[1:])


def clean_nulltypes(df: DataFrame) -> DataFrame:
    """Convert Void/NullType columns typically created from DynamicFrame operations to string
    """
    cols = []
    for field in df.schema:
        column = col('`' + field.name + '`')
        if field.dataType == NullType():
            # Cast null and void type columns to string (safest conversion)
            column = column.cast('string')
        cols.append(column)
    return df.select(cols)
lib/glue_scripts/lib/glue_catalog_helpers.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import aws_cdk as cdk
from constructs import Construct
import aws_cdk.aws_s3 as s3
import aws_cdk.aws_logs as logs
import aws_cdk.aws_glue as glue
import aws_cdk.aws_iam as iam
import aws_cdk.aws_kms as kms
import aws_cdk.aws_dynamodb as dynamodb
import aws_cdk.aws_s3_deployment as s3_deployment
import os
from cdk_nag import NagSuppressions

from .stack_import_helper import ImportedBuckets, ImportedVpc
from .configuration import (
    DEV, PROD, TEST, get_logical_id_prefix, get_resource_name_prefix, get_environment_configuration,
)

class GlueStack(cdk.Stack):
    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        target_environment: str,
        hash_values_table: dynamodb.Table,
        value_lookup_table: dynamodb.Table,
        multi_lookup_table: dynamodb.Table,
        dq_results_table: dynamodb.Table,
        data_lineage_table: dynamodb.Table = None,
        **kwargs
    ):
        """CloudFormation stack to create Glue Jobs, Connections,
        Script Bucket, Temporary Bucket, and an IAM Role for permissions.

        Parameters
        ----------
        scope
            Parent of this stack, usually an App or a Stage, but could be any construct
        construct_id
            The construct ID of this stack; if stackName is not explicitly defined,
            this ID (and any parent IDs) will be used to determine the physical ID of the stack
        target_environment
            The target environment for stacks in the deploy stage
        hash_values_table
            The DynamoDB Table for storing original values from hashing function
        value_lookup_table
            The DynamoDB Table for looking up values in the multi lookup transform
        multi_lookup_table
            The DynamoDB Table for looking up values in the multi lookup transform
        dq_results_table
            The DynamoDB Table for storing Glue Data Quality results
        data_lineage_table: optional
            Optional DynamoDB Table for storing custom data lineage information for Glue job
            transforms; omitting the parameter disables custom data lineage tracking
        kwargs: optional
            Optional keyword arguments to pass up to parent Stack class
        """
        super().__init__(scope, construct_id, **kwargs)

        self.target_environment = target_environment
        self.mappings = get_environment_configuration(target_environment)
        self.logical_id_prefix = get_logical_id_prefix()
        self.resource_name_prefix = get_resource_name_prefix()
        if (target_environment == PROD or target_environment == TEST):
            self.removal_policy = cdk.RemovalPolicy.RETAIN
            self.log_retention = logs.RetentionDays.SIX_MONTHS
        else:
            self.removal_policy = cdk.RemovalPolicy.DESTROY
            self.log_retention = logs.RetentionDays.ONE_MONTH

        buckets = ImportedBuckets(self, logical_id_suffix='GlueStack')
        vpc = ImportedVpc(self, logical_id_suffix='GlueStack')

        self.glue_scripts_bucket = self.get_glue_scripts_bucket(
            buckets.s3_kms_key,
            buckets.access_logs
        )
        self.glue_scripts_temp_bucket = self.get_glue_scripts_temporary_bucket(
            buckets.s3_kms_key,
            buckets.access_logs
        )
        glue_role = self.get_glue_role(
            buckets.s3_kms_key,
            buckets=[
                buckets.raw,
                buckets.conformed,
                buckets.purposebuilt,
                self.glue_scripts_bucket,
                self.glue_scripts_temp_bucket,
            ],
            dynamodb_tables=[
                hash_values_table,
                value_lookup_table,
                multi_lookup_table,
                dq_results_table,
                data_lineage_table,
            ],
        )

        # Each module must be listed specifically for --extra-py-files and --extra-jars (or Zip)
        glue_libraries = [
            f's3://{self.glue_scripts_bucket.bucket_name}/etl/lib/{dirent.name}'
                for dirent in os.scandir('lib/glue_scripts/lib')
                    if os.path.splitext(dirent.path)[1] == '.py'
        ]
        spark_libraries = [
            f's3://{self.glue_scripts_bucket.bucket_name}/etl/lib/{dirent.name}'
                for dirent in os.scandir('lib/glue_scripts/lib')
                    if os.path.splitext(dirent.path)[1] == '.jar'
        ]

        job_connections = []
        for subnet_number in range(len(vpc.subnets)):
            job_connections.append(glue.CfnConnection(
                self,
                f'{target_environment}{self.logical_id_prefix}CollectToCleanseWorkflowConnection{subnet_number + 1}',
                catalog_id=self.account,
                connection_input=glue.CfnConnection.ConnectionInputProperty(
                    connection_type="NETWORK",
                    name=f'{target_environment.lower()}-{self.resource_name_prefix}-collect-to-cleanse-connection-{subnet_number + 1}',
                    physical_connection_requirements=glue.CfnConnection.PhysicalConnectionRequirementsProperty(
                        availability_zone=vpc.subnets[subnet_number].availability_zone,
                        subnet_id=vpc.subnets[subnet_number].subnet_id,
                        security_group_id_list=[vpc.shared_security_group.security_group_id]
                    )
                )
            ))

        self.collect_to_cleanse_job = glue.CfnJob(
            self,
            f'{target_environment}{self.logical_id_prefix}CollectToCleanseJob',
            name=f'{target_environment.lower()}-{self.resource_name_prefix}-collect-to-cleanse-job',
            description='PySpark Glue job data processing logic to cleanse and curate collected source system data',
            command=glue.CfnJob.JobCommandProperty(
                name='glueetl',
                python_version='3',
                script_location=f's3://{self.glue_scripts_bucket.bucket_name}/etl/etl_collect_to_cleanse.py'
            ),
            # Used if Glue job needs connections to resources in VPCs (incurs VPC costs and may trigger IP limitations)
            connections=glue.CfnJob.ConnectionsListProperty(
                connections=[ job_connection.connection_input.name for job_connection in job_connections ],
            ) if job_connections else None,
            # These arguments are common to all Glue job runs and are overlayed by the arguments
            # definition in the calling Step Functions GlueStartJobRun
            default_arguments={
                '--enable-auto-scaling': 'true',
                '--enable-continuous-cloudwatch-log': 'true',
                '--enable-metrics': 'true',
                '--enable-glue-datacatalog': 'true',
                '--user-jars-first': 'true',
                '--extra-jars': ','.join(spark_libraries) if spark_libraries else None,
                '--additional-python-modules': 'rapidfuzz',
                '--extra-py-files': ','.join(glue_libraries),
                '--environment': self.target_environment,
                '--TempDir': f's3://{self.glue_scripts_temp_bucket.bucket_name}/etl/collect-to-cleanse',
                '--txn_bucket': f's3://{self.glue_scripts_bucket.bucket_name}',
                '--txn_spec_prefix_path': '/etl/transformation-spec/',
                '--source_bucket': f's3://{buckets.raw.bucket_name}',
                '--target_bucket': f's3://{buckets.conformed.bucket_name}',
                '--hash_value_table': hash_values_table.table_name,
                '--value_lookup_table': value_lookup_table.table_name,
                '--multi_lookup_table': multi_lookup_table.table_name,
                '--dq_results_table': dq_results_table.table_name,
                '--data_lineage_table': data_lineage_table.table_name if data_lineage_table else None,
            },
            execution_property=glue.CfnJob.ExecutionPropertyProperty(
                max_concurrent_runs=10,
            ),
            glue_version='4.0',
            max_retries=0,
            # With auto-scaling, this represents the maximum number of workers
            # If using a Connection, there must be enough IP addresses for each worker
            number_of_workers=50,
            role=glue_role.role_arn,
            worker_type='G.1X',
            # TODO: Allow the user to specify a user-managed, out-of-stack security group name
            # Glue security configurations cannot be updated by Cloudformation and break all stack updates
            #security_configuration='',
        )

        self.cleanse_to_consume_job = glue.CfnJob(
            self,
            f'{target_environment}{self.logical_id_prefix}CleanseToConsumeJob',
            name=f'{target_environment.lower()}-{self.resource_name_prefix}-cleanse-to-consume-job',
            description='PySpark Glue job data processing logic to prepare cleansed Data Lake tables for specific analytics consumption',
            command=glue.CfnJob.JobCommandProperty(
                name='glueetl',
                python_version='3',
                script_location=f's3://{self.glue_scripts_bucket.bucket_name}/etl/etl_cleanse_to_consume.py'
            ),
            # Used if Glue job needs connections to resources in VPCs (incurs VPC costs and may trigger IP limitations)
            connections=glue.CfnJob.ConnectionsListProperty(
                connections=[ job_connection.connection_input.name for job_connection in job_connections ],
            ) if job_connections else None,
            # These arguments are common to all Glue job runs and are overlayed by the arguments
            # definition in the calling Step Functions GlueStartJobRun
            default_arguments={
                '--enable-auto-scaling': 'true',
                '--enable-glue-datacatalog': 'true',
                '--enable-continuous-cloudwatch-log': 'true',
                '--enable-metrics': 'true',
                '--user-jars-first': 'true',
                '--extra-jars': ','.join(spark_libraries) if spark_libraries else None,
                '--extra-py-files': ','.join(glue_libraries),
                '--environment': self.target_environment,
                '--TempDir': f's3://{self.glue_scripts_temp_bucket.bucket_name}/etl/cleanse_to_consume',
                '--txn_bucket': f's3://{self.glue_scripts_bucket.bucket_name}',
                '--txn_sql_prefix_path': '/etl/transformation-sql/',
                '--source_bucket': f's3://{buckets.conformed.bucket_name}',
                '--target_bucket': f's3://{buckets.purposebuilt.bucket_name}',
                '--data_lineage_table': data_lineage_table.table_name if data_lineage_table else None,
            },
            execution_property=glue.CfnJob.ExecutionPropertyProperty(
                max_concurrent_runs=10,
            ),
            glue_version='4.0',
            max_retries=0,
            # With auto-scaling, this represents the maximum number of workers
            # If using a Connection, there must be enough IP addresses for each worker
            number_of_workers=50,
            role=glue_role.role_arn,
            worker_type='G.1X',
            # TODO: Allow the user to specify a user-managed, out-of-stack security group name
            # Glue security configurations cannot be updated by Cloudformation and break all stack updates
            #security_configuration='',
        )

        # Recommended encryption settings for account Glue Data Catalog
        # Applies to all databases and tables in the account; uncomment to apply
        # glue.CfnDataCatalogEncryptionSettings(
        #     self,
        #     f'{target_environment}{self.logical_id_prefix}DataCatalogEncryptionSettings',
        #     catalog_id=f'{self.account}',
        #     data_catalog_encryption_settings=glue.CfnDataCatalogEncryptionSettings.DataCatalogEncryptionSettingsProperty(
        #         connection_password_encryption=glue.CfnDataCatalogEncryptionSettings.ConnectionPasswordEncryptionProperty(
        #             kms_key_id=buckets.s3_kms_key.key_arn,
        #             return_connection_password_encrypted=True,
        #         ),
        #         encryption_at_rest=glue.CfnDataCatalogEncryptionSettings.EncryptionAtRestProperty(
        #             catalog_encryption_mode="SSE-KMS",
        #             sse_aws_kms_key_id=buckets.s3_kms_key.key_arn,
        #         )
        #     )
        # )

        NagSuppressions.add_resource_suppressions(self, [
            {
                'id': 'AwsSolutions-GL1',
                'reason': 'Creating a Glue security configuration in CDK prevents the stack from updating; log messages are free from sensitive data'
            },
            {
                'id': 'AwsSolutions-GL3',
                'reason': 'Creating a Glue security configuration in CDK prevents the stack from updating; the Glue job has bookmarks disabled'
            }
        ], apply_to_children=True)


    def get_glue_scripts_bucket(
        self,
        s3_kms_key: kms.Key,
        access_logs_bucket: s3.Bucket
    ) -> s3.Bucket:
        """Creates S3 Bucket that contains glue scripts used in Job execution

        Parameters
        ----------
        s3_kms_key
            The KMS Key to use for encryption of data at rest
        access_logs_bucket
            The access logs target for this bucket

        Returns
        -------
        s3.Bucket
            Glue scripts bucket construct
        """
        bucket_name = f'{self.target_environment.lower()}-{self.resource_name_prefix}-{self.account}-{self.region}-etl-scripts'
        bucket = s3.Bucket(
            self,
            f'{self.target_environment}{self.logical_id_prefix}GlueScriptsBucket',
            bucket_name=bucket_name,
            access_control=s3.BucketAccessControl.PRIVATE,
            block_public_access=s3.BlockPublicAccess.BLOCK_ALL,
            enforce_ssl=True,
            bucket_key_enabled=s3_kms_key is not None,
            encryption=s3.BucketEncryption.KMS,
            encryption_key=s3_kms_key,
            public_read_access=False,
            removal_policy=self.removal_policy,
            versioned=True,
            object_ownership=s3.ObjectOwnership.OBJECT_WRITER,
            server_access_logs_bucket=access_logs_bucket,
            server_access_logs_prefix=f'{bucket_name}-',
        )

        # Dynamically upload resources to the script target
        # TODO: Apply resource retention policy to Custom Resource Lambda log group
        s3_deployment.BucketDeployment(
            self,
            'DeployGlueJobScript',
            # This path is relative to the root of the project
            sources=[s3_deployment.Source.asset('./lib/glue_scripts')],
            destination_bucket=bucket,
            destination_key_prefix='etl',
            prune=self.removal_policy == cdk.RemovalPolicy.DESTROY,
            retain_on_delete=self.removal_policy == cdk.RemovalPolicy.RETAIN,
            log_retention=self.log_retention,
        )

        NagSuppressions.add_resource_suppressions(self, [
            {
                'id': 'AwsSolutions-IAM4',
                'reason': 'S3 Deployment CustomResource used only during stack deployment and deletion'
            },
            {
                'id': 'AwsSolutions-IAM5',
                'reason': 'S3 Deployment CustomResource used only during stack deployment and deletion'
            },
            {
                'id': 'AwsSolutions-L1',
                'reason': 'Deployment CustomResource Lambda is maintained by the CDK team and that team needs to update the runtime'
            },
        ], apply_to_children=True)

        return bucket


    def get_glue_scripts_temporary_bucket(
        self,
        s3_kms_key: kms.Key,
        access_logs_bucket: s3.Bucket
    ) -> s3.Bucket:
        """Creates S3 Bucket used as a temporary file store in Job execution

        Parameters
        ----------
        s3_kms_key
            The KMS Key to use for encryption of data at rest
        access_logs_bucket
            The access logs target for this bucket

        Returns
        -------
        s3.Bucket
            Glue scripts bucket construct
        """
        bucket_name = f'{self.target_environment.lower()}-{self.resource_name_prefix}-{self.account}-{self.region}-glue-temp'
        bucket = s3.Bucket(
            self,
            f'{self.target_environment}{self.logical_id_prefix}GlueScriptsTemporaryBucket',
            bucket_name=bucket_name,
            access_control=s3.BucketAccessControl.PRIVATE,
            block_public_access=s3.BlockPublicAccess.BLOCK_ALL,
            enforce_ssl=True,
            bucket_key_enabled=s3_kms_key is not None,
            encryption=s3.BucketEncryption.KMS if s3_kms_key else s3.BucketEncryption.S3_MANAGED,
            encryption_key=s3_kms_key if s3_kms_key else None,
            public_read_access=False,
            removal_policy=cdk.RemovalPolicy.DESTROY,
            versioned=True,
            object_ownership=s3.ObjectOwnership.OBJECT_WRITER,
            server_access_logs_bucket=access_logs_bucket,
            server_access_logs_prefix=f'{bucket_name}-',
        )

        return bucket


    def get_glue_role(
        self,
        s3_kms_key: kms.Key,
        buckets: list,
        dynamodb_tables: list,
    ) -> iam.Role:
        """Creates the role used during Glue Job execution

        Parameters
        ----------
        s3_kms_key
            The KMS Key to use for encryption of data at rest
        buckets
            List of S3 Bucket constructs to use for S3 policy resources
        dynamodb_tables
            List of DynamoDB Table constructs to use for DynamoDB policy resources

        Returns
        -------
        iam.Role
            The IAM role that was created
        """
        bucket_object_resources = [ bucket.arn_for_objects('*') for bucket in buckets ]
        bucket_resources = [ bucket.bucket_arn for bucket in buckets ]
        dynamodb_resources = [ table.table_arn for table in dynamodb_tables if table is not None ]

        glue_role = iam.Role(
            self,
            f'{self.target_environment}{self.logical_id_prefix}GlueRole',
            description='Role for InsuranceLake ETL pipeline Glue Jobs',
            role_name=f'{self.target_environment.lower()}-{self.resource_name_prefix}-{self.region}-glue-role',
            assumed_by=iam.ServicePrincipal('glue.amazonaws.com'),
            inline_policies={
                'S3BucketReadAccess':
                iam.PolicyDocument(statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        actions=[
                            's3:ListBucketVersions',
                            's3:ListBucket',
                            's3:GetBucketNotification',
                            's3:GetBucketLocation',
                        ],
                        resources=bucket_resources
                    )
                ]),
                # TODO: Remove glue scripts bucket from the resource list
                'S3BucketWriteAccess':
                iam.PolicyDocument(statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        actions=[
                            's3:ReplicateObject',
                            's3:PutObject',
                            's3:GetObject',
                            's3:DeleteObject',
                        ],
                        resources=bucket_object_resources
                    )
                ]),
                # This is required due to bucket level encryption on S3 Buckets
                'KmsAccess':
                iam.PolicyDocument(statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        actions=[
                            'kms:Decrypt',
                            'kms:GenerateDataKey',
                        ],
                        resources=[
                            s3_kms_key.key_arn,
                        ]
                    )
                ]),
                'DynamoDBAccess':
                iam.PolicyDocument(statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        actions=[
                            'dynamodb:PutItem',
                            'dynamodb:GetItem',
                            'dynamodb:DeleteItem',
                            'dynamodb:UpdateItem',
                            'dynamodb:DescribeTable',
                            'dynamodb:BatchWriteItem',
                            'dynamodb:BatchReadItem',
                            'dynamodb:Query',
                        ],
                        resources=dynamodb_resources,
                    )
                ]),
                # TODO: Use a custom workgroup to restrict resource access
                'AthenaAccess':
                iam.PolicyDocument(statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        actions=[
                            'athena:GetQueryExecution',
                            'athena:GetQueryResults',
                            'athena:StartQueryExecution',
                        ],
                        resources=[
                            '*',
                        ]
                    )
                ]),
            },
            # TODO: Convert this attached policy to specifically needed permissions and remove nag suppression
            managed_policies=[
                iam.ManagedPolicy.from_aws_managed_policy_name('service-role/AWSGlueServiceRole'),
            ]
        )

        NagSuppressions.add_resource_suppressions(glue_role, [
            {
                'id': 'AwsSolutions-IAM4',
                'reason': 'Glue Job should use built-in Glue Service Role'
            },
            {
                'id': 'AwsSolutions-IAM5',
                'reason': 'AWS provided Glue Service Role uses wildcard permissions'
            },
        ])

        return glue_role
lib/glue_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import aws_cdk as cdk
from constructs import Construct
from .step_functions_stack import StepFunctionsStack
from .glue_stack import GlueStack
from .dynamodb_stack import DynamoDbStack
from .athena_helper_stack import AthenaHelperStack
from .tagging import tag
from .configuration import (
    get_logical_id_prefix,
)

class PipelineDeployStage(cdk.Stage):
    def __init__(
        self, scope: Construct, construct_id: str,
        target_environment: str, env: cdk.Environment=None,
        **kwargs
    ):
        """Adds deploy stage to CodePipeline

        Parameters
        ----------
        scope
            Parent of this stack, usually an App or a Stage, but could be any construct
        construct_id
            The construct ID of this stack; if stackName is not explicitly defined,
            this ID (and any parent IDs) will be used to determine the physical ID of the stack
        target_environment
            The target environment for stacks in the deploy stage
        env: optional
            AWS environment definition (account, region) to pass to stacks
        kwargs: optional
            Optional keyword arguments
        """
        super().__init__(scope, construct_id, **kwargs)
        logical_id_prefix = get_logical_id_prefix()

        dynamodb_stack = DynamoDbStack(
            self,
            f'{logical_id_prefix}EtlDynamoDb',
            description='InsuranceLake stack for DynamoDB tables to store pipeline state and processing history (uksb-1tu7mtee2)',
            target_environment=target_environment,
            env=env,
            **kwargs,
        )

        glue_stack = GlueStack(
            self,
            f'{logical_id_prefix}EtlGlue',
            description='InsuranceLake stack for Glue jobs to support the data pipeline (uksb-1tu7mtee2)',
            target_environment=target_environment,
            env=env,
            hash_values_table=dynamodb_stack.hash_values_table,
            value_lookup_table=dynamodb_stack.value_lookup_table,
            multi_lookup_table=dynamodb_stack.multi_lookup_table,
            dq_results_table=dynamodb_stack.dq_results_table,
            data_lineage_table=dynamodb_stack.data_lineage_table,
            **kwargs,
        )

        step_function_stack = StepFunctionsStack(
            self,
            f'{logical_id_prefix}EtlStepFunctions',
            description='InsuranceLake stack for Step Functions and supporting Lambda functions to orchestrate data pipeline steps (uksb-1tu7mtee2)',
            target_environment=target_environment,
            env=env,
            collect_to_cleanse_job=glue_stack.collect_to_cleanse_job,
            cleanse_to_consume_job=glue_stack.cleanse_to_consume_job,
            job_audit_table=dynamodb_stack.job_audit_table,
            **kwargs,
        )

        athena_helper_stack = AthenaHelperStack(
            self,
            f'{logical_id_prefix}EtlAthenaHelper',
            description='InsuranceLake stack for Athena helper Workgroup (uksb-1tu7mtee2)',
            target_environment=target_environment,
            env=env,
            glue_scripts_temp_bucket=glue_stack.glue_scripts_temp_bucket,
            **kwargs,
        )

        tag(step_function_stack, target_environment)
        tag(dynamodb_stack, target_environment)
        tag(glue_stack, target_environment)
        tag(athena_helper_stack, target_environment)
lib/pipeline_deploy_stage.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import aws_cdk as cdk
from constructs import Construct
import aws_cdk.pipelines as Pipelines
import aws_cdk.aws_s3 as s3
import aws_cdk.aws_iam as iam
import aws_cdk.aws_logs as logs
import aws_cdk.aws_codepipeline as CodePipeline
import aws_cdk.aws_codepipeline_actions as CodePipelineActions
import aws_cdk.aws_codebuild as CodeBuild
import aws_cdk.aws_codecommit as CodeCommit
from cdk_nag import AwsSolutionsChecks, NagSuppressions

from .configuration import (
    ACCOUNT_ID, CODECOMMIT_MIRROR_REPOSITORY_NAME, DEPLOYMENT, PROD, TEST,
	GITHUB_REPOSITORY_NAME, GITHUB_REPOSITORY_OWNER_NAME, GITHUB_TOKEN,
	CODECOMMIT_REPOSITORY_NAME, CODECOMMIT_MIRROR_REPOSITORY_NAME,
    get_logical_id_prefix, get_resource_name_prefix, get_all_configurations
)
from .pipeline_deploy_stage import PipelineDeployStage


class PipelineStack(cdk.Stack):

    def __init__(
        self, scope: Construct, construct_id: str,
        target_environment: str, target_branch: str, target_aws_env: dict,
        **kwargs
    ):
        """CloudFormation stack to create CDK Pipeline resources (Code Pipeline, Code Build, and
        ancillary resources).

        Parameters
        ----------
        scope
            Parent of this stack, usually an App or a Stage, but could be any construct
        construct_id
            The construct ID of this stack; if stackName is not explicitly defined,
            this ID (and any parent IDs) will be used to determine the physical ID of the stack
        target_environment
            The target environment for stacks in the deploy stage
        target_branch
            The source branch for polling
        target_aws_env
            The CDK env variables used for stacks in the deploy stage
        kwargs: optional
            Optional keyword arguments to pass up to parent Stack class
        """
        super().__init__(scope, construct_id, **kwargs)

        self.mappings = get_all_configurations()

        if (target_environment == PROD or target_environment == TEST):
            self.removal_policy = cdk.RemovalPolicy.RETAIN
            self.log_retention = logs.RetentionDays.SIX_MONTHS
        else:
            self.removal_policy = cdk.RemovalPolicy.DESTROY
            self.log_retention = logs.RetentionDays.ONE_MONTH

        self.create_environment_pipeline(
            target_environment,
            target_branch,
            target_aws_env
        )

    def create_environment_pipeline(
        self,
        target_environment: str, target_branch: str, target_aws_env: dict,
    ):
        """Creates CloudFormation stack to create CDK Pipeline resources such as:
        Code Pipeline, Code Build, and ancillary resources.

        Parameters
        ----------
        target_environment
            The target environment for stacks in the deploy stage
        target_branch
            The source branch for polling
        target_aws_env
            The CDK env variables used for stacks in the deploy stage
        """
        logical_id_prefix = get_logical_id_prefix()
        resource_name_prefix = get_resource_name_prefix()

        code_build_env = CodeBuild.BuildEnvironment(
            build_image=CodeBuild.LinuxBuildImage.STANDARD_6_0,
            privileged=False
        )
        
        code_build_opt = Pipelines.CodeBuildOptions(
            build_environment=code_build_env,
            role_policy=[
                iam.PolicyStatement(
                    sid='EtlPipelineSecretsManagerPolicy',
                    effect=iam.Effect.ALLOW,
                    actions=[
                        'secretsmanager:GetSecretValue',
                    ],
                    resources=[
                        f'arn:aws:secretsmanager:{self.region}:{self.account}:secret:/InsuranceLake/*',
                    ],
                ),
                iam.PolicyStatement(
                    sid='EtlPipelineSTSAssumeRolePolicy',
                    effect=iam.Effect.ALLOW,
                    actions=[
                        'sts:AssumeRole',
                    ],
                    resources=[
                        '*',
                    ],
                ),
            ]
        )


        if self.mappings[DEPLOYMENT][GITHUB_REPOSITORY_NAME] != '':
            # Github
            codepipeline_source = Pipelines.CodePipelineSource.git_hub(
                repo_string=f"{self.mappings[DEPLOYMENT][GITHUB_REPOSITORY_OWNER_NAME]}/{self.mappings[DEPLOYMENT][GITHUB_REPOSITORY_NAME]}",
                branch=target_branch,
                authentication=cdk.SecretValue.secrets_manager(
                    self.mappings[DEPLOYMENT][GITHUB_TOKEN]
                ),
                trigger=CodePipelineActions.GitHubTrigger.POLL,
            )
        else:
            # CodeCommit
            if self.mappings[DEPLOYMENT][CODECOMMIT_MIRROR_REPOSITORY_NAME] != '':
                repo = CodeCommit.Repository.from_repository_name(
                    self,
                    f'{DEPLOYMENT}{logical_id_prefix}EtlMirrorRepository',
                    repository_name=self.mappings[DEPLOYMENT][CODECOMMIT_MIRROR_REPOSITORY_NAME],
                )
            else:
                repo = CodeCommit.Repository.from_repository_name(
                    self,
                    f'{DEPLOYMENT}{logical_id_prefix}EtlRepository',
                    repository_name=self.mappings[DEPLOYMENT][CODECOMMIT_REPOSITORY_NAME],
                )

            codepipeline_source = Pipelines.CodePipelineSource.code_commit(
                repository=repo,
                branch=target_branch,
                code_build_clone_output=True,
            )

        pipeline = Pipelines.CodePipeline(
            self,
            f'{target_environment}{logical_id_prefix}EtlPipeline',
            pipeline_name=f'{target_environment.lower()}-{resource_name_prefix}-etl-pipeline',
            code_build_defaults=code_build_opt,
            self_mutation=True,
            synth=Pipelines.ShellStep(
                'Synth',
                input=codepipeline_source,
                commands=[
                    'npm install -g aws-cdk',
                    'python -m pip install -r requirements.txt --root-user-action=ignore',
                    'cdk synth'
                ],
            ),
            cross_account_keys=True
        )

        pipeline_deploy_stage = PipelineDeployStage(
                self,
                target_environment,
                target_environment=target_environment,
                env=cdk.Environment(
                    account=target_aws_env['account'],
                    region=target_aws_env['region']
                )
            )

        # Enable cdk-nag for environment stacks before adding to
        # pipeline, which are deployed with CodePipeline
        cdk.Aspects.of(pipeline_deploy_stage).add(AwsSolutionsChecks())

        pipeline.add_stage(pipeline_deploy_stage)

        # Force Pipeline construct creation during synth so we can add
        # Nag Supressions. Artifact bucket policies, and access Build
        # stages
        pipeline.build_pipeline()

        # Loop through Stages and Actions looking for Build actions
        # that write to CloudWatch logs
        for stage in pipeline.pipeline.stages:
            for action in stage.actions:
                if action.action_properties.category == CodePipeline.ActionCategory.BUILD:
                    logs.LogGroup(
                        self, 
                        id=f'CodeBuildAction{action.action_properties.action_name}LogGroup',
                        # Name the log after the project name so it
                        # matches where CodeBuild writes
                        # resource object is a PipelineProject
                        log_group_name=f'/aws/codebuild/{action.action_properties.resource.project_name}',
                        removal_policy=self.removal_policy,
                        retention=self.log_retention,
                    )

        # Apply stack removal policy to Artifact Bucket
        pipeline.pipeline.artifact_bucket.apply_removal_policy(self.removal_policy)

        # Enable server access logs in the same bucket using escape hatch
        cfn_artifact_bucket = pipeline.pipeline.artifact_bucket.node.default_child
        cfn_artifact_bucket.logging_configuration = s3.CfnBucket.LoggingConfigurationProperty(
            # TODO: Convert to separate bucket that is part of the Pipeline stack
            log_file_prefix='access-logs'
        )
        # Enable artifact bucket encryption key rotation using escape hatch
        cfn_artifact_bucket_encryption_key = pipeline.pipeline.artifact_bucket.encryption_key.node.default_child
        cfn_artifact_bucket_encryption_key.enable_key_rotation = True

        # Apply Nag Suppression to all Pipeline resources (many role and policies)
        NagSuppressions.add_resource_suppressions(pipeline, [
            {
                'id': 'AwsSolutions-IAM5',
                'reason': 'Wildcard IAM permissions are used by auto-created Codepipeline policies and custom policies to allow flexible creation of resources'
            },
        ], apply_to_children=True)
lib/pipeline_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import aws_cdk as cdk
import aws_cdk.aws_s3 as s3
import aws_cdk.aws_ec2 as ec2
import aws_cdk.aws_kms as kms

from .configuration import (
    AVAILABILITY_ZONE_1, AVAILABILITY_ZONE_2, AVAILABILITY_ZONE_3, 
    ROUTE_TABLE_1, ROUTE_TABLE_2, ROUTE_TABLE_3,
	SUBNET_ID_1, SUBNET_ID_2, SUBNET_ID_3, SHARED_SECURITY_GROUP_ID, VPC_ID, VPC_CIDR,
    S3_RAW_BUCKET, S3_CONFORMED_BUCKET, S3_PURPOSE_BUILT_BUCKET, S3_ACCESS_LOG_BUCKET, S3_KMS_KEY,
)

class ImportedBuckets():
     def __init__(
        self,
        stack: cdk.Stack,
        logical_id_suffix: str,
    ):
        """Add imported buckets from Infrastructure Cloudformation exports to provided stack

        Parameters
        ----------
        stack
            Stack in which to add buckets and use for mappings
        logical_id_suffix
            Suffix to differentiate construct IDs of imported resources
        """
        # Mapping should be a defined property in the calling stack
        s3_kms_key_parameter = cdk.Fn.import_value(stack.mappings[S3_KMS_KEY])
        self.s3_kms_key = kms.Key.from_key_arn(stack, f'ImportedKmsKey{logical_id_suffix}', s3_kms_key_parameter)

        access_logs_bucket_name = cdk.Fn.import_value(stack.mappings[S3_ACCESS_LOG_BUCKET])
        self.access_logs = s3.Bucket.from_bucket_attributes(
            stack,
            f'ImportedAccessLogsBucket{logical_id_suffix}',
            bucket_name=access_logs_bucket_name
        )
        raw_bucket_name = cdk.Fn.import_value(stack.mappings[S3_RAW_BUCKET])
        self.raw = s3.Bucket.from_bucket_name(
            stack,
            id=f'ImportedCollectBucket{logical_id_suffix}',
            bucket_name=raw_bucket_name
        )
        conformed_bucket_name = cdk.Fn.import_value(stack.mappings[S3_CONFORMED_BUCKET])
        self.conformed = s3.Bucket.from_bucket_name(
            stack,
            id=f'ImportedCleanseBucket{logical_id_suffix}',
            bucket_name=conformed_bucket_name
        )
        purposebuilt_bucket_name = cdk.Fn.import_value(stack.mappings[S3_PURPOSE_BUILT_BUCKET])
        self.purposebuilt = s3.Bucket.from_bucket_name(
            stack,
            id=f'ImportedConsumeBucket{logical_id_suffix}',
            bucket_name=purposebuilt_bucket_name
        )

class ImportedVpc():
    def __init__(
        self,
        stack: cdk.Stack,
        logical_id_suffix: str,
    ):
        """Add imported VPC and related resources from Infrastructure Cloudformation exports
        to provided stack

        Parameters
        ----------
        stack
            Stack in which to add VPC and and related constructs, and use for mappings
        logical_id_suffix
            Suffix to differentiate construct IDs of imported resources
        """
        # Mapping should be a defined property in the calling stack
        shared_security_group_parameter = cdk.Fn.import_value(stack.mappings[SHARED_SECURITY_GROUP_ID])
        self.shared_security_group = ec2.SecurityGroup.from_security_group_id(
            stack,
            f'ImportedSecurityGroup{logical_id_suffix}',
            shared_security_group_parameter
        )

        self.subnets = []

        if VPC_CIDR not in stack.mappings:
            # Environment configuration does not have a VPC CIDR defined, so VPC exports are not expected
            return

        vpc_id = cdk.Fn.import_value(stack.mappings[VPC_ID])

        # VPC Stack from Infrastructure ensures that 3 AZs and associated valid outputs are always created
        vpc_azs = []
        for az_number in range(3):
            az_mapping_element = globals()[f'AVAILABILITY_ZONE_{az_number + 1}']
            vpc_azs.append(cdk.Fn.import_value(stack.mappings[az_mapping_element]))

        vpc_subnet_ids = []
        for subnet_number in range(3):
            subnet_mapping_element = globals()[f'SUBNET_ID_{subnet_number + 1}']
            vpc_subnet_ids.append(cdk.Fn.import_value(stack.mappings[subnet_mapping_element]))

        vpc_route_tables = []
        for rt_number in range(3):
            rt_mapping_element = globals()[f'ROUTE_TABLE_{rt_number + 1}']
            vpc_route_tables.append(cdk.Fn.import_value(stack.mappings[rt_mapping_element]))

        # Manually construct the VPC because it lives in the target
        # account, not the Deployment account where the synth is ran
        self.vpc = ec2.Vpc.from_vpc_attributes(
            stack,
            f'ImportedVpc{logical_id_suffix}',
            vpc_id=vpc_id,
            availability_zones=vpc_azs,
            private_subnet_ids=vpc_subnet_ids,
            private_subnet_route_table_ids=vpc_route_tables,
        )

        # Some resources require specific subnets to access the VPC, such a Glue Connections
        for az_number in range(len(vpc_azs)):
            self.subnets.append(ec2.Subnet.from_subnet_attributes(
                stack,
                f'ImportedSubnet{logical_id_suffix}{az_number + 1}',
                subnet_id=vpc_subnet_ids[az_number],
                availability_zone=vpc_azs[az_number],
                route_table_id=vpc_route_tables[az_number]
            ))
lib/stack_import_helper.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import json
import boto3
import botocore
import os
import logging
import uuid
import re
from datetime import datetime
import dateutil.tz
from dateutil import parser as dateparser
from urllib.parse import unquote_plus

# Logger initiation
logger = logging.getLogger()
logger.setLevel(logging.INFO)


def record_etl_job_run(
        audit_table_name: str,
        sfn_arn: str,
        execution_id: str,
        execution_name: str,
        execution_input: str,
        principal_id: str,
        source_ipaddress: str,
    ):
    """
    Function to insert entry in dynamodb table for audit trail

    table_name
        DynamoDB table name to use for job audit data
    sfn_arn
        AWS ARN of state machine being triggered
    execution_id
        State Machine execution unique ID
    execution_name
        Name of Step Functions State Machine being triggered
    execution_input
        State Machine execution input parameters
    principal_id
        S3 event principalId from userIdentity record
    source_ipaddress
        S3 event sourceIPAddress from requestParameters record
    """
    logger.info('record_etl_job_run() called')
    now = datetime.now(tz=dateutil.tz.gettz('UTC'))
    record_time = now.strftime('%Y-%m-%d %H:%M:%S.%f')
    item = {}
    item['execution_id'] = execution_id
    item['sfn_execution_name'] = execution_name
    item['sfn_arn'] = sfn_arn
    item['sfn_input'] = execution_input
    item['job_latest_status'] = 'STARTED'
    item['job_start_date'] = record_time
    item['job_last_updated_timestamp'] = record_time
    item['principal_id'] = principal_id
    item['source_ipaddress'] = source_ipaddress

    dynamo_client = boto3.resource('dynamodb')
    try:
        table = dynamo_client.Table(audit_table_name)
        table.put_item(Item=item)
    except botocore.exceptions.ClientError as error:
        logger.error(f'record_etl_job_run() DynamoDB put_item failed: {error}')
        raise error
    logger.info('record_etl_job_run() execution completed')
    print('Job audit record insert completed successfully')


def lambda_handler(event: dict, context: dict) -> dict:
    """Lambda function's entry point. This function receives a PutObject event
    from S3, prepares Step Function State machine inputs, logs the job start in
    DynamoDB table, and initiates the State Machine

    Parameters
    ----------
    event
        S3 PutObject event dictionary
    context
        Lambda context dictionary

    Returns
    -------
    dict
        Lambda result dictionary
    """
    sfn_client = boto3.client('stepfunctions')
    sfn_arn = os.environ['SFN_STATE_MACHINE_ARN']
    audit_table_name = os.environ['DYNAMODB_TABLE_NAME']

    print(event)
    lambda_message = event['Records'][0]
    source_bucket_name = lambda_message['s3']['bucket']['name']
    object_full_path = unquote_plus(lambda_message['s3']['object']['key'])
    event_time = dateparser.parse(lambda_message['eventTime'])
    principal_id = lambda_message['userIdentity']['principalId']
    source_ipaddress = lambda_message['requestParameters']['sourceIPAddress']

    if object_full_path[-1] == '/':
        logger.error(f'Ignoring folder creation: {object_full_path}')
        return {
            'statusCode': 400,
            'body': json.dumps(f'Received PutObject for folder; cannot process')
        }

    try:
        # Bucket/key format: s3://<bucketname>/<source_system_name>/<table_name>
        object_file_dir = os.path.dirname(object_full_path)
        object_base_file_name = os.path.basename(object_full_path)
        # First object/folder name after bucket name will be used as source system/database name
        object_source_system_name = object_file_dir.split('/')[0]
        # Second object/folder name after bucket name will be used as table name
        object_table_name = object_file_dir.split('/')[1]
    except IndexError:
        logger.error(f'File object {object_full_path} cannot be processed without 2 levels of directory structure')
        return {
            'statusCode': 400,
            'body': json.dumps(f'File object {object_full_path} ignored due to unexpected naming convention')
        }

    logger.info(f'Bucket: {source_bucket_name}')
    logger.info(f'Key: {object_full_path}')
    logger.info(f'Source System Name: {object_source_system_name}')
    logger.info(f'Table Name: {object_table_name}')
    logger.info(f'File Path: {object_file_dir}')
    logger.info(f'File Base Name: {object_base_file_name}')
    logger.info(f'State Machine ARN: {sfn_arn}')

    logger.info(f'Date/time used for partitioning: {event_time}')
    p_year = event_time.strftime('%Y')
    p_month = event_time.strftime('%m')
    p_day = event_time.strftime('%d')
    logger.info(f'Partition Year (p_year): {p_year}')
    logger.info(f'Partition Month (p_month): {p_month}')
    logger.info(f'Partition Day (p_day): {p_day}')

    # Ensure state machine execution logging by following CloudWatch log group name constraints
    safe_object_base_file_name = re.sub('[^a-zA-Z0-9_-]', '', object_base_file_name)
    # Max length of execution name is 80 characters, and 21 will be added, so truncate to 59
    execution_name = f"{safe_object_base_file_name[:59]}-{event_time.strftime('%Y%m%d%H%M%S%f')}"
    logger.info(f'Step Function Execution Name: {execution_name}')
    print('Executing step function')
    execution_id = str(uuid.uuid4())
    execution_input = json.dumps(
        {
            'target_database_name': object_source_system_name,
            'source_bucketname': source_bucket_name,
            'source_key': object_file_dir, 
            'base_file_name': object_base_file_name,
            'p_year': p_year,
            'p_month': p_month,
            'p_day': p_day,
            'table_name': object_table_name,
            'execution_id': execution_id,
        }
    )
    print(f'SFN Input: {execution_input}')
    try:
        sfn_response = sfn_client.start_execution(
            stateMachineArn=sfn_arn,
            name=execution_name,
            input=execution_input
        )
        print(f'SFN Reponse: {sfn_response}')
    except botocore.exceptions.ClientError as error:
        logger.error(f'Step function client process failed: {error}')
        raise error

    record_etl_job_run(audit_table_name, sfn_arn, execution_id, execution_name, execution_input, principal_id, source_ipaddress)

    return {
        'statusCode': 200,
        'body': json.dumps('Step function triggered successfully')
    }
lib/state_machine_trigger/lambda_handler.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import os
import aws_cdk as cdk
from constructs import Construct
import aws_cdk.aws_dynamodb as dynamodb
import aws_cdk.aws_logs as logs
import aws_cdk.aws_glue as glue
import aws_cdk.aws_iam as iam
import aws_cdk.aws_lambda as _lambda
import aws_cdk.aws_s3 as s3
import aws_cdk.aws_s3_notifications as s3_notifications
import aws_cdk.aws_sns as sns
import aws_cdk.aws_stepfunctions as stepfunctions
import aws_cdk.aws_stepfunctions_tasks as stepfunctions_tasks
from cdk_nag import NagSuppressions

from .stack_import_helper import ImportedBuckets
from .configuration import (
    DEV, PROD, TEST, STATE_MACHINE, NOTIFICATION_TOPIC,
    get_logical_id_prefix, get_resource_name_prefix, get_environment_configuration,
)


class StepFunctionsStack(cdk.Stack):
    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        target_environment: str,
        collect_to_cleanse_job: glue.CfnJob,
        cleanse_to_consume_job: glue.CfnJob,
        job_audit_table: dynamodb.Table,
        **kwargs
    ):
        """CloudFormation stack to create Step Functions, Lambdas, and SNS Topics

        Parameters
        ----------
        scope
            Parent of this stack, usually an App or a Stage, but could be any construct
        construct_id
            The construct ID of this stack; if stackName is not explicitly defined,
            this ID (and any parent IDs) will be used to determine the physical ID of the stack
        target_environment
            The target environment for stacks in the deploy stage
        collect_to_cleanse_job
            Collect to Cleanse Glue job construct to invoke
        cleanse_to_consume_job
            Cleanse to Consume Glue job construct to invoke
        job_audit_table
            The DynamoDB Table construct for storing Job Audit results
        kwargs: optional
            Optional keyword arguments to pass up to parent Stack class
        """
        super().__init__(scope, construct_id, **kwargs)

        self.target_environment = target_environment
        self.mappings = get_environment_configuration(target_environment)
        self.logical_id_prefix = get_logical_id_prefix()
        self.resource_name_prefix = get_resource_name_prefix()
        if (target_environment == PROD or target_environment == TEST):
            self.removal_policy = cdk.RemovalPolicy.RETAIN
            self.log_retention = logs.RetentionDays.SIX_MONTHS
        else:
            self.removal_policy = cdk.RemovalPolicy.DESTROY
            self.log_retention = logs.RetentionDays.ONE_MONTH

        buckets = ImportedBuckets(self, logical_id_suffix='StepFunctionsStack')

        cloudwatch_step_function_log_group = logs.LogGroup(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlStateMachineLogGroup',
            retention=self.log_retention,
            removal_policy=self.removal_policy,
        )

        notification_topic = sns.Topic(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlNotificationTopic',
            topic_name=f'{target_environment.lower()}-{self.resource_name_prefix}-etl-notification-topic',
            display_name='InsuranceLake ETL Notifications Topic',
            master_key=buckets.s3_kms_key,
        )

        status_function = self.lambda_function_for_etl(
            logical_id_suffix='EtlStatusUpdate',
            resource_name_suffix='etl-status-update',
            function_description='ETL Step Functions workflow triggered handler to update DynamoDB in case of success or failure',
            lambda_code_relative_path='etl_job_auditor',
            lambda_environment={
                'DYNAMODB_TABLE_NAME': job_audit_table.table_name,
            },
            job_audit_table=job_audit_table,
        )

        fail_state = stepfunctions.Fail(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlFailedState',
            cause='End of failure path, root cause from Glue task step',
            error='Error'
        )
        success_state = stepfunctions.Succeed(self, f'{target_environment}{self.logical_id_prefix}EtlSucceededState')

        failure_function_task = stepfunctions_tasks.LambdaInvoke(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlFailureStatusUpdateTask',
            lambda_function=status_function,
            retry_on_service_exceptions=True,
            payload=stepfunctions.TaskInput.from_object({'Input.$': '$'}),
            result_path='$.status_update_result',
            output_path='$',
        )
        failure_notification_task = stepfunctions_tasks.SnsPublish(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlFailurePublishTask',
            topic=notification_topic,
            subject='Job Failed',
            message=stepfunctions.TaskInput.from_json_path_at('$')
        )
        failure_function_task.next(failure_notification_task)
        failure_notification_task.next(fail_state)

        success_function_task = stepfunctions_tasks.LambdaInvoke(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlSuccessStatusUpdateTask',
            lambda_function=status_function,
            retry_on_service_exceptions=True,
            payload=stepfunctions.TaskInput.from_object({'Input.$': '$'}),
            result_path='$.status_update_result',
            output_path='$',
        )
        success_task = stepfunctions_tasks.SnsPublish(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlSuccessPublishTask',
            topic=notification_topic,
            subject='Job Completed',
            message=stepfunctions.TaskInput.from_json_path_at('$')
        )
        success_function_task.next(success_task)
        success_task.next(success_state)

        glue_collect_task = stepfunctions_tasks.GlueStartJobRun(
            self,
            f'{target_environment}{self.logical_id_prefix}GlueCollectJobTask',
            glue_job_name=collect_to_cleanse_job.name,
            comment='Collect to Cleanse data load and transform',
            arguments=stepfunctions.TaskInput.from_object({
                # These arguments overlay base arguments from the Glue Job definition
                '--state_machine_name.$': '$$.StateMachine.Name',
                '--execution_id.$': '$.execution_id',
                '--source_key.$': '$.source_key',
                '--target_database_name.$': '$.target_database_name',
                '--table_name.$': '$.table_name',
                '--base_file_name.$': '$.base_file_name',
                '--p_year.$': '$.p_year',
                '--p_month.$': '$.p_month',
                '--p_day.$': '$.p_day',
            }),
            integration_pattern=stepfunctions.IntegrationPattern.RUN_JOB,
            result_path='$.taskresult',
            output_path='$',
        )
        glue_collect_task.add_catch(failure_function_task, result_path='$.taskresult')

        glue_cleanse_task = stepfunctions_tasks.GlueStartJobRun(
            self,
            f'{target_environment}{self.logical_id_prefix}GlueCleanseJobTask',
            glue_job_name=cleanse_to_consume_job.name,
            comment='Cleanse to Consume data load and transform',
            arguments=stepfunctions.TaskInput.from_object({
                # These arguments overlay base arguments from the Glue Job definition
                '--state_machine_name.$': '$$.StateMachine.Name',
                '--execution_id.$': '$.execution_id',
                '--database_name_prefix.$': '$.target_database_name',
                '--table_name.$': '$.table_name',
                '--base_file_name.$': '$.base_file_name',
                '--p_year.$': '$.p_year',
                '--p_month.$': '$.p_month',
                '--p_day.$': '$.p_day',
            }),
            integration_pattern=stepfunctions.IntegrationPattern.RUN_JOB,
            result_path='$.taskresult',
            output_path='$',
        )
        glue_cleanse_task.add_catch(failure_function_task, result_path='$.taskresult')

        machine_definition = glue_collect_task.next(
            glue_cleanse_task.next(
                success_function_task
            )
        )


        machine = stepfunctions.StateMachine(
            self,
            f'{target_environment}{self.logical_id_prefix}EtlStateMachine',
            state_machine_name=f'{target_environment.lower()}-{self.resource_name_prefix}-etl-state-machine',
            tracing_enabled=True,
            definition=machine_definition,
            logs=stepfunctions.LogOptions(
                destination=cloudwatch_step_function_log_group,
                level=stepfunctions.LogLevel.ALL,
            ),
        )

        # State Machine needs access to KMS key to publish topics
        machine.add_to_role_policy(
            iam.PolicyStatement(
                actions=[
                    'kms:GenerateDataKey',
                ],
                resources=[buckets.s3_kms_key.key_arn],
            )
        )

        # The State Machine generated role is otherwise specific and uses only the needed actions
        # and resources for the Lambda functions, Glue Jobs, and SNS Topic. And builds these
        # policies automatically based on the state machine definition.
        NagSuppressions.add_resource_suppressions(machine.role, [
            {
                'id': 'AwsSolutions-IAM5',
                'reason': 'State Machine generated role uses wildcard permissions for CloudWatch and X-Ray'
            },
        ], apply_to_children=True)


        trigger_function = self.lambda_function_for_etl(
            logical_id_suffix='EtlTrigger',
            resource_name_suffix='etl-trigger',
            function_description='Collect S3 Bucket triggered handler to trigger Step Functions workflow',
            lambda_code_relative_path='state_machine_trigger',
            lambda_environment={
                'DYNAMODB_TABLE_NAME': job_audit_table.table_name,
                'SFN_STATE_MACHINE_ARN': machine.state_machine_arn,
            },
            job_audit_table=job_audit_table,
            state_machine=machine,
        )

        # Will create CustomResource and Lambda to add event handler to imported bucket
        # TODO: Apply rotation and retention policies to Custom Resource Lambda log group
        buckets.raw.add_event_notification(
            s3.EventType.OBJECT_CREATED,
            s3_notifications.LambdaDestination(trigger_function),
        )

        NagSuppressions.add_resource_suppressions(self, [
            {
                'id': 'AwsSolutions-IAM4',
                'reason': 'Bucket Notification CustomResource used only during stack deployment and deletion'
            },
            {
                'id': 'AwsSolutions-IAM5',
                'reason': 'Bucket Notification CustomResource used only during stack deployment and deletion'
            },
        ], apply_to_children=True)

        cdk.CfnOutput(
            self,
            f'{target_environment}{self.logical_id_prefix}StateMachineName',
            value=machine.state_machine_name,
            export_name=self.mappings[STATE_MACHINE]
        )

        cdk.CfnOutput(
            self,
            f'{target_environment}{self.logical_id_prefix}SnsTopicName',
            value=notification_topic.topic_name,
            export_name=self.mappings[NOTIFICATION_TOPIC]
        )


    def get_lambda_role(
        self,
        logical_id_suffix: str,
        resource_name_suffix: str,
        log_group: logs.LogGroup,
        job_audit_table: dynamodb.Table,
        state_machine: stepfunctions.StateMachine = None,
    ) -> iam.Role:
        """Creates the role used during Lambda execution

        Parameters
        ----------
        logical_id_suffix
            Suffix to append to Logical ID of IAM Role to differentiate between Lambdas
        resource_name_suffix
            Suffix to append to IAM Role name to differentiate between Lambdas
        log_group
            Log group for the Lambda function
        job_audit_table
            The DynamoDB Table construct for storing Job Audit results
        state_machine: optional
            Optional Step Function StateMachine for execution (if used by the Lambda)

        Returns
        -------
        iam.Role
            The IAM role that was created
        """
        policies = {
            'CloudWatchLogAccess':
            iam.PolicyDocument(statements=[
                iam.PolicyStatement(
                    effect=iam.Effect.ALLOW,
                    actions=[
                        'logs:CreateLogStream',
                        'logs:PutLogEvents',
                    ],
                    resources=[log_group.log_group_arn]
                )
            ]),
            'JobAuditTableAccess':
            iam.PolicyDocument(statements=[
                iam.PolicyStatement(
                    effect=iam.Effect.ALLOW,
                    actions=[
                        'dynamodb:PutItem',
                        'dynamodb:GetItem',
                        'dynamodb:UpdateItem',
                    ],
                    resources=[job_audit_table.table_arn]
                )
            ]),
        }

        if state_machine is not None:
            policies.update({
                'StateMachineAccess':
                iam.PolicyDocument(statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        actions=[
                            'states:StartExecution',
                        ],
                        resources=[state_machine.state_machine_arn]
                    )
                ]),
            })

        return iam.Role(
            self,
            f'{self.target_environment}{self.logical_id_prefix}{logical_id_suffix}LambdaRole',
            description='Role for InsuranceLake ETL Lambda Functions',
            role_name=f'{self.target_environment.lower()}-{self.resource_name_prefix}-{self.region}-{resource_name_suffix}-lambda',
            assumed_by=iam.ServicePrincipal('lambda.amazonaws.com'),
            inline_policies=policies,
        )


    def lambda_function_for_etl(
        self,
        logical_id_suffix: str,
        resource_name_suffix: str,
        function_description: str,
        lambda_code_relative_path: str,
        job_audit_table: dynamodb.Table,
        state_machine: stepfunctions.StateMachine = None,
        lambda_environment: dict = {},
    ) -> _lambda.Function:
        """Creates a Lambda Function to support the ETL process

        Parameters
        ----------
        logical_id_suffix
            Suffix to append to Logical ID of IAM Role to differentiate between Lambdas
        resource_name_suffix
            Suffix to append to IAM Role name to differentiate between Lambdas
        function_description
            Description of Lambda to use for deployment
        lambda_code_relative_path
            Relative local path to the Lambda source code
        job_audit_table
            The DynamoDB Table construct for storing Job Audit results
        state_machine: optional
            Optional Step Function StateMachine for execution (if used by the Lambda)
        lambda_environment: optional
            Optional environment variable key, value pairs to pass to Lambda

        Returns
        -------
        lambda.Function
            The Lambda Function that was created
        """
        lambda_function_name = f'{self.target_environment.lower()}-{self.resource_name_prefix}-{resource_name_suffix}'

        # Associate the Log group to the Lambda function by giving it the same name
        cloudwatch_log_group = logs.LogGroup(
            self,
            f'{self.target_environment}{self.logical_id_prefix}{logical_id_suffix}LambdaLogGroup',
            log_group_name=f'/aws/lambda/{lambda_function_name}',
            retention=self.log_retention,
            removal_policy=self.removal_policy,
        )

        return _lambda.Function(
            self,
            f'{self.target_environment}{self.logical_id_prefix}{logical_id_suffix}',
            function_name=lambda_function_name,
            description=function_description,
            runtime=_lambda.Runtime.PYTHON_3_11,
            handler='lambda_handler.lambda_handler',
            code=_lambda.Code.from_asset(f'{os.path.dirname(__file__)}/{lambda_code_relative_path}'),
            architecture=_lambda.Architecture.ARM_64,
            environment=lambda_environment,
            role=self.get_lambda_role(
                logical_id_suffix,
                resource_name_suffix,
                cloudwatch_log_group,
                job_audit_table,
                state_machine
            ),
        )
lib/step_functions_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import aws_cdk as cdk

from .configuration import (
    get_logical_id_prefix, get_resource_name_prefix, get_all_configurations
)


COST_CENTER = 'COST_CENTER'
TAG_ENVIRONMENT = 'TAG_ENVIRONMENT'
TEAM = 'TEAM'
APPLICATION = 'APPLICATION'


def tag(stack: cdk.Stack, target_environment: str):
    """Adds a tag to all constructs in the stack

    stack
        CDK stack construct to tag
    target_environment
        The environment the stack is deployed to (for tag values)
    """
    cdk.Tags.of(stack).add(*get_tag(COST_CENTER, target_environment))
    cdk.Tags.of(stack).add(*get_tag(TAG_ENVIRONMENT, target_environment))
    cdk.Tags.of(stack).add(*get_tag(TEAM, target_environment))
    cdk.Tags.of(stack).add(*get_tag(APPLICATION, target_environment))


def get_tag(tag_name: str, target_environment: str) -> dict:
    """Get a tag for a given parameter and target environment.

    tag_name
        The name of the tag (must exist in tag_map)
    target_environment
        The environment the tag is applied to (for tag values)

    Raises
    ------
    AttributeError
        If target environment or tag name is not present in the tag_map below

    Returns
    -------
    dict
        key, value pair for each tag and tag value
    """
    mapping = get_all_configurations()
    if target_environment not in mapping:
        raise AttributeError(f'Target environment {target_environment} not found in environment configurations')

    logical_id_prefix = get_logical_id_prefix()
    resource_name_prefix = get_resource_name_prefix()
    tag_map = {
        COST_CENTER: [
            f'{resource_name_prefix}:cost-center',
            f'{logical_id_prefix}Etl',
        ],
        TAG_ENVIRONMENT: [
            f'{resource_name_prefix}:environment',
            target_environment,
        ],
        TEAM: [
            f'{resource_name_prefix}:team',
            f'{logical_id_prefix}Admin',
        ],
        APPLICATION: [
            f'{resource_name_prefix}:application',
            f'{logical_id_prefix}Etl',
        ],
    }
    if tag_name not in tag_map:
        raise AttributeError(f'Tag map does not contain a key/value for {tag_name}')

    return tag_map[tag_name]
lib/tagging.py
aws-samples/aws-insurancelake-etl
#!/usr/bin/env python3
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
"""
Script to convert 2 column CSV lookup data to JSON structure suitable for loading into
DynamoDB lookup_value table with load_dynamodb_lookup_table.py. Expects CSV input from stdin,
outputs JSON to stdout.

Usage: csv_to_json.py column_name < input.csv > output.json

Options:
  -h, --help           show this help message and exit
  column_name          Value to use for column_name in the DynamoDB JSON
  input.csv            2 column CSV input file (other columns will be ignored)
  output.json          JSON output suitable for load_dynamodb_lookup_table.py

"""
import sys
import csv
import json
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('column_name',
	help="Value to use for column_name in the DynamoDB JSON")
args = parser.parse_args()

data = sys.stdin.readlines()
csvreader = csv.reader(data)
next(csvreader)
data_dict = { args.column_name: { row[0]: row[1] for row in csvreader } }
print(json.dumps(data_dict))
resources/csv_to_json.py
aws-samples/aws-insurancelake-etl
#!/usr/bin/env python3
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
"""
Script to clean-up artifacts left from running the pipeline and leave anything deployed by CDK

*** THIS SCRIPT DELETES DATA AND RESOURCES THAT CANNOT BE RECOVERED ***

Usage: etl_cleanup.py [-h] [--mode reallydelete]

Options:
  -h, --help           show this help message and exit
  --mode reallydelete  Must specify to actually delete anything, otherwise runs dry-run

Ensure you have the correct account access and tokens in your environment, as well as the default
region specified in the AWS_DEFAULT_REGION environment variable.

Script is hard-coded to only work on the InsuranceLake Development environment.

Running without --reallydelete will return a count of objects only (except S3 buckets, which will
report 0).

"""
import boto3
import argparse
import re
import sys
import os

sys.path.append(os.path.dirname(__file__) + '/../lib')
from configuration import (
    S3_ACCESS_LOG_BUCKET, S3_CONFORMED_BUCKET, S3_PURPOSE_BUILT_BUCKET, S3_RAW_BUCKET, DEV, TEST, PROD,
    get_environment_configuration, get_resource_name_prefix, get_logical_id_prefix
)

# Strongly recommend only using this on the DEV environment
environment = DEV


parser = argparse.ArgumentParser()
parser.add_argument('--mode', dest='mode', metavar='reallydelete', required = False,
	choices=[ 'reallydelete' ],
	help='Must specify to actually delete anything, otherwise runs dry-run')
args = parser.parse_args()
reallydelete = args.mode == 'reallydelete'

if not reallydelete:
	print('**** DRY RUN MODE ENABLED. NO REAL DELETING WILL OCCUR ****')

mappings = get_environment_configuration(environment)
resource_prefix = get_resource_name_prefix()
logical_id_prefix = get_logical_id_prefix()

bucket_exports = [
   mappings[S3_ACCESS_LOG_BUCKET], 
   mappings[S3_CONFORMED_BUCKET], 
   mappings[S3_PURPOSE_BUILT_BUCKET],
   mappings[S3_RAW_BUCKET],
]

# Gather Cloudformation outputs for existing stack
buckets = []
dynamodb_tables = []
cf = boto3.client('cloudformation')
response = cf.list_exports()
while True:
	for export in response['Exports']:
		if export['Name'] in bucket_exports:
			buckets.append(export['Value'])
		match = re.search('ExportsOutputRef\w+Table', export['Name'])
		if match:
			dynamodb_tables.append(export['Value'])
	if 'NextToken' in response:
		response = cf.list_exports(nextToken=response['NextToken'])
	else:
		break


if not buckets and not dynamodb_tables:
	print('No Cloudformation exports found for S3 buckets or DynamoDB tables; something is probably wrong (check account and region)')
	exit(1)


print (f'Emptying Collect/Cleanse/Consume S3 buckets...')
s3 = boto3.resource('s3')
for bucket_name in buckets:
	s3_bucket = s3.Bucket(bucket_name)
	print(f'Emptying S3 bucket: {bucket_name}: ', end='')
	bucket_versioning = s3.BucketVersioning(bucket_name)
	versions_deleted = 0
	objects_deleted = 0
	if bucket_versioning.status == 'Enabled':
		if reallydelete:
			deleted = s3_bucket.object_versions.delete()
			if deleted:
				versions_deleted = len(deleted[0]['Deleted'])
	else:
		if reallydelete:
			deleted = s3_bucket.objects.all().delete()
			if deleted:
				objects_deleted = len(deleted[0]['Deleted'])
	print(f'{objects_deleted} Objects deleted, and {versions_deleted} Object Versions deleted')


print('Deleting Cleanse/Consume Glue Data Catalogs...')
glue = boto3.client('glue')
response_databases = glue.get_databases()
db_names = {}
while True:
	if 'DatabaseList' in response_databases:
		for database in response_databases['DatabaseList']:
			response_tables = glue.get_tables(DatabaseName=database['Name'])
			while True:
				if 'TableList' in response_tables:
					for table in response_tables['TableList']:
						match = re.search('^s3://([^\/]+)', table['StorageDescriptor']['Location'])
						if match and match.group(1) in buckets:
							# Table location matches an InsuranceLake S3 bucket name
							# Use a dict key to remove duplicates
							db_names[table['DatabaseName']] = 1
				if 'NextToken' in response_tables:
					response_tables = glue.get_tables(
						DatabaseName=database['Name'],
						nextToken=response_tables['NextToken']
					)
				else:
					break
	if 'NextToken' in response_databases:
		response_databases = glue.get_databases(nextToken=response_databases['NextToken'])
	else:
		break

for database in db_names.keys():
	print(f'Deleting Glue Catalog Database: {database}')
	if reallydelete:
		# Deleting the DB deletes all tables in the DB
		glue.delete_database(Name=database)


print('Deleting Glue Cloudwatch Log Groups...')
glue_log_groups = ['/aws-glue/jobs/logs-v2', '/aws-glue/jobs/output', '/aws-glue/jobs/error']
logs = boto3.client('logs')
for log_group_name in glue_log_groups:
	try:
		print(f'Deleting Log Group: {log_group_name}')
		if reallydelete:
			logs.delete_log_group(logGroupName=log_group_name)
	except logs.exceptions.ResourceNotFoundException as error:
		pass

print('Emptying Cloudwatch Log Groups for Lambda functions. State Machine, and VPC Flow Log...')
logs_to_empty = [
	f'/aws/lambda/{environment.lower()}-{resource_prefix}-etl-status-update',
	f'/aws/lambda/{environment.lower()}-{resource_prefix}-etl-trigger',
]

log_groups_response = logs.describe_log_groups()
if 'NextToken' in log_groups_response:
	print('Not all log groups were parsed. Support for multiple pages is not implemented.')
for log_group in log_groups_response['logGroups']:
	if re.search(f'{environment}-{logical_id_prefix}\w+-\w+StateMachineLogGroup', log_group['logGroupName']):
		logs_to_empty.append(log_group['logGroupName'])
	if re.search(f'{environment}-{logical_id_prefix}\w+-\w+VpcFlowLogGroup', log_group['logGroupName']):
		logs_to_empty.append(log_group['logGroupName'])
		print("NOTE: VPC Flow Logs are created continuously; you will see some log streams immediately after deleting")

for log_group_name in logs_to_empty:
	print(f'Emptying Log Group: {log_group_name}: ', end='')
	log_streams_deleted = 0
	log_streams_response = logs.describe_log_streams(logGroupName=log_group_name)
	while True:
		for stream in log_streams_response['logStreams']:
			if reallydelete:
				logs.delete_log_stream(logGroupName=log_group_name, logStreamName=stream['logStreamName'])
			log_streams_deleted += 1
		if 'nextToken' in log_streams_response:
			log_streams_response = logs.describe_log_streams(
				logGroupName=log_group_name, 
				nextToken=log_streams_response['nextToken']
			)
		else:
			break
	print(f'{log_streams_deleted} Log Streams deleted')


print('Deleting all items in DynamoDB tables...')
dynamodb = boto3.resource('dynamodb')
for table_name in dynamodb_tables:
	table = dynamodb.Table(table_name)
	key_names = [ key['AttributeName'] for key in table.key_schema ]
	print(f'Emptying DynamoDB Table: {table_name}: ', end='')
	flag = False
	scan = table.scan()
	while True:
		items_deleted = 0
		with table.batch_writer() as batch:
			for each in scan['Items']:
				item_key = { key_name: each[key_name] for key_name in key_names }
				if reallydelete:
					batch.delete_item(Key=item_key)
				items_deleted += 1
			if 'LastEvaluatedKey' in scan:
				scan = table.scan(ExclusiveStartKey=scan['LastEvaluatedKey'])
			else:
				break
	print(f'{items_deleted} Items deleted')
resources/etl_cleanup.py
aws-samples/aws-insurancelake-etl
#!/usr/bin/env python3
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
"""
Script to populate the DynamoDB table for the lookup transform with
lookup data. Requires a JSON file with lookup data in key: value
format for each column on which the transform is being applied.
"""
import boto3
import argparse
import json

parser = argparse.ArgumentParser()
parser.add_argument('source_system',
	help="Name of source data system that will be used to match against uploaded file path to the Collect bucket")
parser.add_argument('table_name',
	help="Name of DynamoDB lookup table to be populated by the script")
parser.add_argument('data_file',
	help="Path/file name of JSON file to use for populating the DynamoDB table")
args = parser.parse_args()
 
# Read data from JSON in the format:
# {
#	"column_name1": { "lookup_value1": "lookedup_value1", "lookup_value2": "lookedup_value2", ... },
#	"column_name2": { ... }
# }
with open(args.data_file) as lookup_data_file:
	lookup_data = json.load(lookup_data_file)

# Prepare DynamoDB data in the form of the table schema
dynamodb_write_data = [
	{
		'source_system': args.source_system,
		'column_name': lookup,
		'lookup_data': json.dumps(lookup_data[lookup]),
	}
	for lookup in lookup_data
]

# Write data to DynamoDB table
dynamodb = boto3.resource('dynamodb')
lookup_table = dynamodb.Table(args.table_name)

with lookup_table.batch_writer() as batch:
	for data_row in dynamodb_write_data:
		response = batch.put_item(data_row)
resources/load_dynamodb_lookup_table.py
aws-samples/aws-insurancelake-etl
#!/usr/bin/env python3
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
"""
Script to populate the DynamoDB table for the multilookup transform
with lookup data. Requires a CSV file with lookup data and return
values. Script requires lookup columns to be specified. All other
columns will assumed to be return values. Only return value
columns will be written to the DynamoDB table separately. The lookup
columns will be combined to a single lookup_item key.

Example usage:
	./load_dynamodb_multilookup_table.py dev-insurancelake-etl-multi-lookup lookups.csv PolicyData-LOBCoverage originalprogram originalcoverage
"""
import boto3
import argparse
import csv

parser = argparse.ArgumentParser()
parser.add_argument('table_name',
	help='Name of DynamoDB multilookup table to be populated by the script')
parser.add_argument('data_file',
	help='Path/file name of CSV file to use for populating the DynamoDB table')
parser.add_argument('lookup_group',
	help='Value to use for lookup_group in DynamoDB table (grouping of lookup_items)')
parser.add_argument('lookup_columns', nargs='+',
	help='Lookup columns to use for combined lookup_item value')
args = parser.parse_args()

with open(args.data_file, mode="r", encoding="utf-8-sig") as lookup_data_file:
	lookup_data_file = csv.DictReader(lookup_data_file)

	dynamodb_write_data = []
	for row in lookup_data_file:
		# lookup_columns are used to create the single lookup key and removed from the row
		lookup_item = '-'.join([row.pop(lookup_column) for lookup_column in args.lookup_columns])

		# Remaining columns are considered all return columns
		row_to_write = { return_column: row[return_column] for return_column in row }

		# Set partition key and sort key
		row_to_write.update({
			'lookup_group': args.lookup_group,
			'lookup_item': lookup_item,
			
		})

		dynamodb_write_data.append(row_to_write)

# Write data to DynamoDB table
dynamodb = boto3.resource('dynamodb')
lookup_table = dynamodb.Table(args.table_name)

with lookup_table.batch_writer() as batch:
	for data_row in dynamodb_write_data:
		response = batch.put_item(data_row)

print(f'{len(dynamodb_write_data)} rows written to {args.table_name} for lookup_group {args.lookup_group}')
resources/load_dynamodb_multilookup_table.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import setuptools

with open("README.md") as fp:
    long_description = fp.read()

setuptools.setup(
    name="aws-insurancelake-etl",
    version="2.5.0",
    description="A CDK Python app for deploying ETL jobs that operate data pipelines for InsuranceLake in AWS",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/aws-samples/aws-insurancelake-etl",
    author="Cory Visi <cvisi@amazon.com>, Ratnadeep Bardhan Roy <rdbroy@amazon.com>, Isaiah Grant <igrant@2ndwatch.com>, Ravi Itha <itharav@amazon.com>, Zahid Muhammad Ali <zhidli@amazon.com>",
    packages=setuptools.find_packages(),
    install_requires=[
        "aws-cdk-lib>=2.80.0",
        "constructs>=10.1.0",
    ],
    python_requires=">=3.9",
    keywords='aws-insurancelake-etl aws cdk insurance datalake etl glue spark python pyspark',
    license='MIT-0',
    classifiers=[
        "Development Status :: 4 - Beta",
        "Intended Audience :: Developers",
        "Programming Language :: JavaScript",
        "Programming Language :: Python :: 3 :: Only",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Topic :: Software Development :: Code Generators",
        "Topic :: Utilities",
        "Typing :: Typed",
    ],
)
setup.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import datetime
from dateutil import tz

mock_account_id = 'notrealaccountid'
mock_region = 'us-east-1'

mock_athena_execute_query_response = {
    'QueryExecutionId': '27029f2d-83d9-4f74-9ec5-4c8a188a67d1',
    'ResponseMetadata': {
        'RequestId': 'c6dff4ba-ac1d-435b-bcb7-420e67ae8f9b',
        'HTTPStatusCode': 200,
        'HTTPHeaders': {'content-type': 'application/x-amz-json-1.1', 'date': 'Thu, 22 Sep 2022 14:29:30 GMT', 'x-amzn-requestid': 'c6dff4ba-ac1d-435b-bcb7-420e67ae8f9b', 'content-length': '59', 'connection': 'keep-alive'},
        'RetryAttempts': 0
    }}

mock_athena_get_query_execute_failed_response =  {
    'QueryExecution': {
        'QueryExecutionId': '308cb419-d6a9-4bcc-a01e-9efa702e6b94',
        'Query': 'CREATE OR REPLACE VIEW "test_view" AS \nSELECT\n  generatedate\nFROM\n  AwsDataCatalog.default.mytable',
        'StatementType': 'DDL',
        'ResultConfiguration': {
            'OutputLocation': 's3://fakebucket/results/308cb419-d6a9-4bcc-a01e-9efa702e6b94.txt'
        },
        'QueryExecutionContext': {'Database': 'mydb', 'Catalog': 'fakeaccount'},
        'Status': {
            'State': 'FAILED',
            'StateChangeReason': "line 3:3: Column 'generatedate' cannot be resolved",
            'SubmissionDateTime': datetime.datetime(2022, 9, 22, 15, 29, 51, 159000, tzinfo=tz.tzlocal()),
            'CompletionDateTime': datetime.datetime(2022, 9, 22, 15, 29, 51, 535000, tzinfo=tz.tzlocal())
        },
        'Statistics': {'EngineExecutionTimeInMillis': 246, 'DataScannedInBytes': 0, 'TotalExecutionTimeInMillis': 376, 'QueryQueueTimeInMillis': 102, 'ServiceProcessingTimeInMillis': 28},
        'WorkGroup': 'primary',
		'EngineVersion': {'SelectedEngineVersion': 'AUTO', 'EffectiveEngineVersion': 'Athena engine version 2'}
    },
    'ResponseMetadata': {
        'RequestId': 'ef15ebbf-2482-49ef-ae2d-9397e5efdaaa',
        'HTTPStatusCode': 200,
        'HTTPHeaders': {'content-type': 'application/x-amz-json-1.1', 'date': 'Thu, 22 Sep 2022 15:29:51 GMT', 'x-amzn-requestid': 'ef15ebbf-2482-49ef-ae2d-9397e5efdaaa', 'content-length': '1961', 'connection': 'keep-alive'},
        'RetryAttempts': 0
    }}


class mock_client_sts:

	@staticmethod
	def get_caller_identity():
		return { 'Account': mock_account_id }

class mock_client_sfn:

	@staticmethod
	def start_execution(stateMachineArn: str, name: str, input: str):
		return { 
			'executionArn': f'arn:aws:states:{mock_region}:{mock_account_id}:execution:test-state-machine:{name}-plusuniqueid',
			'startDate': '2023-01-01 23:00:00',
		}

# There's no way to use moto's athena mock to cause a failed query
class mock_client_athena:
	@staticmethod
	def start_query_execution(QueryExecutionContext: dict, QueryString: str, ResultConfiguration: dict):
		return mock_athena_execute_query_response

	@staticmethod
	def get_query_execution(QueryExecutionId: str):
		return mock_athena_get_query_execute_failed_response

def mock_boto3_client(client: str):
	if client == 'sts':
		return mock_client_sts
	elif client == 'stepfunctions':
		return mock_client_sfn
	elif client == 'athena':
		return mock_client_athena
	else:
		raise RuntimeError(f'boto3 client {client} requested from mock but not implemented')
test/boto_mocking_helper.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys
import os
from pyspark.sql import SparkSession
from pyspark.sql.dataframe import DataFrameWriter

mock_region = 'us-east-1'
mock_database_name = 'TestGlueCatalogDatabase'
mock_table_name = 'TestGlueCatalogTable'

mock_resource_prefix = 'unittest-insurancelake'
mock_temp_bucket = f'file:///tmp/{mock_resource_prefix}-temp-bucket'
mock_scripts_bucket = f'file:///tmp/{mock_resource_prefix}-etl-scripts-bucket'
mock_collect_bucket = f'file:///tmp/{mock_resource_prefix}-collect-bucket'
mock_cleanse_bucket = f'file:///tmp/{mock_resource_prefix}-cleanse-bucket'
mock_consume_bucket = f'file:///tmp/{mock_resource_prefix}-consume-bucket'


def write_local_file(path: str, filename: str, data: str):
    path = path.replace('file://', '')
    os.makedirs(path, exist_ok=True)
    with open(path + '/' + filename, 'w') as input_file:
        input_file.writelines(data)

def mock_spark_context():
    return spark.sparkContext

class mock_lineage:
    def __init__(self, args: list, uniid: str = None):
        # Not implemented/needed for mocks
        pass

    def update_lineage(self, df, dataset, operation, **kwargs):
        # Not implemented/needed for mocks
        pass

def mock_create_database(database_name: str):
    # Not implemented/needed for mocks
    pass

def mock_upsert_catalog_table(
        df: any, target_database: str, table_name: str, partition_keys: list, storage_location: str,
        allow_schema_change: str = 'strict'):
    # Not implemented/needed for mocks
    pass

def mock_put_s3_object(uri: str, data: any):
    # Not implemented/needed for mocks
    pass

def mock_spark_sql(self, query: str):
    # Simplified mock behavior designed to get past the sql statements
    if 'alter table' in query.lower():
        return None
    else:
        return spark.createDataFrame([], schema='year int, month int, day int, data string')

def mock_athena_execute_query(database: str, query: str, max_attempts: int = 15) -> str:
    return 'SUCCEEDED'

def mock_dataframe_saveastable(self, target_table: str, **kwargs):
    # Collapse saveAsTable calls to save(), skipping the catalog update
    storage_location = kwargs.pop('path', '')
    self.save(storage_location, **kwargs)

class mock_glue_job:
    def __init__(self, glue_job_module):
        self.glue_job_module = glue_job_module

    def __call__(self, func):
        def inner(monkeypatch, *args, **kwargs):
            monkeypatch.setattr(SparkSession, 'sql', mock_spark_sql)
            monkeypatch.setattr(DataFrameWriter, 'saveAsTable', mock_dataframe_saveastable)
            # Both scripts use these classes/functions
            monkeypatch.setattr(self.glue_job_module, 'SparkContext', mock_spark_context)
            monkeypatch.setattr(self.glue_job_module, 'DataLineageGenerator', mock_lineage)
            # Only collect_to_cleanse uses these functions
            if 'collect_to_cleanse' in self.glue_job_module.__name__:
                monkeypatch.setattr(self.glue_job_module, 'upsert_catalog_table', mock_upsert_catalog_table)
                monkeypatch.setattr(self.glue_job_module, 'put_s3_object', mock_put_s3_object)
            if 'cleanse_to_consume' in self.glue_job_module.__name__:
                monkeypatch.setattr(self.glue_job_module, 'create_database', mock_create_database)
            func(monkeypatch, *args, **kwargs)
        return inner


# Shared Spark session for all tests
if 'pyspark' in sys.modules:
    spark = SparkSession.builder.master('local[*]').appName('Unit-tests').getOrCreate()
test/glue_job_mocking_helper.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
"""
Stub module to allow successful import of Glue job test scripts so
that tests can be explicitly skipped when there is no Spark
environment; companion to glue_job_mocking_helper.py
"""
mock_region = ''
mock_database_name = ''
mock_table_name = ''
mock_resource_prefix = ''
mock_temp_bucket = ''
mock_scripts_bucket = ''
mock_collect_bucket = ''
mock_cleanse_bucket = ''
mock_consume_bucket = ''

def etl_cleanse_to_consume():
	# Stub function so that decorator call will succeed
	pass

def etl_collect_to_cleanse():
	# Stub function so that decorator call will succeed
	pass

def mock_glue_job(func):
	# Stub decorator used by unit tests
	def inner(*args):
		func()
	return inner
test/glue_job_mocking_helper_stub.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk
from aws_cdk.assertions import Template

from test.boto_mocking_helper import *
from lib.dynamodb_stack import DynamoDbStack
from lib.glue_stack import GlueStack
from lib.athena_helper_stack import AthenaHelperStack

import lib.configuration as configuration
from lib.configuration import (
	DEV, PROD, TEST
)


def test_resource_types_and_counts(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	# Use one DynamoDbStack stack for all 3 environments since it is
	# not the test subject
	dynamodb_stack = DynamoDbStack(
		app,
		'DynamoDbStackForTests',
		target_environment=DEV
	)

	# Use one GlueStack stack for all 3 environments since it is not
	# the test subject
	glue_stack = GlueStack(
		app,
		'GlueStackForTests',
		target_environment=DEV,
		hash_values_table=dynamodb_stack.hash_values_table,
		value_lookup_table=dynamodb_stack.value_lookup_table,
		multi_lookup_table=dynamodb_stack.value_lookup_table,
		dq_results_table=dynamodb_stack.dq_results_table,
	)

	athena_stacks = {}
	for environment in [DEV, TEST, PROD]:
		athena_stacks[environment] = AthenaHelperStack(
			app,
			f'{environment}-AthenaHelperStackForTests',
			target_environment=environment,
			glue_scripts_temp_bucket=glue_stack.glue_scripts_temp_bucket
		)

	# All stacks should be generated before calling Template methods
	for environment in athena_stacks.keys():
		template = Template.from_stack(athena_stacks[environment])

		template.resource_count_is('AWS::Athena::WorkGroup', 1)
test/test_athena_helper_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk
from aws_cdk.assertions import Template, Match

from test.boto_mocking_helper import *
from lib.code_commit_stack import CodeCommitStack

import lib.configuration as configuration
from lib.configuration import (
    DEPLOYMENT, get_logical_id_prefix
)


def test_resource_types_and_counts(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	codecommit_stack = CodeCommitStack(
		app,
		f'Deploy-CodeCommitStackForTests',
		target_environment='Deploy',
	)

	template = Template.from_stack(codecommit_stack)
	template.resource_count_is('AWS::CodeCommit::Repository', 1)
	template.resource_count_is('AWS::IAM::User', 1)


def test_stack_has_correct_outputs(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	codecommit_stack = CodeCommitStack(
		app,
		f'Deploy-CodeCommitStackForTests',
		target_environment='Deploy',
	)

	template = Template.from_stack(codecommit_stack)
	stack_outputs = template.find_outputs('*')

	repository_output = False
	mirror_user_output = False
	for output_id in stack_outputs.keys():
		output_name = stack_outputs[output_id]['Export']['Name']

		if output_name.find('EtlMirrorRepository') != -1:
			repository_output = True
		if output_name.find('EtlMirrorUser') != -1:
			mirror_user_output = True

	assert repository_output, 'Missing CF output for CodeCommit mirror repository'
	assert mirror_user_output, 'Missing CF output for mirror repository user'


def test_mirror_user_can_access_repository(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	codecommit_stack = CodeCommitStack(
		app,
		'Deploy-CodeCommitStackForTests',
		target_environment='Deploy',
	)

	template = Template.from_stack(codecommit_stack)
	template.has_resource_properties(
		'AWS::IAM::Policy',
		Match.object_like(
			{
				"PolicyDocument": {
					"Statement": [
						{
							"Action": [
								'codecommit:GitPull',
								'codecommit:GitPush'
							],
							"Effect": "Allow",
							"Resource": {
								"Fn::GetAtt": [
									Match.string_like_regexp(r'EtlMirrorRepository\S+'),
									"Arn"
								],
							}
						}
					],
				},
				"PolicyName": Match.any_value(),
				'Users': [
					{ 
						"Ref": Match.string_like_regexp(r'EtlGitExternalMirrorUser\S+')
					}
				]
			}
		)
	)
test/test_codecommit_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest

from test.boto_mocking_helper import *
import lib.configuration as configuration
from lib.configuration import (
    AVAILABILITY_ZONE_1, AVAILABILITY_ZONE_2, AVAILABILITY_ZONE_3,
	ROUTE_TABLE_1, ROUTE_TABLE_2, ROUTE_TABLE_3,
    SHARED_SECURITY_GROUP_ID, SUBNET_ID_1, SUBNET_ID_2, SUBNET_ID_3, VPC_ID,
	S3_KMS_KEY, S3_PURPOSE_BUILT_BUCKET,
	ACCOUNT_ID, REGION,
	ENVIRONMENT, DEPLOYMENT, DEV, PROD, TEST, RESOURCE_NAME_PREFIX
)


def test_get_local_configuration_returns_four_valid_environments(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	for environment in [DEPLOYMENT, DEV, TEST, PROD]:
		local_config = configuration.get_local_configuration(environment)

		assert ACCOUNT_ID in local_config, f'Missing AWS Account from {environment} environment'
		assert REGION in local_config, f'Missing AWS Region from {environment} environment'


def test_get_local_configuration_catches_invalid_prefix(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	with pytest.raises(AttributeError) as e_info:
		configuration.get_local_configuration(DEPLOYMENT, local_mapping={
			DEPLOYMENT: {
				ACCOUNT_ID: mock_account_id,
				REGION: mock_region,
				RESOURCE_NAME_PREFIX: 'Bad_Prefix'
			}
		})

	assert e_info.match('names may only contain'), \
		'Expected Attribute Error for invalid characters in resource name prefix not raised'


def test_get_local_configuration_catches_long_prefix(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	with pytest.raises(AttributeError) as e_info:
		configuration.get_local_configuration(DEPLOYMENT, local_mapping={
			DEPLOYMENT: {
				ACCOUNT_ID: '12digitslong',
				REGION: mock_region,
				RESOURCE_NAME_PREFIX: 'really-long-resource-name-that-will-break-s3-buckets'
			}
		})

	assert e_info.match('prefix is too long'), \
		'Expected Attribute Error for long resource name prefix not raised'


def test_get_local_configuration_catches_bad_environment(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	with pytest.raises(AttributeError) as e_info:
		configuration.get_local_configuration('BadEnvironment', local_mapping={
			DEPLOYMENT: {
				ACCOUNT_ID: mock_account_id,
				REGION: mock_region,
				RESOURCE_NAME_PREFIX: 'testlake'
			}
		})

	assert e_info.match('does not exist in local mappings'), \
		'Expected Attribute Error for invalid environment not raised'


def test_get_environment_configuration_has_outputs_and_environment(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	test_config = configuration.get_environment_configuration(PROD)
	for parameter in [
			ACCOUNT_ID, REGION,
			ENVIRONMENT, VPC_ID, SHARED_SECURITY_GROUP_ID, 
			S3_KMS_KEY, S3_PURPOSE_BUILT_BUCKET
		]:
		assert parameter in test_config, f'Missing {parameter} from PROD environment configuration'


def test_get_all_configurations_has_all_environments(monkeypatch):
	def mock_get_environment_configuration(environment: str):
		return { ENVIRONMENT: environment }

	monkeypatch.setattr(configuration, 'get_environment_configuration', mock_get_environment_configuration)
	# The same mock can work for both functions in this test
	monkeypatch.setattr(configuration, 'get_local_configuration', mock_get_environment_configuration)

	all_config = configuration.get_all_configurations()
	for environment in [DEPLOYMENT, DEV, TEST, PROD]:
		assert environment in all_config


def test_get_logical_id_prefix_returns_string(monkeypatch):
	# Patch boto3, not get_local_configuration() so we can test the structure of local_mapping
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	test_logic_id_prefix = configuration.get_logical_id_prefix()
	assert type(test_logic_id_prefix) == str
	assert len(test_logic_id_prefix) > 0


def test_get_resource_name_prefix_returns_string(monkeypatch):
	# Patch boto3, not get_local_configuration() so we can test the structure of local_mapping
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	test_resource_name_prefix = configuration.get_resource_name_prefix()
	assert type(test_resource_name_prefix) == str
	assert len(test_resource_name_prefix) > 0
test/test_configuration.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys

try:
    from test.glue_job_mocking_helper import *
    from lib.glue_scripts.lib.custom_mapping import *
except ModuleNotFoundError as e:
    if 'pyspark' not in e.msg:
        raise e
    from test.glue_job_mocking_helper_stub import *

pytestmark = pytest.mark.skipif('pyspark' not in sys.modules, reason='No pySpark environment found')


mock_args = {
    'source_key': mock_database_name + '/' + mock_table_name,
}

mock_table_data = [ ( 1, '1/1/2022', 1000, 'test' ), ( 2, '12/31/2022', 2000, 'test' ) ]
mock_table_schema = 'test_column_1 int, test_column_2 string, test_column_3 int, test_column_4 string'

mock_mapping = [
    {
        'sourcename': 'test_column_1',
        'destname': 'renamed_column_1'
    },
    {
        'sourcename': 'test_column_2',
        'destname': 'null'
    },
    {
        'sourcename': 'test_column_3',
        'destname': 'renamed_column_3'
    }
]


def test_custommapping_renames_field():
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    lineage = mock_lineage([])
    df = custommapping(df, mock_mapping, mock_args, lineage)
    mapped_fieldnames = [ field.name for field in df.schema ]
    assert 'renamed_column_1' in mapped_fieldnames
    assert 'renamed_column_3' in mapped_fieldnames


def test_custommapping_drops_fields():
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    lineage = mock_lineage([])
    df = custommapping(df, mock_mapping, mock_args, lineage)
    mapped_fieldnames = [ field.name for field in df.schema ]
    assert 'test_column_2' not in mapped_fieldnames, 'Explicit field drop failed'
    assert 'test_column_4' not in mapped_fieldnames, 'Implicit field drop failed'


def test_custommapping_fuzzymatching_tokenratio():
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    lineage = mock_lineage([])
    df = custommapping(
        df, 
        [
            {
                'sourcename': '1_column_test',
                'destname': 'renamed_column_1',
                'threshold': 90,
                'scorer': 'token_sort_ratio'
            },
            {
                'sourcename': 'test_column_2',
                'destname': 'renamed_column_2',
            },
            {
                'sourcename': 'test_column_3',
                'destname': 'renamed_column_3',
            },
            {
                'sourcename': 'test_column_4',
                'destname': 'renamed_column_4',
            }
        ],
        mock_args, lineage)
    mapped_fieldnames = [ field.name for field in df.schema ]
    assert 'renamed_column_1' in mapped_fieldnames


def test_custommapping_ignores_extra_mapping_data(monkeypatch):
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    lineage = mock_lineage([])
    df = custommapping(
        df, 
        [
            {
                'sourcename': '1_column_test',
                'destname': 'renamed_column_1',
                'threshold': 90,
                'scorer': 'token_set_ratio',
                'comment': 'nothinghere',
                'notes': 'something',
                'owner': 'tom',
            },
            {
                'sourcename': 'test_column_2',
                'destname': 'renamed_column_2',
                'comment': 'nothinghere',
            },
            {
                'sourcename': 'test_column_3',
                'destname': 'renamed_column_3',
            },
            {
                'sourcename': 'test_column_4',
                'destname': 'renamed_column_4',
                'notes': 'somenotes',
            }
        ],
        mock_args, lineage)
    mapped_fieldnames = [ field.name for field in df.schema ]
    for new_name in [ 'renamed_column_1', 'renamed_column_2', 'renamed_column_3', 'renamed_column_4' ]:
        assert new_name in mapped_fieldnames


def test_custommapping_logs_dropped_fields(capsys):
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    lineage = mock_lineage([])
    df = custommapping(df, mock_mapping, mock_args, lineage)
    captured = capsys.readouterr()
    assert 'discarded unmapped fields' in captured.out.lower()
test/test_custom_mapping.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys
import boto3
from moto import mock_dynamodb

try:
    from test.glue_job_mocking_helper import *
    from lib.glue_scripts.lib.datalineage import DataLineageGenerator
except ModuleNotFoundError as e:
    if 'pyspark' not in e.msg:
        raise e
    from test.glue_job_mocking_helper_stub import *

pytestmark = pytest.mark.skipif('pyspark' not in sys.modules, reason='No pySpark environment found')

test_table_name = 'test-data-lineage'

mock_args = {
    'JOB_NAME': 'UnitTest',
    'JOB_RUN_ID': 'jr_unittest',
    'state_machine_name': f'{mock_resource_prefix}-etl-state-machine',
    'execution_id': 'test_execution_id',
    'data_lineage_table': test_table_name,
    'source_key': mock_database_name + '/' + mock_table_name,
}

mock_args_no_table = {
    'JOB_NAME': 'UnitTest',
    'JOB_RUN_ID': 'jr_unittest',
    'state_machine_name': f'{mock_resource_prefix}-etl-state-machine',
    'execution_id': 'test_execution_id',
    'source_key': mock_database_name + '/' + mock_table_name,
}

mock_table_columns = [ 'id', 'date' ]
mock_table_data = [ ( 1, '1/1/2022' ), ( 2, '12/31/2022' ) ]


@pytest.fixture
def use_moto():
    @mock_dynamodb
    def dynamodb_client_and_lineage_table():
        dynamodb = boto3.resource('dynamodb')
 
        # KeySchema, AttributeDefinitions, and BillingMode should match
        # dynamodb table creation in dynamodb_stack
        table = dynamodb.create_table(
            TableName=test_table_name,
            KeySchema=[
                { 'AttributeName': 'step_function_execution_id', 'KeyType': 'HASH' },
                { 'AttributeName': 'job_id_operation_seq', 'KeyType': 'RANGE' }
            ],
            AttributeDefinitions=[
                { 'AttributeName': 'step_function_execution_id', 'AttributeType': 'S' },
                { 'AttributeName': 'job_id_operation_seq', 'AttributeType': 'S' }
            ],
            BillingMode='PAY_PER_REQUEST'
        )
        return table
    return dynamodb_client_and_lineage_table


def test_lineage_is_skipped_when_no_table():
    df = spark.createDataFrame(mock_table_data).toDF(*mock_table_columns)
    lineage = DataLineageGenerator(mock_args_no_table)
    lineage.update_lineage(df, 'TestDb/TestTable', 'read')


@mock_dynamodb
def test_update_lineage_writes_to_table(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    lineage_table = use_moto()

    df = spark.createDataFrame(mock_table_data).toDF(*mock_table_columns)
    lineage = DataLineageGenerator(mock_args)
    lineage.update_lineage(df, 'TestDb/TestTable', 'read')

    item = lineage_table.get_item(
        Key={
            'step_function_execution_id': 'test_execution_id',
            'job_id_operation_seq': 'jr_unittest-1'
        }
    )
    assert item['Item']['dataset'] == 'TestDb/TestTable'


@mock_dynamodb
def test_update_lineage_handles_any_transform(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    lineage_table = use_moto()

    df = spark.createDataFrame(mock_table_data).toDF(*mock_table_columns)
    lineage = DataLineageGenerator(mock_args)
    lineage.update_lineage(df, 'TestDb/TestTable', 'anytransform',
        transform=[
            { 'field': 'id' },
            { 'field': 'date' }
        ]
    )

    item = lineage_table.get_item(
        Key={
            'step_function_execution_id': 'test_execution_id',
            'job_id_operation_seq': 'jr_unittest-1'
        }
    )
    assert item['Item']['dataset'] == 'TestDb/TestTable'
    assert item['Item']['dataset_operation'] == 'anytransform'
test/test_datalineage.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys
import json
from moto import mock_dynamodb

try:
    from test.glue_job_mocking_helper import *
    from lib.glue_scripts.lib.datatransform_lookup import *
except ModuleNotFoundError as e:
    if 'pyspark' not in e.msg:
        raise e
    from test.glue_job_mocking_helper_stub import *

pytestmark = pytest.mark.skipif('pyspark' not in sys.modules, reason='No pySpark environment found')


mock_value_lookup_table = 'unittest-insurancelake-value-lookup-table'
mock_multi_lookup_table = 'unittest-insurancelake-multi-lookup-table'

mock_args = {
    'target_database_name': mock_database_name,
    'source_key': mock_database_name + '/' + mock_table_name,
    'value_lookup_table': mock_value_lookup_table,
    'multi_lookup_table': mock_multi_lookup_table,
}

mock_schema = [
    { 'Name': 'test_column_1', 'Type': 'date' },
    { 'Name': 'test_column_2', 'Type': 'string' },
    { 'Name': 'test_column_3', 'Type': 'int' },
]
mock_field = { 'Name': 'test_column', 'Type': 'string' }
mock_table_columns = [ 'id', 'date' ]
mock_table_data = [ ( 1, '1/1/2022' ), ( 2, '12/31/2022' ) ]
mock_table_schema = 'id int, date string'


@pytest.fixture
def dynamodb_table_for_lookup():
    @mock_dynamodb
    def inner():
        dynamodb = boto3.resource('dynamodb')

        # KeySchema, AttributeDefinitions, and BillingMode should match
        # dynamodb table creation in dynamodb_stack
        table = dynamodb.create_table(
            TableName=mock_value_lookup_table,
            KeySchema=[
                { 'AttributeName': 'source_system', 'KeyType': 'HASH' },
                { 'AttributeName': 'column_name', 'KeyType': 'RANGE' }
            ],
            AttributeDefinitions=[
                { 'AttributeName': 'source_system', 'AttributeType': 'S' },
                { 'AttributeName': 'column_name', 'AttributeType': 'S' }
            ],
            BillingMode='PAY_PER_REQUEST'
        )
        table.put_item(Item={
            'source_system': mock_database_name,
            'column_name': 'unittest',
            'lookup_data': json.dumps({ "1": "First", "2": "Second" })
        })
        return table
    return inner

@mock_dynamodb
def test_transform_lookup_finds_match(monkeypatch, dynamodb_table_for_lookup):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    lineage = mock_lineage([])
    dynamodb_table_for_lookup()
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    df = transform_lookup(df, [ { 'field': 'order', 'source': 'id', 'lookup': 'unittest' } ], mock_args, lineage, spark.sparkContext)
    assert 'order' in df.columns
    assert df.filter('`order` is null').count() == 0


@pytest.fixture
def dynamodb_table_for_multilookup():
    @mock_dynamodb
    def inner():
        dynamodb = boto3.resource('dynamodb')
 
        # KeySchema, AttributeDefinitions, and BillingMode should match
        # dynamodb table creation in dynamodb_stack
        table = dynamodb.create_table(
            TableName=mock_multi_lookup_table,
            KeySchema=[
                { 'AttributeName': 'lookup_group', 'KeyType': 'HASH' },
                { 'AttributeName': 'lookup_item', 'KeyType': 'RANGE' }
            ],
            AttributeDefinitions=[
                { 'AttributeName': 'lookup_group', 'AttributeType': 'S' },
                { 'AttributeName': 'lookup_item', 'AttributeType': 'S' }
            ],
            BillingMode='PAY_PER_REQUEST'
        )
        table.put_item(Item={
            'lookup_group': 'unittest-lookup',
            'lookup_item': '1-1/1/2022',
            'a': '1',
            'b': '2',
        })
        table.put_item(Item={
            'lookup_group': 'unittest-lookup',
            'lookup_item': '2-12/31/2022',
            'a': '3',
            'b': '4',
        })
        return table
    return inner

@mock_dynamodb
def test_transform_multilookup_finds_matches(monkeypatch, dynamodb_table_for_multilookup):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    lineage = mock_lineage([])
    dynamodb_table_for_multilookup()
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    df = transform_multilookup(
        df,
        [{
            'lookup_group': 'unittest-lookup',
            'match_columns': [ 'id', 'date' ],
            'return_attributes': [ 'a', 'b' ],
        }],
        mock_args, lineage, spark.sparkContext)
    assert 'a' in df.columns and 'b' in df.columns
    assert df.filter('`a` is null').count() == 0
    assert df.filter('`b` is null').count() == 0

@mock_dynamodb
def test_transform_multilookup_nomatch_default(monkeypatch, dynamodb_table_for_multilookup):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    lineage = mock_lineage([])
    dynamodb_table_for_multilookup()
    # Test dataframe will have no matches with mock lookup items
    df = spark.createDataFrame([ ( 1, '1/1/2021' ), ( 2, '12/31/2021' ) ], schema=mock_table_schema)
    df = transform_multilookup(
        df,
        [{
            'lookup_group': 'unittest-lookup',
            'match_columns': [ 'id', 'date' ],
            'return_attributes': [ 'a', 'b' ],
            'nomatch': 'N/A'
        }],
        mock_args, lineage, spark.sparkContext)
    assert 'a' in df.columns and 'b' in df.columns
    assert df.filter("`a` = 'N/A'").count() == 2
    assert df.filter("`b` = 'N/A'").count() == 2

@mock_dynamodb
def test_transform_multilookup_error_on_missing_lookup_data(monkeypatch):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    lineage = mock_lineage([])
    # Create empty lookup table
    dynamodb = boto3.resource('dynamodb')
    dynamodb.create_table(
        TableName=mock_multi_lookup_table,
        KeySchema=[
            { 'AttributeName': 'lookup_group', 'KeyType': 'HASH' },
            { 'AttributeName': 'lookup_item', 'KeyType': 'RANGE' }
        ],
        AttributeDefinitions=[
            { 'AttributeName': 'lookup_group', 'AttributeType': 'S' },
            { 'AttributeName': 'lookup_item', 'AttributeType': 'S' }
        ],
        BillingMode='PAY_PER_REQUEST'
    )
    # Test dataframe will have no matches with mock lookup items
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    with pytest.raises(RuntimeError) as e_info:
        df = transform_multilookup(
            df,
            [{
                'lookup_group': 'unittest-lookup',
                'match_columns': [ 'id', 'date' ],
                'return_attributes': [ 'a', 'b' ],
                'nomatch': 'N/A'
            }],
            mock_args, lineage, spark.sparkContext)
    e_info.match('lookup_group')

@mock_dynamodb
def test_get_multilookup_data_paginates(monkeypatch, dynamodb_table_for_multilookup):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    dynamodb_table_for_multilookup()
    items = get_multilookup_data(
        mock_multi_lookup_table,
        'unittest-lookup',
        [ 'a', 'b' ],
        limit=1,
    )
    assert len(items) == 2
test/test_datatransform_lookup.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys
from decimal import Decimal

try:
    from pyspark.sql.types import DateType, StringType
    from test.glue_job_mocking_helper import *
    from lib.glue_scripts.lib.datatransform_typeconversion import *
except ModuleNotFoundError as e:
    if 'pyspark' not in e.msg:
        raise e
    from test.glue_job_mocking_helper_stub import *

pytestmark = pytest.mark.skipif('pyspark' not in sys.modules, reason='No pySpark environment found')

mock_args = {
    'source_key': mock_database_name + '/' + mock_table_name
}

mock_schema = [
    { 'Name': 'test_column_1', 'Type': 'date' },
    { 'Name': 'test_column_2', 'Type': 'string' },
    { 'Name': 'test_column_3', 'Type': 'int' },
]
mock_field = { 'Name': 'test_column', 'Type': 'string' }
mock_table_columns = [ 'id', 'date' ]
mock_table_data = [ ( 1, '1/1/2022' ), ( 2, '12/31/2022' ) ]
mock_table_schema = 'id int, date string'


def test_transform_date_converts_schema():
    lineage = mock_lineage([])
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    assert df.schema['date'].dataType == StringType()
    df = transform_date(df, [ { 'field': 'date', 'format': 'M/d/yyyy' } ], mock_args, lineage)
    assert df.schema['date'].dataType == DateType()

def test_transform_date_converts_data():
    lineage = mock_lineage([])
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    df = transform_date(df, [ { 'field': 'date', 'format': 'M/d/yyyy' } ], mock_args, lineage)
    assert df.filter('`date` is null').count() == 0

def test_transform_date_throws_error_on_bad_convert():
    lineage = mock_lineage([])
    df = spark.createDataFrame(mock_table_data, schema=mock_table_schema)
    # This date format will fail to convert 1/1/2022 because of the single digit month and day
    df = transform_date(df, [ { 'field': 'date', 'format': 'MM/dd/yyyy' } ], mock_args, lineage)
    with pytest.raises(Exception) as e_info:
        df.show(5)
    # Look for mention of this setting in Spark Exception to modify date parsing behavior
    assert e_info.match('spark.sql.legacy.timeParserPolicy')


def test_transform_currency_converts_schema():
    mock_currency_schema = 'money string'
    mock_currency_data = [ ( '100.00', ), ( '$500.00', ), ( '$1,234,567.89', ), ( '-$2,000', ) ]
    lineage = mock_lineage([])
    df = spark.createDataFrame(mock_currency_data, schema=mock_currency_schema)
    assert df.schema['money'].dataType == StringType()
    df = transform_currency(df, [ { 'field': 'money', 'format': '10,2' } ], mock_args, lineage)
    assert str(df.schema['money'].dataType) == 'DecimalType(10,2)'

def test_transform_currency_converts_data():
    mock_currency_schema = 'money string'
    mock_currency_data = [ ( '100.00', ), ( '$500.00', ), ( '$1,234,567.89', ), ( '-$2,000', ) ]
    lineage = mock_lineage([])
    df = spark.createDataFrame(mock_currency_data, schema=mock_currency_schema)
    df = transform_currency(df, [ { 'field': 'money' } ], mock_args, lineage)
    assert df.filter('`money` is null').count() == 0
    assert df.agg({ 'money': 'sum' }).first()['sum(money)'] == Decimal('1233167.89')

def test_transform_currency_converts_euro_data():
    mock_currency_schema = 'money string'
    mock_currency_data = [ ( '100,00', ), ( '$500,00', ), ( '$1.234.567,89', ), ( '-$2.000', ) ]
    lineage = mock_lineage([])
    df = spark.createDataFrame(mock_currency_data, schema=mock_currency_schema)
    df = transform_currency(df, [ { 'field': 'money', 'euro': True } ], mock_args, lineage)
    assert df.filter('`money` is null').count() == 0
    assert df.agg({ 'money': 'sum' }).first()['sum(money)'] == Decimal('1233167.89')

def test_transform_titlecase_converts_string():
    lineage = mock_lineage([])
    df = spark.createDataFrame([ ( 'test string', ) ], schema='text string')
    df = transform_titlecase(df, [ 'text' ], mock_args, lineage)
    assert df.filter('`text` = "Test String"').count() == 1

def test_transform_bigint_converts_string():
    lineage = mock_lineage([])
    df = spark.createDataFrame([ ( '100', ) ], schema='amount string')
    assert df.schema['amount'].dataType == StringType()
    df = transform_bigint(df, [ 'amount' ], mock_args, lineage)
    assert str(df.schema['amount'].dataType) == 'LongType()'
    assert df.filter('`amount` = 100').count() == 1
test/test_datatransform_typeconversion.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk
from aws_cdk.assertions import Template, Match

from test.boto_mocking_helper import *
from lib.dynamodb_stack import DynamoDbStack

import lib.configuration as configuration
from lib.configuration import (
    DEV, PROD, TEST, ACCOUNT_ID, REGION, LOGICAL_ID_PREFIX, RESOURCE_NAME_PREFIX, LINEAGE,
)


mock_environment = {
	ACCOUNT_ID: mock_account_id,
	REGION: mock_region,
	# Mix Deploy environment variables so we can return one dict for all environments
	LOGICAL_ID_PREFIX: 'TestLake',
	RESOURCE_NAME_PREFIX: 'testlake',
}

def mock_get_local_configuration_with_lineage(environment, local_mapping = None):
	lineage_environment = mock_environment.copy()
	lineage_environment.update({LINEAGE: True})
	return lineage_environment

def mock_get_local_configuration_without_lineage(environment, local_mapping = None):
	nolineage_environment = mock_environment.copy()
	nolineage_environment.update({LINEAGE: False})
	return nolineage_environment


def test_resource_types_and_counts_with_lineage(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)
	monkeypatch.setattr(configuration, 'get_local_configuration', mock_get_local_configuration_with_lineage)

	app = cdk.App()
	dynamodb_stack = DynamoDbStack(
		app,
		'Dev-DynamoDbStackForTests',
		target_environment=DEV,
	)
	template = Template.from_stack(dynamodb_stack)

	# Job audit table, lookup value data, multi lookup value, hash value table, dq results table, lineage table
	template.resource_count_is('AWS::DynamoDB::Table', 6)


def test_resource_types_and_counts_without_lineage(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)
	monkeypatch.setattr(configuration, 'get_local_configuration', mock_get_local_configuration_without_lineage)

	app = cdk.App()
	dynamodb_stack = DynamoDbStack(
		app,
		'Dev-DynamoDbStackForTests',
		target_environment=DEV,
	)
	template = Template.from_stack(dynamodb_stack)

	# Job audit table, lookup value data, multi lookup value, hash value table, dq results table
	template.resource_count_is('AWS::DynamoDB::Table', 5)


def test_resource_types_and_counts_all_environments(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()
	dynamodb_stacks = {}
	for environment in [DEV, TEST, PROD]:
		dynamodb_stacks[environment] = DynamoDbStack(
			app,
			f'{environment}-DynamoDbStackForTests',
			target_environment=environment,
		)

	# All stacks should be generated before calling Template methods
	for environment in dynamodb_stacks.keys():
		template = Template.from_stack(dynamodb_stacks[environment])

		# Job audit table, lookup value data, hash value table, dq results table regardless of configuration
		template.has_resource_properties(
			'AWS::DynamoDB::Table',
			Match.object_like(
				{
					"TableName": Match. string_like_regexp('job-audit')
				}
			)
		)
		template.has_resource_properties(
			'AWS::DynamoDB::Table',
			Match.object_like(
				{
					"TableName": Match. string_like_regexp('value-lookup')
				}
			)
		)
		template.has_resource_properties(
			'AWS::DynamoDB::Table',
			Match.object_like(
				{
					"TableName": Match. string_like_regexp('multi-lookup')
				}
			)
		)
		template.has_resource_properties(
			'AWS::DynamoDB::Table',
			Match.object_like(
				{
					"TableName": Match. string_like_regexp('hash-values')
				}
			)
		)
		template.has_resource_properties(
			'AWS::DynamoDB::Table',
			Match.object_like(
				{
					"TableName": Match. string_like_regexp('dq-results')
				}
			)
		)
test/test_dynamodb_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys
from urllib.parse import urlparse
import shutil
from moto import mock_athena
from test.boto_mocking_helper import *

try:
    from test.glue_job_mocking_helper import *
    from awsglue.utils import GlueArgumentError
    import lib.glue_scripts.etl_cleanse_to_consume as etl_cleanse_to_consume
except ModuleNotFoundError as e:
    if 'pyspark' not in e.msg:
        raise e
    from test.glue_job_mocking_helper_stub import *

pytestmark = pytest.mark.skipif('pyspark' not in sys.modules, reason='No pySpark environment found')

mock_args = [
    'etl_cleanse-to-consume.py',
    '--JOB_NAME=UnitTest',
    '--JOB_RUN_ID=jr-unittest',
    f'--state_machine_name={mock_resource_prefix}-etl-state-machine',
    '--execution_id=test_execution_id',
    '--environment=unittest',
    f'--TempDir={mock_temp_bucket}',
    f'--txn_bucket={mock_scripts_bucket}',
    '--txn_sql_prefix_path=/etl/transformation-sql/',
    f'--source_bucket={mock_cleanse_bucket}',
    f'--target_bucket={mock_consume_bucket}',
    f'--database_name_prefix={mock_database_name}',
    f'--table_name={mock_table_name}',
    '--base_file_name=testfile.csv',
    '--p_year=2022',
    '--p_month=12',
    '--p_day=6',
]

mock_spark_sql_file = """select
    1 as year,
    2 as month,
    3 as day,
    1234 as data
    """

mock_create_table_spark_sql_file = """create table differenttable as
    select
    1 as year,
    2 as month,
    3 as day,
    1234 as data
    """

mock_athena_sql_file = """create view test_view as
    select
    1 as c1,
    2 as c2
    """


@mock_glue_job(etl_cleanse_to_consume)
def test_job_execution_and_commit(monkeypatch):
    monkeypatch.setattr(sys, 'argv', mock_args)
    etl_cleanse_to_consume.main()


@mock_glue_job(etl_cleanse_to_consume)
def test_consume_bucket_write(monkeypatch):
    monkeypatch.setattr(sys, 'argv', mock_args)
    write_local_file(f'{mock_scripts_bucket}/etl/transformation-sql', f'spark-{mock_database_name}-{mock_table_name}.sql', mock_spark_sql_file)

    consume_path = f'{mock_consume_bucket}/{mock_database_name}/{mock_table_name}'
    parsed_uri = urlparse(consume_path)
    shutil.rmtree(parsed_uri.path, ignore_errors=True)

    etl_cleanse_to_consume.main()
    assert os.path.isdir(parsed_uri.path), \
        'Cleanse-to-Consume Glue job failed to write to Consume bucket with supplied Spark SQL'


@mock_glue_job(etl_cleanse_to_consume)
def test_missing_argument_exception(monkeypatch):
    """
    Check error is raised if base_file_name argument is missing
    """
    mock_args_without_base_file_name = mock_args.copy()
    mock_args_without_base_file_name.remove('--base_file_name=testfile.csv')
    monkeypatch.setattr(sys, 'argv', mock_args_without_base_file_name)
    with pytest.raises(GlueArgumentError):
        etl_cleanse_to_consume.main()


@mock_athena()
def test_athena_execute_query_success(monkeypatch):
    """
    Test simple create view and expect success
    """
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    result = etl_cleanse_to_consume.athena_execute_query(mock_database_name, mock_athena_sql_file, mock_temp_bucket)
    assert result == 'SUCCEEDED'


@mock_athena
def test_athena_execute_query_max_retries_error(monkeypatch):
    """
    Provide 0 attempts and expect failure from exceeding max attempts
    """
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setattr(sys, 'argv', mock_args)
    with pytest.raises(RuntimeError) as e_info:
        etl_cleanse_to_consume.athena_execute_query(mock_database_name, mock_athena_sql_file, mock_temp_bucket, max_attempts = 0)
    assert e_info.match('exceeded max_attempts')


def test_athena_execute_query_fail(monkeypatch):
    """
    Test simple create table and expect failure
    Mock API directly because there is no way to generate a failure with moto's athena mock
    """
    monkeypatch.setattr(sys, 'argv', mock_args)
    monkeypatch.setattr(etl_cleanse_to_consume.boto3, 'client', mock_boto3_client)
    with pytest.raises(RuntimeError) as e_info:
        etl_cleanse_to_consume.athena_execute_query(mock_database_name, mock_athena_sql_file, mock_temp_bucket)
    assert e_info.match('failed with query engine error')
test/test_etl_cleanse_to_consume.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys
import json
from urllib.parse import urlparse
import shutil

try:
    from test.glue_job_mocking_helper import *
    from awsglue.utils import GlueArgumentError
    import lib.glue_scripts.etl_collect_to_cleanse as etl_collect_to_cleanse
except ModuleNotFoundError as e:
    if 'pyspark' not in e.msg:
        raise e
    from test.glue_job_mocking_helper_stub import *

pytestmark = pytest.mark.skipif('pyspark' not in sys.modules, reason='No pySpark environment found')

mock_args = [
    'etl_collect-to-cleanse.py',
    '--JOB_NAME=UnitTest',
    '--JOB_RUN_ID=jr-unittest',
    f'--state_machine_name={mock_resource_prefix}-etl-state-machine',
    '--execution_id=test_execution_id',
    '--environment=unittest',
    f'--TempDir=file:///tmp/{mock_resource_prefix}-temp',
    f'--txn_bucket={mock_scripts_bucket}',
    '--txn_spec_prefix_path=/etl/transformation-spec/',
    f'--source_bucket={mock_collect_bucket}',
    f'--target_bucket=file:///tmp/{mock_resource_prefix}-cleanse-bucket',
    f'--source_key={mock_database_name}/{mock_table_name}',
    f'--target_database_name={mock_database_name}',
    f'--table_name={mock_table_name}',
    '--base_file_name=testfile.csv',
    '--p_year=2022',
    '--p_month=12',
    '--p_day=6',
    f'--hash_value_table={mock_resource_prefix}-hash-value-table',
    f'--value_lookup_table={mock_resource_prefix}-value-lookup-table',
    f'--multi_lookup_table={mock_resource_prefix}-multi-lookup-table',
    f'--dq_results_table={mock_resource_prefix}-dq-results-table',
]

mock_input_file_csv = """field1,field2
1,1/1/2022
2,12/31/2022
"""

mock_input_file_tsv = """field1\tfield2
1\t1/1/2022
2\t12/31/2022
"""

mock_mapping_file = """SourceName,DestName
field1,id
field2,date
"""

mock_spec_file = {}
mock_spec_file_tsv = {
    'input_spec': {
        'tsv': {}
    }
}


@mock_glue_job(etl_collect_to_cleanse)
def test_cleanse_bucket_write(monkeypatch):
    monkeypatch.setattr(sys, 'argv', mock_args)
    write_local_file(f'{mock_collect_bucket}/{mock_database_name}/{mock_table_name}', 'testfile.csv', mock_input_file_csv)

    cleanse_path = f'{mock_cleanse_bucket}/{mock_database_name}/{mock_table_name}'
    parsed_uri = urlparse(cleanse_path)
    shutil.rmtree(parsed_uri.path, ignore_errors=True)

    etl_collect_to_cleanse.main()
    assert os.path.isdir(parsed_uri.path), \
        'Collect-to-Cleanse Glue job failed to write to Cleanse bucket with supplied input file'


@mock_glue_job(etl_collect_to_cleanse)
def test_job_execution_and_commit_csv(monkeypatch):
    monkeypatch.setattr(sys, 'argv', mock_args)
    write_local_file(f'{mock_collect_bucket}/{mock_database_name}/{mock_table_name}', 'testfile.csv', mock_input_file_csv)
    write_local_file(f'{mock_scripts_bucket}/etl/transformation-spec', f'{mock_database_name}-{mock_table_name}.csv', mock_mapping_file)
    write_local_file(f'{mock_scripts_bucket}/etl/transformation-spec', f'{mock_database_name}-{mock_table_name}.json', json.dumps(mock_spec_file))
    etl_collect_to_cleanse.main()


@mock_glue_job(etl_collect_to_cleanse)
def test_job_execution_and_commit_tsv(monkeypatch):
    monkeypatch.setattr(sys, 'argv', mock_args)
    write_local_file(f'{mock_collect_bucket}/{mock_database_name}/{mock_table_name}', 'testfile.csv', mock_input_file_tsv)
    write_local_file(f'{mock_scripts_bucket}/etl/transformation-spec', f'{mock_database_name}-{mock_table_name}.csv', mock_mapping_file)
    write_local_file(f'{mock_scripts_bucket}/etl/transformation-spec', f'{mock_database_name}-{mock_table_name}.json', json.dumps(mock_spec_file_tsv))
    etl_collect_to_cleanse.main()


@mock_glue_job(etl_collect_to_cleanse)
def test_missing_argument_exception(monkeypatch):
    """
    Check error is raised if base_file_name argument is missing
    """
    mock_args_without_base_file_name = mock_args.copy()
    mock_args_without_base_file_name.remove('--base_file_name=testfile.csv')
    monkeypatch.setattr(sys, 'argv', mock_args_without_base_file_name)

    with pytest.raises(GlueArgumentError):
        etl_collect_to_cleanse.main()
test/test_etl_collect_to_cleanse.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
from moto import mock_dynamodb
import boto3
import botocore
import json
from lib.etl_job_auditor.lambda_handler import lambda_handler

mock_region = 'us-east-1'
test_table_name = 'test-etl-job-audit'
test_success_execution_id = '3b278a83-8b04-468a-983e-3749b6609dcb'
test_failure_execution_id = '884423e7-2dba-4597-9907-371d1f424e16'

state_machine_success_event = {
    'eventTime': '2023-01-02T23:00:00.000Z', 
    'Input': {
        'execution_id': test_success_execution_id, 
        'taskresult': {
    		'AllocatedCapacity': 50, 
            'Attempt': 0, 
            'CompletedOn': 1675701338733, 
            'ExecutionTime': 30, 
            'GlueVersion': '3.0', 
            'Id': 'jr_5ae791637444c02e3c26b8c97750446cbead7bf184494f2aaaf5be713e7e55f6', 
            'JobName': 'test_run_success', 
            'JobRunState': 'SUCCEEDED', 
            'LastModifiedOn': 1675701338733, 
            'LogGroupName': '/aws-glue/jobs', 
            'MaxCapacity': 50.0, 
            'NumberOfWorkers': 50, 
            'PredecessorRuns': [], 
            'StartedOn': 1675701301319, 
            'Timeout': 2880, 
            'WorkerType': 'G.1X'
        }
    }
}

state_machine_failure_event = {
    'eventTime': '2023-01-03T23:00:00.000Z', 
    'Input': {
        'execution_id': test_failure_execution_id, 
        'taskresult': {
            'Error': 'States.TaskFailed', 
            'Cause': json.dumps(
                {
                    "AllocatedCapacity": 50,
                    "Attempt": 0,
                    "CompletedOn": 1675434244707,
                    "DPUSeconds": 120.0,
                    "ErrorMessage": "RuntimeError: Error for testing",
                    "ExecutionTime": 24,
                    "GlueVersion": "3.0",
                    "Id": "jr_bde060a78592786933a57ade0a282f39eec026236b82ba3b678e8154592d666c",
                    "JobName": "test_run_failure",
                    "JobRunState": "FAILED",
                    "LastModifiedOn": 1675434244707,
                    "LogGroupName": "/aws-glue/jobs",
                    "MaxCapacity": 50.0,
                    "NumberOfWorkers": 50,
                    "PredecessorRuns": [],
                    "StartedOn": 1675434215855,
                    "Timeout": 2880,
                    "WorkerType":"G.1X"
                })
        }
    }
}

test_context = {
    'aws_request_id': 'd7d36c0e-a1bc-11ed-a8fc-0242ac120002',
    'function_name': 'testfunction',
    'invoked_function_arn': 'arn:aws:lambda:::function:notrealarn',
    'identity': None,
    'client_context': None
}

@pytest.fixture
def use_moto():
    @mock_dynamodb
    def dynamodb_client_and_audit_table():
        dynamodb = boto3.resource('dynamodb')
 
        # KeySchema, AttributeDefinitions, and BillingMode should match
        # dynamodb table creation in dynamodb_stack
        table = dynamodb.create_table(
            TableName=test_table_name,
            KeySchema=[
                { 'AttributeName': 'execution_id', 'KeyType': 'HASH' }
            ],
            AttributeDefinitions=[
                { 'AttributeName': 'execution_id', 'AttributeType': 'S' }
            ],
            BillingMode='PAY_PER_REQUEST'
        )
        return table
    return dynamodb_client_and_audit_table


@mock_dynamodb
def test_handler_success_event_returns_200(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    use_moto()

    result = lambda_handler(state_machine_success_event, test_context)
    assert result['statusCode'] == 200


@mock_dynamodb
def test_handler_success_event_returns_state(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    use_moto()

    result = lambda_handler(state_machine_success_event, test_context)
    assert result['body'].find('SUCCEEDED') != -1


@mock_dynamodb
def test_handler_success_event_records_state(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    table = use_moto()

    lambda_handler(state_machine_success_event, test_context)

    item = table.get_item(
        TableName=test_table_name,
        Key={ 'execution_id': test_success_execution_id }
    )
    assert item['Item']['job_latest_status'] == 'SUCCEEDED'


@mock_dynamodb
def test_handler_failure_event_returns_200(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    use_moto()

    result = lambda_handler(state_machine_failure_event, test_context)
    assert result['statusCode'] == 200


@mock_dynamodb
def test_handler_failure_event_returns_state_and_error(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    use_moto()

    result = lambda_handler(state_machine_failure_event, test_context)
    assert result['body'].find('FAILED') != -1
    assert result['body'].find('Error for testing') != -1


@mock_dynamodb
def test_handler_failure_event_records_error_message(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    table = use_moto()

    lambda_handler(state_machine_failure_event, test_context)

    item = table.get_item(
        TableName=test_table_name,
        Key={ 'execution_id': test_failure_execution_id }
    )
    assert item['Item']['error_message'] == 'RuntimeError: Error for testing'


@mock_dynamodb
def test_wrong_table_returns_error(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', 'wrong_table_name')
    use_moto()

    with pytest.raises(botocore.exceptions.ClientError) as e_info:
        lambda_handler(state_machine_success_event, test_context)
        assert e_info.match('Requested resource not found')
test/test_etl_job_auditor.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import sys
import boto3
from moto import mock_glue

try:
    from test.glue_job_mocking_helper import *
    from lib.glue_scripts.lib.glue_catalog_helpers import *
except ModuleNotFoundError as e:
    if 'pyspark' not in e.msg:
        raise e
    from test.glue_job_mocking_helper_stub import *

pytestmark = pytest.mark.skipif('pyspark' not in sys.modules, reason='No pySpark environment found')

mock_schema = [
    { 'Name': 'test_column_1', 'Type': 'date' },
    { 'Name': 'test_column_2', 'Type': 'string' },
    { 'Name': 'test_column_3', 'Type': 'int' },
]
mock_field = { 'Name': 'test_column', 'Type': 'string' }
mock_table_columns = [ 'id', 'date' ]
mock_table_data = [ ( 1, '1/1/2022' ), ( 2, '12/31/2022' ) ]


@mock_glue
def test_table_exists_true(monkeypatch):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    glue_client = boto3.client('glue')
    glue_client.create_database(DatabaseInput={ 'Name': mock_database_name })
    glue_client.create_table(
        DatabaseName=mock_database_name,
        TableInput={ 'Name': mock_table_name }
    )
    assert table_exists(mock_database_name, mock_table_name) == True

@mock_glue
def test_table_exists_false(monkeypatch):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    glue_client = boto3.client('glue')
    glue_client.create_database(DatabaseInput={ 'Name': mock_database_name })
    # Purposely do not create the table
    assert table_exists(mock_database_name, mock_table_name) == False

@mock_glue
def test_create_database(monkeypatch):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    glue_client = boto3.client('glue')
    create_database(mock_database_name)
    assert glue_client.get_database(Name=mock_database_name)

def test_check_schema_change_permissive():
    new_schema = mock_schema.copy()
    new_schema.append(mock_field)
    assert check_schema_change(mock_schema, new_schema, 'permissive')

def test_check_schema_change_strict_same():
    assert check_schema_change(mock_schema, mock_schema, 'strict')

def test_check_schema_change_strict_different():
    new_schema = mock_schema.copy()
    new_schema.append(mock_field)
    assert not check_schema_change(mock_schema, new_schema, 'strict')

def test_check_schema_change_reorder_same():
    new_schema = mock_schema.copy()
    assert check_schema_change(mock_schema, new_schema, 'reorder')

def test_check_schema_change_reorder_reordered():
    new_schema = list(reversed(mock_schema.copy()))
    assert check_schema_change(mock_schema, new_schema, 'reorder')

def test_check_schema_change_reorder_nonpermissible():
    new_schema = mock_schema.copy()
    new_schema.append(mock_field)
    assert not check_schema_change(mock_schema, new_schema, 'reorder')

@mock_glue
def test_upsert_catalog_table_insert(monkeypatch):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    glue_client = boto3.client('glue')
    # create_database tests will test that function, so skip it by creating db now
    glue_client.create_database(DatabaseInput={ 'Name': mock_database_name })
    df = spark.createDataFrame(mock_table_data).toDF(*mock_table_columns)

    upsert_catalog_table(df, mock_database_name, mock_table_name, [], 's3://fake-bucket', 'strict')

    assert glue_client.get_table(DatabaseName=mock_database_name, Name=mock_table_name)

@mock_glue
def test_upsert_catalog_table_update(monkeypatch):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    glue_client = boto3.client('glue')
    glue_client.create_database(DatabaseInput={ 'Name': mock_database_name })
    glue_client.create_table(
        DatabaseName=mock_database_name,
        TableInput={
            'Name': mock_table_name,
            'StorageDescriptor': { 'Columns': [ mock_field ] }
        }
    )

    df = spark.createDataFrame(mock_table_data).toDF(*mock_table_columns)
    upsert_catalog_table(df, mock_database_name, mock_table_name, [], 's3://fake-bucket', 'permissive')

    df_schema = [ { 'Name': field[0], 'Type': field[1] } for field in df.dtypes ]
    response = glue_client.get_table(DatabaseName=mock_database_name, Name=mock_table_name)

    for field in df_schema:
        assert field in response['Table']['StorageDescriptor']['Columns']

@mock_glue
def test_upsert_catalog_table_update_raises_error(monkeypatch):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    glue_client = boto3.client('glue')
    glue_client.create_database(DatabaseInput={ 'Name': mock_database_name })
    glue_client.create_table(
        DatabaseName=mock_database_name,
        TableInput={
            'Name': mock_table_name,
            'StorageDescriptor': { 'Columns': [ mock_field ] }
        }
    )

    df = spark.createDataFrame(mock_table_data).toDF(*mock_table_columns)

    with pytest.raises(RuntimeError) as e_info:
        upsert_catalog_table(df, mock_database_name, mock_table_name, [], 's3://fake-bucket', 'strict')
    assert e_info.match('Nonpermissible')
test/test_glue_catalog_helpers.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk
from aws_cdk.assertions import Template

from test.boto_mocking_helper import *
from lib.dynamodb_stack import DynamoDbStack
from lib.glue_stack import GlueStack

import lib.configuration as configuration
from lib.configuration import (
    DEV, PROD, TEST, ACCOUNT_ID, REGION, LOGICAL_ID_PREFIX, RESOURCE_NAME_PREFIX, VPC_CIDR, LINEAGE,
)

def mock_get_local_configuration_with_vpc(environment, local_mapping = None):
	return {
		ACCOUNT_ID: mock_account_id,
		REGION: mock_region,
		LINEAGE: False,
		# Mix Deploy environment variables so we can return one dict for all environments
		LOGICAL_ID_PREFIX: 'TestLake',
		RESOURCE_NAME_PREFIX: 'testlake',
		VPC_CIDR: '10.0.0.0/24',
	}


def test_resource_types_and_counts(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	# Use one DynamoDbStack stack for all 3 environments since it is not the test subject
	dynamodb_stack = DynamoDbStack(
		app,
		'DynamoDbStackForTests',
		target_environment=DEV
	)

	glue_stacks = {}
	for environment in [DEV, TEST, PROD]:
		glue_stacks[environment] = GlueStack(
			app,
			f'{environment}-GlueStackForTests',
			target_environment=environment,
			hash_values_table=dynamodb_stack.hash_values_table,
			value_lookup_table=dynamodb_stack.value_lookup_table,
			multi_lookup_table=dynamodb_stack.value_lookup_table,
			dq_results_table=dynamodb_stack.dq_results_table,
		)

	# All stacks should be generated before calling Template methods
	for environment in glue_stacks.keys():
		template = Template.from_stack(glue_stacks[environment])

		# Collect-Cleanse, Cleanse-Consume
		template.resource_count_is('AWS::Glue::Job', 2)
		# Glue Scripts, Glue Temp
		template.resource_count_is('AWS::S3::Bucket', 2)
		# S3 Deployment CustomResource handler, CustomResource Log Retention handler
		template.resource_count_is('AWS::Lambda::Function', 2)
		# Glue Job role, S3 Deployment CustomResource handler role, CustomResource Log Retention handler role
		template.resource_count_is('AWS::IAM::Role', 3)


def test_glue_connections_with_vpc(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)
	monkeypatch.setattr(configuration, 'get_local_configuration', mock_get_local_configuration_with_vpc)

	app = cdk.App()

	# Use one DynamoDbStack stack for all 3 environments since it is not the test subject
	dynamodb_stack = DynamoDbStack(
		app,
		'DynamoDbStackForTests',
		target_environment=DEV
	)

	glue_stack = GlueStack(
		app,
		'Dev-GlueStackForTests',
		target_environment=DEV,
		hash_values_table=dynamodb_stack.hash_values_table,
		value_lookup_table=dynamodb_stack.value_lookup_table,
		multi_lookup_table=dynamodb_stack.value_lookup_table,
		dq_results_table=dynamodb_stack.dq_results_table,
	)
	template = Template.from_stack(glue_stack)

	# 3 Glue Connections, one for each AZ
	template.resource_count_is('AWS::Glue::Connection', 3)
test/test_glue_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk
from aws_cdk.assertions import Template, Match

from test.boto_mocking_helper import *
from lib.pipeline_stack import PipelineStack

import lib.configuration as configuration
from lib.configuration import (
    DEV, PROD, TEST
)


def test_resource_types_and_counts(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	pipeline_stacks = {}
	for environment in [DEV, TEST, PROD]:
		pipeline_stacks[environment] = PipelineStack(
			app,
			f'{environment}-PipelineStackForTests',
			target_environment=environment,
			target_branch='main', 
			# Target and Pipeline account/region are the same - not testing cross-account/cross-region
			target_aws_env={ 'account': mock_account_id, 'region': mock_region },
			env=cdk.Environment(
				account=mock_account_id,
				region=mock_region
			),
		)

	# 3 stacks expected (dev, test, prod), no cross-pipeline support stack
	assert len(app.node.children) == 3, 'Unexpected number of stacks'

	# All stacks should be generated before calling Template methods
	for environment in pipeline_stacks.keys():
		template = Template.from_stack(pipeline_stacks[environment])

		template.resource_count_is('AWS::CodePipeline::Pipeline', 1)
		# Project for cdk synth, and pipeline update/self-mutate, 6 file asset pipeline steps
		template.resource_count_is('AWS::CodeBuild::Project', 8)
		# Artifact bucket
		template.resource_count_is('AWS::S3::Bucket', 1)
		# Artifact bucket encryption key
		template.resource_count_is('AWS::KMS::Key', 1)
		# LogGroup for each build action (includes 6 file asset pipeline steps)
		template.resource_count_is('AWS::Logs::LogGroup', 8)
		# CodePipeline role, 3 CodeBuild roles, 2 Pipeline action roles, Pipeline event role
		template.resource_count_is('AWS::IAM::Role', 7)


def test_cross_region_number_of_stacks(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	pipeline_stacks = {}
	for environment in [DEV, TEST, PROD]:
		pipeline_stacks[environment] = PipelineStack(
			app,
			f'{environment}-PipelineStackForTests',
			target_environment=environment,
			target_branch='main', 
			# Different fake region for each environment to trigger pipeline support stack
			target_aws_env={
				'account': mock_account_id, 
				'region': f'{environment.lower()}-region'
			},
			env=cdk.Environment(
				account=mock_account_id,
				region=mock_region
			),
		)

	# 3 infrastructure stacks (dev, test, prod), 3 pipeline support stacks
	assert len(app.node.children) == 6, 'Unexpected number of stacks'


def test_cross_account_number_of_stacks(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	pipeline_stacks = {}
	for environment in [DEV, TEST, PROD]:
		pipeline_stacks[environment] = PipelineStack(
			app,
			f'{environment}-PipelineStackForTests',
			target_environment=environment,
			target_branch='main', 
			# Different accounts for each environment
			target_aws_env={ 
				'account': f'{environment.lower()}notrealaccount', 
				'region': mock_region 
			},
			env=cdk.Environment(
				account=mock_account_id,
				region=mock_region
			),
		)

	# 3 stacks expected (dev, test, prod)
	assert len(app.node.children) == 3, 'Unexpected number of stacks'


def test_pipeline_self_mutates(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	stack_logical_id = 'Dev-PipelineStackForTests'
	pipeline_stack = PipelineStack(
		app,
		stack_logical_id,
		target_environment=DEV,
		target_branch='main', 
		target_aws_env={ 'account': mock_account_id, 'region': mock_region },
		env=cdk.Environment(
			account=mock_account_id,
			region=mock_region
		),
	)

	template = Template.from_stack(pipeline_stack)
	template.has_resource_properties(
		'AWS::CodeBuild::Project',
		Match.object_like(
			{
				"Source": {
					"BuildSpec": Match.serialized_json(
						{
							"version": Match.any_value(),
							"phases": {
								"install": Match.any_value(),
								"build": {
									"commands": [ 
										Match.string_like_regexp(fr'cdk -a . deploy {stack_logical_id} \S+')
									]
								}
							}
						}
					)
				}
			}
		)
	)


def test_codebuild_runs_synth(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	pipeline_stack = PipelineStack(
		app,
		'Dev-PipelineStackForTests',
		target_environment=DEV,
		target_branch='main', 
		target_aws_env={ 'account': mock_account_id, 'region': mock_region },
		env=cdk.Environment(
			account=mock_account_id,
			region=mock_region
		),
	)

	template = Template.from_stack(pipeline_stack)
	template.has_resource_properties(
		'AWS::CodeBuild::Project',
		Match.object_like(
			{
				"Source": {
					"BuildSpec": Match.serialized_json(
						{
							"version": Match.any_value(),
							"phases": {
								"build": {
									"commands": Match.array_with(['cdk synth'])
								}
							},
							"artifacts": Match.any_value()
						}
					)
				}
			}
		)
	)
test/test_pipeline_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk

from test.boto_mocking_helper import *
from lib.stack_import_helper import ImportedBuckets, ImportedVpc

import lib.configuration as configuration
from lib.configuration import (
    DEPLOYMENT, DEV, PROD, TEST, ACCOUNT_ID, REGION, VPC_CIDR, RESOURCE_NAME_PREFIX,
)


def test_importedbuckets_resource_counts(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	for environment in [DEV, TEST, PROD]:
		app = cdk.App()
		stack = cdk.Stack(app, 'StackForTests')
		stack.mappings = configuration.get_environment_configuration(environment)
		ImportedBuckets(stack, logical_id_suffix='TestSuffix')

		bucket_proxy_count = 0
		key_proxy_count = 0
		for resource in stack.node.children:
			if str(type(resource)) == "<class 'jsii._reference_map.InterfaceDynamicProxy'>":
				for child_resource in resource._delegates:
					# Use string comparison because _IKeyProxy is a private class
					if str(type(child_resource)) == "<class 'aws_cdk.aws_kms._IKeyProxy'>":
						key_proxy_count += 1

			# Use string comparison because _BucketBaseProxy is a private class
			if str(type(resource)) == "<class 'aws_cdk.aws_s3._BucketBaseProxy'>":
				bucket_proxy_count += 1

		assert bucket_proxy_count == 4, \
			f'Unexpected number of imported S3 Bucket proxies in {environment} environment'
		assert key_proxy_count == 1, \
			f'Unexpected number of imported KMS key proxies in {environment} environment'


def test_importedvpcs_resource_counts(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()
	stack = cdk.Stack(app, 'StackForTests')
	stack.mappings = configuration.get_environment_configuration(DEV, local_mapping = {
		DEPLOYMENT: {
			ACCOUNT_ID: mock_account_id,
			REGION: mock_region,
			RESOURCE_NAME_PREFIX: 'testlake',
		},
		DEV: {
			ACCOUNT_ID: mock_account_id,
			REGION: mock_region,
			VPC_CIDR: '10.0.0.0/24'
		}
	})
	ImportedVpc(stack, logical_id_suffix='TestSuffix')

	vpc_proxy_count = 0
	subnet_proxy_count = 0
	securitygroup_proxy_count = 0
	for jsii_resource in stack.node.children:
		# Go another level deeper to get better resource type visibility
		for resource in jsii_resource._delegates:
			# Use string comparison because Proxies are private classes
			if str(type(resource)) == "<class 'aws_cdk.aws_ec2._IVpcProxy'>":
				vpc_proxy_count += 1
			if str(type(resource)) == "<class 'aws_cdk.aws_ec2._ISubnetProxy'>":
				subnet_proxy_count += 1
			if str(type(resource)) == "<class 'aws_cdk.aws_ec2._ISecurityGroupProxy'>":
				securitygroup_proxy_count += 1

	# Imported values are used and exposed as 3 specific resources on the stack
	assert vpc_proxy_count == 1, \
		f'Unexpected number of imported VPC proxies'
	assert subnet_proxy_count == 3, \
		f'Unexpected number of imported Subnet proxies'
	assert securitygroup_proxy_count == 1, \
		f'Unexpected number of imported SecurityGroup proxies'
test/test_stack_import_helper.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
from moto import mock_dynamodb
import boto3
import botocore
import logging

from test.boto_mocking_helper import *
import lib.state_machine_trigger.lambda_handler as lambda_handler

test_table_name = 'test-etl-job-audit'
test_state_machine_name = 'test-state-machine'
test_state_machine_arn = f'arn:aws:states:{mock_region}:{mock_account_id}:stateMachine:{test_state_machine_name}'
test_bucket_arn = 'arn:aws:s3:::test-bucket'
test_execution_id = '3b278a83-8b04-468a-983e-3749b6609dcb'
test_event_time = '2023-01-01T23:30:00.000Z'

test_success_event = {
    'Records': [{
        'eventTime': test_event_time, 
        's3': {
            'bucket': { 'name': 'test-bucket', 'arn': test_bucket_arn }, 
            'object': { 'key': 'level1/level2/test-file.csv' }
        },
        'userIdentity': { 'principalId': 'testuser' },
        'requestParameters': { 'sourceIPAddress': 'testipaddress'},
    }]
}

test_bad_file_path_event1 = {
    'Records': [{
        'eventTime': test_event_time, 
        's3': {
            'bucket': { 'name': 'test-bucket', 'arn': test_bucket_arn }, 
            'object': { 'key': 'top-level-test-file.csv' }
        },
        'userIdentity': { 'principalId': 'testuser' },
        'requestParameters': { 'sourceIPAddress': 'testipaddress'},
    }]
}

test_bad_file_path_event2 = {
    'Records': [{
        'eventTime': test_event_time, 
        's3': {
            'bucket': { 'name': 'test-bucket', 'arn': test_bucket_arn }, 
            'object': { 'key': 'test-source-system/first-level-test-file.csv' }
        },
        'userIdentity': { 'principalId': 'testuser' },
        'requestParameters': { 'sourceIPAddress': 'testipaddress'},
    }]
}

test_bad_folder_event = {
    'Records': [{
        'eventTime': test_event_time, 
        's3': {
            'bucket': { 'name': 'test-bucket', 'arn': test_bucket_arn }, 
            'object': { 'key': 'level1/level2/' }
        },
        'userIdentity': { 'principalId': 'testuser' },
        'requestParameters': { 'sourceIPAddress': 'testipaddress'},
    }]
}

test_context = {
    'aws_request_id': 'd7d36c0e-a1bc-11ed-a8fc-0242ac120002',
    'function_name': 'testfunction',
    'invoked_function_arn': 'arn:aws:lambda:::function:notrealarn',
    'identity': None,
    'client_context': None
}


def mock_record_etl_job_run(audit_table_name: str, sfn_arn: str, execution_id: str,
        execution_name: str, execution_input: str, principal_id: str, source_ipaddress: str):
    # Mock function does not need to do any work
    pass


@pytest.fixture
def use_moto():
    @mock_dynamodb
    def dynamodb_client_and_audit_table():
        dynamodb = boto3.resource('dynamodb')
 
        # KeySchema, AttributeDefinitions, and BillingMode should match
        # dynamodb table creation in dynamodb_stack
        table = dynamodb.create_table(
            TableName=test_table_name,
            KeySchema=[
                { 'AttributeName': 'execution_id', 'KeyType': 'HASH' }
            ],
            AttributeDefinitions=[
                { 'AttributeName': 'execution_id', 'AttributeType': 'S' }
            ],
            BillingMode='PAY_PER_REQUEST'
        )
        return table
    return dynamodb_client_and_audit_table


def test_handler_success_returns_200(monkeypatch):
    monkeypatch.setattr(lambda_handler.boto3, 'client', mock_boto3_client)
    monkeypatch.setattr(lambda_handler, 'record_etl_job_run', mock_record_etl_job_run)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    monkeypatch.setenv('SFN_STATE_MACHINE_ARN', test_state_machine_arn)

    result = lambda_handler.lambda_handler(test_success_event, test_context)
    assert result['statusCode'] == 200


def test_handler_bad_file_returns_400(monkeypatch):
    monkeypatch.setattr(lambda_handler.boto3, 'client', mock_boto3_client)
    monkeypatch.setattr(lambda_handler, 'record_etl_job_run', mock_record_etl_job_run)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    monkeypatch.setenv('SFN_STATE_MACHINE_ARN', test_state_machine_arn)

    result1 = lambda_handler.lambda_handler(test_bad_file_path_event1, test_context)
    assert result1['statusCode'] == 400

    result2 = lambda_handler.lambda_handler(test_bad_file_path_event2, test_context)
    assert result2['statusCode'] == 400



def test_handler_folder_putobject_returns_400(monkeypatch):
    monkeypatch.setattr(lambda_handler.boto3, 'client', mock_boto3_client)
    monkeypatch.setattr(lambda_handler, 'record_etl_job_run', mock_record_etl_job_run)
    monkeypatch.setenv('DYNAMODB_TABLE_NAME', test_table_name)
    monkeypatch.setenv('SFN_STATE_MACHINE_ARN', test_state_machine_arn)

    result = lambda_handler.lambda_handler(test_bad_folder_event, test_context)
    assert result['statusCode'] == 400


@mock_dynamodb
def test_record_etl_job_run_records_status(monkeypatch, use_moto):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    table = use_moto()

    lambda_handler.record_etl_job_run(test_table_name, test_state_machine_arn, test_execution_id, 'test-execution', '{}', 'testUser', 'testipaddress')

    item = table.get_item(
        TableName=test_table_name,
        Key={ 'execution_id': test_execution_id }
    )
    assert item['Item']['job_latest_status'] == 'STARTED'


@mock_dynamodb
def test_record_etl_jon_run_logs_no_table(monkeypatch, caplog):
    monkeypatch.setenv('AWS_DEFAULT_REGION', mock_region)
    # Purposely do not call use_moto() to create the table

    with pytest.raises(botocore.exceptions.ClientError) as e_info:
        lambda_handler.record_etl_job_run(test_table_name, test_state_machine_arn, test_execution_id, 'test-execution', '{}', 'testUser', 'testipaddress')

    assert e_info.match('ResourceNotFoundException'), 'Expected Boto3 Client Error not raised'

    with caplog.at_level(logging.ERROR):
        assert caplog.text.find('Requested resource not found') != -1
test/test_state_machine_trigger.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk
from aws_cdk.assertions import Template

from test.boto_mocking_helper import *
from lib.step_functions_stack import StepFunctionsStack
from lib.dynamodb_stack import DynamoDbStack
from lib.glue_stack import GlueStack

import lib.configuration as configuration
from lib.configuration import (
    DEV, PROD, TEST
)


def test_resource_types_and_counts(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	# Use one DynamoDbStack stack for all 3 environments since it is
	# not the test subject
	dynamodb_stack = DynamoDbStack(
		app,
		'DynamoDbStackForTests',
		target_environment='Dev'
	)

	# Use one GlueStack stack for all 3 environments since it is not
	# the test subject
	glue_stack = GlueStack(
		app,
		'GlueStackForTests',
		target_environment='Dev',
        hash_values_table=dynamodb_stack.hash_values_table,
        value_lookup_table=dynamodb_stack.value_lookup_table,
        multi_lookup_table=dynamodb_stack.value_lookup_table,
		dq_results_table=dynamodb_stack.dq_results_table,
	)

	step_functions_stacks = {}
	for environment in [DEV, TEST, PROD]:
		step_functions_stacks[environment] = StepFunctionsStack(
			app,
			f'{environment}-StepFunctionsStackForTests',
			target_environment=environment,
			collect_to_cleanse_job=glue_stack.collect_to_cleanse_job,
			cleanse_to_consume_job=glue_stack.cleanse_to_consume_job,
			job_audit_table=dynamodb_stack.job_audit_table,
		)

	# All stacks should be generated before calling Template methods
	for environment in step_functions_stacks.keys():
		template = Template.from_stack(step_functions_stacks[environment])

		template.resource_count_is('AWS::StepFunctions::StateMachine', 1)
		# ETL Trigger, Job Audit Log, S3 Deployment CustomResource handler
		template.resource_count_is('AWS::Lambda::Function', 3)
		# 2 Lambda logs, 1 Statemachine log (future: include the CustomResource lambda log)
		template.resource_count_is('AWS::Logs::LogGroup', 3)
		# 3 Lambda roles, 1 Statemachine role
		template.resource_count_is('AWS::IAM::Role', 4)


def test_stack_has_correct_outputs(monkeypatch):
	monkeypatch.setattr(configuration.boto3, 'client', mock_boto3_client)

	app = cdk.App()

	# Use one DynamoDbStack stack for all 3 environments since it is
	# not the test subject
	dynamodb_stack = DynamoDbStack(
		app,
		'DynamoDbStackForTests',
		target_environment='Dev'
	)

	# Use one GlueStack stack for all 3 environments since it is not
	# the test subject
	glue_stack = GlueStack(
		app,
		'GlueStackForTests',
		target_environment='Dev',
        hash_values_table=dynamodb_stack.hash_values_table,
        value_lookup_table=dynamodb_stack.value_lookup_table,
        multi_lookup_table=dynamodb_stack.value_lookup_table,
		dq_results_table=dynamodb_stack.dq_results_table,
	)

	step_functions_stack = StepFunctionsStack(
		app,
		f'Dev-StepFunctionsStackForTests',
		target_environment='Dev',
		collect_to_cleanse_job=glue_stack.collect_to_cleanse_job,
		cleanse_to_consume_job=glue_stack.cleanse_to_consume_job,
		job_audit_table=dynamodb_stack.job_audit_table,
	)

	template = Template.from_stack(step_functions_stack)
	stack_outputs = template.find_outputs('*')

	state_machine_output = False
	sns_topic_output = False
	for output_id in stack_outputs.keys():
		output_name = stack_outputs[output_id]['Export']['Name']

		if output_name.find('StateMachineName') != -1:
			state_machine_output = True
		if output_name.find('SnsTopicName') != -1:
			sns_topic_output = True

	assert state_machine_output, 'Missing CF output for step functions state machine'
	assert sns_topic_output, 'Missing CF output for SNS topic'
test/test_step_functions_stack.py
aws-samples/aws-insurancelake-etl
# Copyright Amazon.com and its affiliates; all rights reserved. This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
# SPDX-License-Identifier: MIT-0
import pytest
import aws_cdk as cdk
import aws_cdk.aws_s3 as s3
from aws_cdk.assertions import Template, Match

import lib.tagging as tagging
from lib.tagging import (
	COST_CENTER, TAG_ENVIRONMENT, TEAM, APPLICATION
)
from lib.configuration import (
	ENVIRONMENT, DEPLOYMENT, DEV, PROD, TEST
)

test_environment = DEPLOYMENT
test_id_prefix = 'TestPrefix'
test_resource_prefix = 'testprefix'

def mock_get_all_configurations():
	return { test_environment: {
				ENVIRONMENT: test_environment
			}
		}

def mock_get_logical_id_prefix():
	return test_id_prefix

def mock_get_resource_name_prefix():
	return test_resource_prefix


def test_get_tag(monkeypatch):
	monkeypatch.setattr(tagging, 'get_all_configurations', mock_get_all_configurations)
	monkeypatch.setattr(tagging, 'get_logical_id_prefix', mock_get_logical_id_prefix)
	monkeypatch.setattr(tagging, 'get_resource_name_prefix', mock_get_resource_name_prefix)

	test_tags = tagging.get_tag(APPLICATION, test_environment)
	assert f'{test_id_prefix}Etl' in test_tags

	test_tags = tagging.get_tag(TAG_ENVIRONMENT, test_environment)
	assert test_environment in test_tags


def test_get_tag_missing_environment_error(monkeypatch):
	monkeypatch.setattr(tagging, 'get_all_configurations', mock_get_all_configurations)
	monkeypatch.setattr(tagging, 'get_logical_id_prefix', mock_get_logical_id_prefix)
	monkeypatch.setattr(tagging, 'get_resource_name_prefix', mock_get_resource_name_prefix)

	with pytest.raises(AttributeError) as e_info:
		tagging.get_tag(APPLICATION, 'BadEnvironment')

	assert e_info.match('not found in environment configurations'), \
		'Expected Attribute Error for missing environment not raised'


def test_tagging_stack_resource(monkeypatch):
	monkeypatch.setattr(tagging, 'get_all_configurations', mock_get_all_configurations)
	monkeypatch.setattr(tagging, 'get_logical_id_prefix', mock_get_logical_id_prefix)
	monkeypatch.setattr(tagging, 'get_resource_name_prefix', mock_get_resource_name_prefix)

	app = cdk.App()
	stack = cdk.Stack(app, 'StackForTests')
	s3.Bucket(stack, 'BucketForTests')
	tagging.tag(stack, test_environment)

	template = Template.from_stack(stack)
	template.has_resource_properties(
		'AWS::S3::Bucket',
		Match.object_like(
			{
				"Tags": [
					{
						"Key": f"{test_resource_prefix}:application",
						"Value": f"{test_id_prefix}Etl"
					},
					{
						"Key": f"{test_resource_prefix}:cost-center",
						"Value": f"{test_id_prefix}Etl"
					},
					{
						"Key": f"{test_resource_prefix}:environment",
						"Value": test_environment
					},
					{
						"Key": f"{test_resource_prefix}:team",
						"Value": f"{test_id_prefix}Admin"
					}
				]
			}
		)
	)
test/test_tagging.py
aws-samples/aws-insurancelake-etl