AWSTemplateFormatVersion: 2010-09-09

Description: >-
  Set up resources and deploy a Databricks workspace in your AWS account.
  If you encounter any errors during the process, reach out to quickstart-support@databricks.com.
  Visit this Databricks Community post for troubleshooting guidance: https://dbricks.co/AWSQuickStartHelp.

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: "Databricks Account Credentials"
        Parameters:
          - AccountId
          - SessionToken
      - Label:
          default: "Workspace Configuration"
        Parameters:
          - WorkspaceName
    ParameterLabels:
      AccountId:
        default: Databricks account ID
      SessionToken:
        default: Browser session token
      WorkspaceName:
        default: Workspace name

Outputs:
  DatabricksAccountId:
    Description: Databricks account ID
    Value: !GetAtt createCredentials.ExternalId
  WorkspaceIamRoleArn:
    Description: ARN of the workspace cross-account IAM role
    Value: !GetAtt workspaceIamRole.Arn
  WorkspaceS3BucketName:
    Description: Name of the workspace S3 root bucket
    Value: !Ref workspaceS3Bucket
  CatalogIamRoleArn:
    Description: ARN of the IAM role for catalog
    Value: !GetAtt catalogIamRole.Arn
  CredentialsId:
    Description: Credential ID
    Value: !GetAtt createCredentials.CredentialsId
  StorageConfigId:
    Description: Storage configuration ID
    Value: !GetAtt createStorageConfiguration.StorageConfigId
  WorkspaceURL:
    Description: URL of the workspace
    Value: !Join
      - ""
      - - "https://"
        - !GetAtt createWorkspace.DeploymentName
        - ".cloud.databricks.com"
  WorkspaceStatus:
    Description: Status of the requested workspace
    Value: !GetAtt createWorkspace.WorkspaceStatus
  WorkspaceStatusMessage:
    Description: Detailed status description of the requested workspace
    Value: !GetAtt createWorkspace.WorkspaceStatusMsg
  PricingTier:
    Description: The pricing tier of the workspace. See https://databricks.com/product/aws-pricing.
    Value: !GetAtt createWorkspace.PricingTier

Parameters:
  AccountId:
    Description: "Find your account ID at https://accounts.cloud.databricks.com."
    AllowedPattern: "^[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}$"
    MinLength: "36"
    Type: String
    Default: aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee
  SessionToken:
    Description: "Browser session token."
    NoEcho: "true"
    Type: String
  WorkspaceName:
    Description: "Human-readable name for this workspace."
    MinLength: "1"
    MaxLength: "64"
    Type: String

