elasticsearch.py

from troposphere import NoValue, Tags, Sub
from troposphere import (
    ec2,
    elasticsearch,
    iam,
)

from stacker.blueprints.base import Blueprint
from stacker.blueprints.variables.types import TroposphereType


class ServiceLinkedRole(Blueprint):
    def create_service_linked_role(self):
        t = self.template
        role = t.add_resource(
            iam.ServiceLinkedRole(
                "ServiceLinkRole",
                AWSServiceName="es.amazonaws.com",
            )
        )

        self.add_output("RoleId", role.Ref())

    def create_template(self):
        self.create_service_linked_role()


class Domain(Blueprint):
    VARIABLES = {
        "EBSOptions": {
            "type": TroposphereType(elasticsearch.EBSOptions, optional=True),
            "description": "Optional EBSOptions config for the Domain.",
            "default": None,
        },
        "ElasticsearchClusterConfig": {
            "type": TroposphereType(
                elasticsearch.ElasticsearchClusterConfig, optional=True
            ),
            "description": "Optional ElasticsearchClusterConfig config for "
                           "the Domain.",
            "default": None,
        },
        "ElasticsearchVersion": {
            "type": str,
            "description": "The version of elasticsearch for the Domain.",
            "default": "",
        },
        "EncryptionAtRestOptions": {
            "type": TroposphereType(
                elasticsearch.EncryptionAtRestOptions, optional=True
            ),
            "description": "Optional EncryptionAtRestOptions config for "
                           "the Domain.",
            "default": None,
        },
        "SnapshotOptions": {
            "type": TroposphereType(
                elasticsearch.SnapshotOptions, optional=True
            ),
            "description": "Optional SnapshotOptions config for "
                           "the Domain.",
            "default": None,
        },
        "Tags": {
            "type": dict,
            "description": "An optional dictionary of tags to put on the "
                           "Domain.",
            "default": {},
        },
        "VpcId": {
            "type": str,
            "description": "The VpcId to deploy the Domain into.  Only "
                           "required if no SecurityGroupIds are specified.",
            "default": "",
        },
        "SecurityGroupIds": {
            "type": list,
            "description": "An optional list of security group ids to add "
                           "the domain to. If not specified, one will be "
                           "automatically created.",
            "default": [],
        },
        "Subnets": {
            "type": list,
            "description": "A required list of subnet Ids to deploy the "
                           "Domain into.",
        },
    }

    @property
    def ebs_options(self):
        return self.get_variables()["EBSOptions"] or NoValue

    @property
    def elasticsearch_cluster_config(self):
        return self.get_variables()["ElasticsearchClusterConfig"] or NoValue

    @property
    def elasticsearch_version(self):
        return self.get_variables()["ElasticsearchVersion"] or NoValue

    @property
    def encryption_at_rest_options(self):
        return self.get_variables()["EncryptionAtRestOptions"] or NoValue

    @property
    def snapshot_options(self):
        return self.get_variables()["SnapshotOptions"] or NoValue

    @property
    def tags(self):
        tags = self.get_variables()["Tags"]
        if tags:
            return Tags(**tags)
        return NoValue

    @property
    def vpc_id(self):
        return self.get_variables()["VpcId"]

    @property
    def security_group_ids(self):
        return self.get_variables()["SecurityGroupIds"]

    @property
    def subnet_ids(self):
        return self.get_variables()["Subnets"]

    def create_security_group(self):
        if self.security_group_ids:
            self.security_groups = self.security_group_ids
            return

        t = self.template
        sg = t.add_resource(
            ec2.SecurityGroup(
                "SecurityGroup",
                GroupDescription=Sub("${AWS::StackName}"),
                VpcId=self.vpc_id,
            )
        )

        self.security_groups = [sg.Ref()]

        self.add_output("SecurityGroupId", sg.Ref())

    def create_domain(self):
        t = self.template

        vpc_options = elasticsearch.VPCOptions(
            SecurityGroupIds=self.security_groups,
            SubnetIds=self.subnet_ids,
        )

        domain = t.add_resource(
            elasticsearch.Domain(
                "Domain",
                EBSOptions=self.ebs_options,
                ElasticsearchClusterConfig=self.elasticsearch_cluster_config,
                ElasticsearchVersion=self.elasticsearch_version,
                EncryptionAtRestOptions=self.encryption_at_rest_options,
                SnapshotOptions=self.snapshot_options,
                Tags=self.tags,
                VPCOptions=vpc_options,
            )
        )

        self.add_output("DomainId", domain.Ref())
        self.add_output("DomainArn", domain.GetAtt("DomainArn"))
        self.add_output("DomainEndpoint", domain.GetAtt("DomainEndpoint"))

    def create_template(self):
        self.create_security_group()
        self.create_domain()