AWSTemplateFormatVersion: '2010-09-09'

Mappings:
  S3Buckets:
    us-west-2:
      bucket: aws-admartech-samples-us-west-2
    us-east-1:
      bucket: aws-admartech-samples-us-east-1
    us-east-2:
      bucket: aws-admartech-samples-us-east-2
    eu-west-1:
      bucket: aws-admartech-samples-eu-west-1

  Constants:
    S3Keys:
      neptuneNotebooks: /identity-resolution/notebooks/identity-graph
      irdata: /identity-resolution/data/
      bulkLoadStack: /identity-resolution/templates/bulk-load-stack.yaml
      neptuneNotebookStack: /identity-resolution/templates/neptune-workbench-stack.yaml

  #------------------------------------------------------------------------------#
  # RESOURCES
  #------------------------------------------------------------------------------#
Resources:
# ---------- CREATING NEPTUNE CLUSTER FROM SNAPSHOT ----------
  NeptuneBaseStack:
    Type: AWS::CloudFormation::Stack
    Properties:
      TemplateURL: https://s3.amazonaws.com/aws-neptune-customer-samples/v2/cloudformation-templates/neptune-base-stack.json
      Parameters:
        NeptuneQueryTimeout: '300000'
        DbInstanceType: db.r5.12xlarge
      TimeoutInMinutes: '360'

# ---------- SETTING UP SAGEMAKER NOTEBOOK INSTANCES ----------
  ExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: "/"
      Policies:
        - PolicyName: "sagemakerneptunepolicy"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - cloudwatch:PutMetricData
                Resource:
                  Fn::Sub: "arn:${AWS::Partition}:cloudwatch:${AWS::Region}:${AWS::AccountId}:*"
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogGroup"
                  - "logs:CreateLogStream"
                  - "logs:DescribeLogStreams"
                  - "logs:PutLogEvents"
                  - "logs:GetLogEvents"
                Resource:
                  Fn::Sub: "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:*"
              - Effect: "Allow"
                Action: "neptune-db:connect"
                Resource:
                  Fn::Sub: "arn:${AWS::Partition}:neptune-db:${AWS::Region}:${AWS::AccountId}:${NeptuneBaseStack.Outputs.DBClusterId}/*"
              - Effect: "Allow"
                Action:
                  - "s3:Get*"
                  - "s3:List*"
                Resource:
                  Fn::Sub: "arn:${AWS::Partition}:s3:::*"

  SageMakerNeptuneStack:
    Type: AWS::CloudFormation::Stack
    Properties:
      TemplateURL:
        Fn::Join: [ "",
          [
            https://s3.amazonaws.com/,
            !FindInMap [ S3Buckets, Ref: 'AWS::Region', bucket ],
            !FindInMap [ Constants, S3Keys, neptuneNotebookStack ]
          ]
        ]
      Parameters:
        SageMakerNotebookName: "id-graph-notebook"
        NotebookInstanceType: ml.m5.xlarge
        NeptuneClusterEndpoint:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.DBClusterEndpoint
        NeptuneClusterPort:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.DBClusterPort
        NeptuneClusterSecurityGroups:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.NeptuneSG
        NeptuneClusterSubnetId:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.PublicSubnet1
        SageMakerNotebookRole:
          Fn::GetAtt:
          - ExecutionRole
          - Arn
        AdditionalNotebookS3Locations: !Join
                                          - ''
                                          - - 's3://'
                                            - !FindInMap
                                              - S3Buckets
                                              - !Ref 'AWS::Region'
                                              - bucket
                                            - !FindInMap
                                              - Constants
                                              - S3Keys
                                              - neptuneNotebooks
      TimeoutInMinutes: '60'

# --------- LOAD DATA INTO NEPTUNE ---------

  NeptuneBulkLoadStack:
    Type: AWS::CloudFormation::Stack
    Properties:
      TemplateURL: !Join
                    - ''
                    - - 'https://s3.'
                      - !Ref 'AWS::Region'
                      - '.amazonaws.com/'
                      - !FindInMap
                        - S3Buckets
                        - !Ref 'AWS::Region'
                        - bucket
                      - !FindInMap
                        - Constants
                        - S3Keys
                        - bulkLoadStack
      Parameters:
        bulkloadNeptuneEndpoint:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.DBClusterEndpoint
        bulkloadNeptuneData: !Join
                          - ''
                          - - 's3://'
                            - !FindInMap
                              - S3Buckets
                              - !Ref 'AWS::Region'
                              - bucket
                            - !FindInMap
                              - Constants
                              - S3Keys
                              - irdata
        bulkloadNeptuneIAMRole:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.NeptuneLoadFromS3IAMRoleArn
        bulkloadNeptuneSecurityGroup:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.NeptuneSG
        bulkloadSubnet1:
          Fn::GetAtt:
          - NeptuneBaseStack
          - Outputs.PrivateSubnet1
        bulkloadBucket: !FindInMap
                          - S3Buckets
                          - !Ref 'AWS::Region'
                          - bucket
          

  #------------------------------------------------------------------------------#
  # OUTPUTS
  #------------------------------------------------------------------------------#

Outputs:
  VPC:
    Description: VPC of the Neptune Cluster
    Value:
      Fn::GetAtt:
        - NeptuneBaseStack
        - Outputs.VPC
  PublicSubnet1:
    Value:
      Fn::GetAtt:
        - NeptuneBaseStack
        - Outputs.PublicSubnet1
  NeptuneSG:
    Description: Neptune Security Group
    Value:
      Fn::GetAtt:
        - NeptuneBaseStack
        - Outputs.NeptuneSG
  SageMakerNotebook:
    Value:
      Fn::GetAtt:
      - SageMakerNeptuneStack
      - Outputs.NeptuneNotebook
  DBClusterEndpoint:
    Description: Master Endpoint for Neptune Cluster
    Value:
      Fn::GetAtt:
      - NeptuneBaseStack
      - Outputs.DBClusterEndpoint
  DBInstanceEndpoint:
    Description: Master Instance Endpoint
    Value:
      Fn::GetAtt:
      - NeptuneBaseStack
      - Outputs.DBInstanceEndpoint
  GremlinEndpoint:
    Description: Gremlin Endpoint for Neptune
    Value:
      Fn::GetAtt:
      - NeptuneBaseStack
      - Outputs.GremlinEndpoint
  LoaderEndpoint:
    Description: Loader Endpoint for Neptune
    Value:
      Fn::GetAtt:
      - NeptuneBaseStack
      - Outputs.LoaderEndpoint
  DBClusterReadEndpoint:
    Description: DB cluster Read Endpoint
    Value:
      Fn::GetAtt:
      - NeptuneBaseStack
      - Outputs.DBClusterReadEndpoint