Resources:
  # Workspace cross-account IAM role for Databricks-created and -managed VPC.
  workspaceIamRole:
    Type: "AWS::IAM::Role"
    Metadata:
      cfn-lint:
        config:
          ignore_checks:
            - EIAMPolicyWildcardResource
          ignore_reasons:
            EIAMPolicyWildcardResource: "Need to manage a databricks workspace"
    Properties:
      RoleName: !Sub ${AWS::StackName}-role
      AssumeRolePolicyDocument:
        Statement:
          - Action: "sts:AssumeRole"
            Condition:
              StringEquals:
                "sts:ExternalId": !Sub ${AccountId}
            Effect: Allow
            Principal:
              AWS: arn:aws:iam::414351767826:root
            Sid: ""
        Version: 2012-10-17
      Path: /
      Policies:
        - PolicyDocument:
            Statement:
              - Sid: Stmt1403287045000
                Effect: Allow
                Action:
                  - "ec2:AllocateAddress"
                  - "ec2:AssignPrivateIpAddresses"
                  - "ec2:AssociateDhcpOptions"
                  - "ec2:AssociateIamInstanceProfile"
                  - "ec2:AssociateRouteTable"
                  - "ec2:AttachInternetGateway"
                  - "ec2:AttachVolume"
                  - "ec2:AuthorizeSecurityGroupEgress"
                  - "ec2:AuthorizeSecurityGroupIngress"
                  - "ec2:CancelSpotInstanceRequests"
                  - "ec2:CreateDhcpOptions"
                  - "ec2:CreateFleet"
                  - "ec2:CreateInternetGateway"
                  - "ec2:CreateLaunchTemplate"
                  - "ec2:CreateLaunchTemplateVersion"
                  - "ec2:CreateNatGateway"
                  - "ec2:CreateRoute"
                  - "ec2:CreateRouteTable"
                  - "ec2:CreateSecurityGroup"
                  - "ec2:CreateSubnet"
                  - "ec2:CreateTags"
                  - "ec2:CreateVolume"
                  - "ec2:CreateVpc"
                  - "ec2:CreateVpcEndpoint"
                  - "ec2:DeleteDhcpOptions"
                  - "ec2:DeleteFleets"
                  - "ec2:DeleteInternetGateway"
                  - "ec2:DeleteLaunchTemplate"
                  - "ec2:DeleteLaunchTemplateVersions"
                  - "ec2:DeleteNatGateway"
                  - "ec2:DeleteRoute"
                  - "ec2:DeleteRouteTable"
                  - "ec2:DeleteSecurityGroup"
                  - "ec2:DeleteSubnet"
                  - "ec2:DeleteTags"
                  - "ec2:DeleteVolume"
                  - "ec2:DeleteVpc"
                  - "ec2:DeleteVpcEndpoints"
                  - "ec2:DescribeAvailabilityZones"
                  - "ec2:DescribeFleetHistory"
                  - "ec2:DescribeFleetInstances"
                  - "ec2:DescribeFleets"
                  - "ec2:DescribeIamInstanceProfileAssociations"
                  - "ec2:DescribeInstanceStatus"
                  - "ec2:DescribeInstances"
                  - "ec2:DescribeInternetGateways"
                  - "ec2:DescribeLaunchTemplates"
                  - "ec2:DescribeLaunchTemplateVersions"
                  - "ec2:DescribeNatGateways"
                  - "ec2:DescribePrefixLists"
                  - "ec2:DescribeReservedInstancesOfferings"
                  - "ec2:DescribeRouteTables"
                  - "ec2:DescribeSecurityGroups"
                  - "ec2:DescribeSpotInstanceRequests"
                  - "ec2:DescribeSpotPriceHistory"
                  - "ec2:DescribeSubnets"
                  - "ec2:DescribeVolumes"
                  - "ec2:DescribeVpcs"
                  - "ec2:DetachInternetGateway"
                  - "ec2:DisassociateIamInstanceProfile"
                  - "ec2:DisassociateRouteTable"
                  - "ec2:GetLaunchTemplateData"
                  - "ec2:GetSpotPlacementScores"
                  - "ec2:ModifyFleet"
                  - "ec2:ModifyLaunchTemplate"
                  - "ec2:ModifyVpcAttribute"
                  - "ec2:ReleaseAddress"
                  - "ec2:ReplaceIamInstanceProfileAssociation"
                  - "ec2:RequestSpotInstances"
                  - "ec2:RevokeSecurityGroupEgress"
                  - "ec2:RevokeSecurityGroupIngress"
                  - "ec2:RunInstances"
                  - "ec2:TerminateInstances"
                Resource:
                  - "*"
              - Effect: Allow
                Action:
                  - "iam:CreateServiceLinkedRole"
                  - "iam:PutRolePolicy"
                Resource:
                  - !Sub arn:${AWS::Partition}:iam::*:role/aws-service-role/spot.amazonaws.com/AWSServiceRoleForEC2Spot
                Condition:
                  StringLike:
                    "iam:AWSServiceName": spot.amazonaws.com
            Version: 2012-10-17
          PolicyName: !Sub ${AWS::StackName}-policy
      Tags:
        - Key: Name
          Value: "databricks-trial-quickstart-cloud-formation"

  # Workspace S3 root bucket requirements.
  workspaceS3Bucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub ${AWS::StackName}-bucket
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
  workspaceS3BucketPolicy:
    Type: "AWS::S3::BucketPolicy"
    Properties:
      PolicyDocument:
        Id: MyPolicy
        Version: 2012-10-17
        Statement:
          - Sid: Grant Databricks Access
            Effect: Allow
            Principal:
              AWS: arn:aws:iam::414351767826:root
            Action:
              - "s3:GetObject"
              - "s3:GetObjectVersion"
              - "s3:PutObject"
              - "s3:DeleteObject"
              - "s3:ListBucket"
              - "s3:GetBucketLocation"
            Resource:
              - !Sub arn:${AWS::Partition}:s3:::${workspaceS3Bucket}/*
              - !Sub arn:${AWS::Partition}:s3:::${workspaceS3Bucket}
            Condition:
              StringEquals:
                "aws:PrincipalTag/DatabricksAccountId": !Sub ${AccountId}
          - Sid: Prevent DBFS from accessing Unity Catalog metastore
            Effect: Deny
            Principal:
              AWS: arn:aws:iam::414351767826:root
            Action:
              - "s3:*"
            Resource:
              - !Sub arn:${AWS::Partition}:s3:::${workspaceS3Bucket}/unity-catalog/*
              - !Sub arn:${AWS::Partition}:s3:::${workspaceS3Bucket}/unity-catalog
      Bucket: !Ref workspaceS3Bucket

  # IAM role for catalog.
  catalogIamRole:
    Type: "AWS::IAM::Role"
    Metadata:
      cfn-lint:
        config:
          ignore_checks:
            - EIAMPolicyWildcardResource
          ignore_reasons:
            EIAMPolicyWildcardResource: "Need to manage databricks metastore"
    Properties:
      RoleName: !Sub ${AWS::StackName}-catalog-role
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              AWS:
                - arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL
            Action:
              - "sts:AssumeRole"
            Condition:
              StringEquals:
                "sts:ExternalId": !Sub ${AccountId}
          - Effect: Allow
            Principal:
              AWS:
                - !Sub arn:aws:iam::${AWS::AccountId}:root
            Action:
              - "sts:AssumeRole"
            Condition:
              StringEquals:
                "AWS:PrincipalArn": !Sub arn:aws:iam::${AWS::AccountId}:role/${AWS::StackName}-catalog-role
                "sts:ExternalId": !Sub ${AccountId}
        Version: 2012-10-17
      Path: /
      Tags:
        - Key: Name
          Value: "databricks-trial-quickstart-cloud-formation"
  catalogIamRolePolicy:
    DependsOn: catalogIamRole
    Type: "AWS::IAM::Policy"
    Properties:
      PolicyName: !Sub ${AWS::StackName}-catalog-policy
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action:
              - "s3:GetObject"
              - "s3:PutObject"
              - "s3:DeleteObject"
              - "s3:ListBucket"
              - "s3:GetBucketLocation"
              - "s3:GetLifecycleConfiguration"
              - "s3:PutLifecycleConfiguration"
            Resource:
              - !Sub arn:aws:s3:::${workspaceS3Bucket}/*
              - !Sub arn:aws:s3:::${workspaceS3Bucket}
          - Effect: Allow
            Action: sts:AssumeRole
            Resource: !Sub arn:aws:iam::${AWS::AccountId}:role/${AWS::StackName}-catalog-role
      Roles:
        - !Sub ${AWS::StackName}-catalog-role

  # Databricks API for workspace credentials creation.
  createCredentials:
    Type: Custom::CreateCredentials
    Properties:
      ServiceToken: !GetAtt "databricksApiFunction.Arn"
      action: "CREATE_CREDENTIALS"
      accountId: !Ref AccountId
      credentials_name: !Sub ${AWS::StackName}-credentials
      role_arn: !GetAtt "workspaceIamRole.Arn"
      encodedbase64: !Ref SessionToken
      user_agent: "databricks-CloudFormation-Trial-inhouse-uc-by-default-via-backend-oauth"

  # Databricks API for workspace storage configuration creation.
  createStorageConfiguration:
    Type: Custom::CreateStorageConfigurations
    Properties:
      ServiceToken: !GetAtt "databricksApiFunction.Arn"
      action: "CREATE_STORAGE_CONFIGURATIONS"
      accountId: !Ref AccountId
      storage_config_name: !Sub ${AWS::StackName}-storage
      s3bucket_name: !Ref workspaceS3Bucket
      role_arn: !GetAtt catalogIamRole.Arn
      encodedbase64: !Ref SessionToken
      user_agent: "databricks-CloudFormation-Trial-inhouse-uc-by-default-via-backend-oauth"

  # Databricks API for workspace creation.
  createWorkspace:
    Type: Custom::CreateWorkspace
    Properties:
      ServiceToken: !GetAtt "databricksApiFunction.Arn"
      action: "CREATE_WORKSPACES"
      accountId: !Ref AccountId
      workspace_name: !Ref WorkspaceName
      deployment_name: ""
      aws_region: !Ref AWS::Region
      credentials_id: !GetAtt createCredentials.CredentialsId
      storage_config_id: !GetAtt createStorageConfiguration.StorageConfigId
      encodedbase64: !Ref SessionToken
      network_id: ""
      customer_managed_key_id: ""
      pricing_tier: ""
      customer_name: ""
      authoritative_user_email: ""
      authoritative_user_full_name: ""
      user_agent: "databricks-CloudFormation-Trial-inhouse-uc-by-default-via-backend-oauth"

  databricksApiFunction:
    DependsOn: CopyZips
    Type: AWS::Lambda::Function
    Properties:
      Description: Databricks Account API
      Handler: rest_client.handler
      Runtime: python3.11
      Role: !GetAtt "functionRole.Arn"
      Timeout: 900
      Code:
        S3Bucket: !Ref LambdaZipsBucket
        S3Key: "functions/packages/uc-by-default-via-backend-oauth/lambda.zip"

  # IAM Role for lambda function execution.
  functionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - !Sub arn:${AWS::Partition}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
        - !Sub arn:${AWS::Partition}:iam::aws:policy/IAMFullAccess
      Tags:
        - Key: Name
          Value: "databricks-CloudFormation-Trial-provider"

  # Resources to stage lambda.zip file.
  LambdaZipsBucket:
    Type: AWS::S3::Bucket
  CopyZips:
    Type: Custom::CopyZips
    Properties:
      ServiceToken: !GetAtt "CopyZipsFunction.Arn"
      DestBucket: !Ref LambdaZipsBucket
      SourceBucket: "databricks-prod-public-cfts"
      Prefix: ""
      Objects:
        - functions/packages/uc-by-default-via-backend-oauth/lambda.zip
  CopyZipsRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - !Sub arn:${AWS::Partition}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Path: /
      Policies:
        - PolicyName: lambda-copier
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - s3:GetObject
                Resource:
                  - !Sub arn:${AWS::Partition}:s3:::databricks-prod-public-cfts/*
              - Effect: Allow
                Action:
                  - s3:GetObject
                  - s3:PutObject
                  - s3:DeleteObject
                Resource:
                  - !Sub arn:${AWS::Partition}:s3:::${LambdaZipsBucket}/*
      Tags:
        - Key: Name
          Value: "databricks-CloudFormation-Trial-provider"
  CopyZipsFunction:
    Type: AWS::Lambda::Function
    Properties:
      Description: Copies objects from a source S3 bucket to a destination
      Handler: index.handler
      Runtime: python3.11
      Role: !GetAtt "CopyZipsRole.Arn"
      Timeout: 240
      Code:
        ZipFile: |
          import json
          import logging
          import threading
          import boto3
          import cfnresponse
          def copy_objects(source_bucket, dest_bucket, prefix, objects):
              s3 = boto3.client('s3')
              for o in objects:
                  key = prefix + o
                  copy_source = {
                      'Bucket': source_bucket,
                      'Key': key
                  }
                  print('copy_source: %s' % copy_source)
                  print('dest_bucket = %s'%dest_bucket)
                  print('key = %s' %key)
                  s3.copy_object(CopySource=copy_source, Bucket=dest_bucket,
                        Key=key)
          def delete_objects(bucket, prefix, objects):
              s3 = boto3.client('s3')
              objects = {'Objects': [{'Key': prefix + o} for o in objects]}
              s3.delete_objects(Bucket=bucket, Delete=objects)
          def timeout(event, context):
              logging.error('Execution is about to time out, sending failure response to CloudFormation')
              cfnresponse.send(event, context, cfnresponse.FAILED, {}, None)
          def handler(event, context):
              # make sure we send a failure to CloudFormation if the function
              # is going to timeout
              timer = threading.Timer((context.get_remaining_time_in_millis()
                        / 1000.00) - 0.5, timeout, args=[event, context])
              timer.start()
              print('Received event: %s' % json.dumps(event))
              status = cfnresponse.SUCCESS
              try:
                  source_bucket = event['ResourceProperties']['SourceBucket']
                  dest_bucket = event['ResourceProperties']['DestBucket']
                  prefix = event['ResourceProperties']['Prefix']
                  objects = event['ResourceProperties']['Objects']
                  if event['RequestType'] == 'Delete':
                      delete_objects(dest_bucket, prefix, objects)
                  else:
                      copy_objects(source_bucket, dest_bucket, prefix, objects)
              except Exception as e:
                  logging.error('Exception: %s' % e, exc_info=True)
                  status = cfnresponse.FAILED
              finally:
                  timer.cancel()
                  cfnresponse.send(event, context, status, {}, None)
