majuna/app/terraform/terraform.py

import json
import subprocess
from typing import Any, Dict, List, Optional, Tuple

import jinja2

from app.terraform import BaseAutomation


class TerraformAutomation(BaseAutomation):
    """
    An abstract class to be extended by automation plugins using Terraform
    providers to deploy resources.
    """

    parallelism = 10
    """
    Default parallelism for remote API calls.
    """

    def automate(self, full: bool = False):
        prehook_result = self.tf_prehook()
        self.tf_generate()
        self.tf_init()
        returncode, logs = self.tf_apply(refresh=full)
        self.tf_posthook(prehook_result=prehook_result)
        return True if returncode == 0 else False, logs

    def tf_apply(self, *,
                 refresh: bool = True,
                 parallelism: Optional[int] = None,
                 lock_timeout: int = 15) -> Tuple[int, List[Dict[str, Any]]]:
        if not parallelism:
            parallelism = self.parallelism
        tf = subprocess.run(
            ['terraform',
             'apply',
             '-auto-approve',
             '-json',
             f'-refresh={str(refresh).lower()}',
             f'-parallelism={str(parallelism)}',
             f'-lock-timeout={str(lock_timeout)}m',
             ],
            cwd=self.working_directory(),
            stdout=subprocess.PIPE)
        logs = []
        for line in tf.stdout.decode('utf-8').split('\n'):
            if line.strip():
                logs.append(json.loads(line))
        return tf.returncode, str(logs)

    def tf_generate(self):
        raise NotImplementedError()

    def tf_init(self, *,
                lock_timeout: int = 15):
        # The init command does not support JSON output
        subprocess.run(
            ['terraform',
             'init',
             f'-lock-timeout={str(lock_timeout)}m',
             ],
            cwd=self.working_directory())

    def tf_output(self) -> Dict[str, Any]:
        tf = subprocess.run(
            ['terraform', 'output', '-json'],
            cwd=self.working_directory(),
            stdout=subprocess.PIPE)
        return json.loads(tf.stdout)

    def tf_plan(self, *,
                refresh: bool = True,
                parallelism: Optional[int] = None,
                lock_timeout: int = 15):
        tf = subprocess.run(
            ['terraform',
             'plan',
             '-json',
             f'-refresh={str(refresh).lower()}',
             f'-parallelism={str(parallelism)}',
             f'-lock-timeout={str(lock_timeout)}m',
             ],
            cwd=self.working_directory())
        logs = []
        for line in tf.stdout.decode('utf-8').split('\n'):
            if line.strip():
                logs.append(json.loads(line))
        return tf.returncode, str(logs)

    def tf_posthook(self, *, prehook_result: Any = None) -> None:
        """
        This hook function is called as part of normal automation, after the
        completion of :func:`tf_apply`.

        The default, if not overridden by a subclass, is to do nothing.

        :param prehook_result: the returned value of :func:`tf_prehook`
        :return: None
        """
        pass

    def tf_prehook(self) -> Optional[Any]:
        """
        This hook function is called as part of normal automation, before generating
        the terraform configuration file. The return value will be passed to
        :func:`tf_posthook` but is otherwise ignored.

        The default, if not overridden by a subclass, is to do nothing.

        :return: state that is useful to :func:`tf_posthook`, if required
        """
        pass

    def tf_show(self) -> Dict[str, Any]:
        terraform = subprocess.run(
            ['terraform', 'show', '-json'],
            cwd=self.working_directory(),
            stdout=subprocess.PIPE)
        return json.loads(terraform.stdout)

    def tf_write(self, template: str, **kwargs):
        tmpl = jinja2.Template(template)
        with open(self.working_directory("main.tf"), 'w') as tf:
            tf.write(tmpl.render(**kwargs))