doczyai-pipelines/archive/terraform.py

from terminal import run_command, prepare_logger, decorate_white_bold, decorate_warn
from git_diff_changes import is_deploy_module

import os
import argparse
import sys
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

script_dir = os.path.dirname(os.path.abspath(__file__))

LOGGER = prepare_logger()

LOGGER.info(sys.argv)

parser = argparse.ArgumentParser(
    description="Run Terraform init with dynamic backend configuration."
)
parser.add_argument(
    "terraform_command",
    type=str,
    choices=["plan", "apply"],
    help="The Terraform command to execute (plan, apply, init).",
)
parser.add_argument(
    "--environment",
    type=str,
    required=True,
    help="The environment name i.e. dev, prod.",
)
parser.add_argument("--module", type=str, required=False, help="The module name.")
parser.add_argument(
    "--init-extra-args", type=str, required=False, help="Terraform init extra arguments"
)
parser.add_argument(
    "--plan-extra-args", type=str, required=False, help="Terraform plan extra arguments"
)
parser.add_argument(
    "--apply-extra-args",
    type=str,
    required=False,
    help="Terraform plan extra arguments",
)
parser.add_argument(
    "--threads", type=int, default=4, help="Number of threads to use (default: 4)"
)

args = parser.parse_args()

raw_terraform_command = args.terraform_command
environment = args.environment
init_extra_args = args.init_extra_args
plan_extra_args = args.plan_extra_args
apply_extra_args = args.apply_extra_args
terraform_threads = args.threads
module_name = args.module

bitbucket_ci = os.getenv("BITBUCKET_CI")


def change_dir(path):
    if path:
        _, git_root = run_command(
            "git rev-parse --show-toplevel", log_cmd=True, log_output=True
        )
        LOGGER.info(f"git_root={git_root}")
        absolute_path = os.path.join(git_root[0], path)
        return absolute_path
        # os.chdir(absolute_path)
    else:
        LOGGER.error("TF_ROOT_MODULE_PATH environment variable is not set.")
        sys.exit(1)


def select_vars_file():
    tf_vars_file = os.getenv("TFVARS_FILE_PATH", f"vars-{environment}.tfvars")
    if tf_vars_file:
        LOGGER.info(f"Using {tf_vars_file} as vars-file...")
    else:
        LOGGER.error("TFVARS_FILE_PATH variable is not set.")
        sys.exit(1)
    return tf_vars_file


def prepare_tf_init_command(
    stack_path, init_extra_args, state_environment, tf_state_bucket_name
):
    command = (
        f"terraform init --reconfigure "
        f'-backend-config="dynamodb_table=doczyai-use2-{get_environment_value(environment)}-infra-dyd-terraform-lock" '
        f'-backend-config="bucket=doczyai-use2-{get_environment_value(environment)}-infra-s3-terraform-state" '
        f'-backend-config="key=terraform/{stack_path}/terraform.tfstate" -backend-config="region=us-east-2"'
    )
    if init_extra_args:
        command = f"{command} {init_extra_args}"
    return command


def get_environment_value(environment):
    environment_map = {"prod": "p", "uat": "u", "qa": "q", "dev": "d"}

    return environment_map.get(
        environment, None
    )  # Returns None if the environment is not found


def prepare_tf_main_command(
    environment,
    stack_path,
    state_environment,
    tf_state_bucket_name,
    vars_file,
    plan_extra_args,
    apply_extra_args,
):
    common_command = f'-var-file=vars-{environment}.tfvars -var "aws_region=us-east-2" -var "environment={environment}"'
    if raw_terraform_command == "plan":
        path_part = stack_path.split("/")[1] if "/" in stack_path else stack_path
        command = (
            f"terraform plan -out .{path_part}.{environment}.tfplan {common_command}"
        )
        if plan_extra_args:
            command = f"{command} {plan_extra_args}"
    elif raw_terraform_command == "apply":
        command = f"terraform apply -auto-approve {common_command}"
        if apply_extra_args:
            command = f"{command} {apply_extra_args}"
    return command


tf_root_module_path = (
    os.getenv("TF_ROOT_MODULE_PATH") if not module_name else module_name
)


def do_terraform(stack_path):
    LOGGER.info(f"[{stack_path}]: stack_path={stack_path}")
    is_deploy_module_path = is_deploy_module(stack_path, "")
    LOGGER.info(
        f"[{stack_path}]: is_deploy_module_path={decorate_white_bold(is_deploy_module_path)}"
    )
    if not is_deploy_module_path:
        LOGGER.info(
            decorate_warn(
                f"[{stack_path}]: Module is not found in git changes. Skipping deploy."
            )
        )
        return True, "OK"

    absolute_path = change_dir(stack_path)
    # TF INIT
    state_environment = "dev-new" if environment == "dev" else environment
    tf_state_bucket_name = f"pi2new-{state_environment}-terraform-state-file"

    terraform_init_command = prepare_tf_init_command(
        stack_path, init_extra_args, state_environment, tf_state_bucket_name
    )
    rc_init, init_result = run_command(
        terraform_init_command,
        log_output=True,
        log_cmd=True,
        log_prefix=stack_path,
        cwd=absolute_path,
    )

    if rc_init > 0:
        raise Exception(f"{stack_path}: Failed to run terraform init")

    vars_file = select_vars_file()
    terraform_main_command = prepare_tf_main_command(
        environment,
        stack_path,
        state_environment,
        tf_state_bucket_name,
        vars_file,
        plan_extra_args,
        apply_extra_args,
    )

    rc_main, command_result = run_command(
        terraform_main_command,
        log_output=True,
        log_cmd=True,
        log_prefix=stack_path,
        cwd=absolute_path,
    )
    if rc_init or rc_main > 0:
        raise RuntimeError(
            f"Unexpected Terraform error for command: {terraform_main_command}"
        )
    else:
        return True, "OK"

    # return True, "OK"


def find_paths_with_file(file_name):
    try:
        completed_process = subprocess.run(
            ["git", "ls-files", f"*{file_name}"],
            check=True,
            text=True,
            stdout=subprocess.PIPE,
        )

        paths = set()
        for file_path in completed_process.stdout.splitlines():
            dir_path = "/".join(file_path.split("/")[:-1])
            if dir_path.startswith("modules/"):
                continue
            if dir_path.startswith("infra/"):
                dir_path = dir_path[len("infra/") :]
            if dir_path.startswith("kinesis") or dir_path.startswith(
                "aws-transfer-sftp"
            ):
                paths.add(dir_path)

        return paths
    except subprocess.CalledProcessError as e:
        LOGGER.error(f"Error executing git command: {e}")
        return set()


def read_dependencies(path):
    depends_on_file = os.path.join(path, ".depends-on.txt")
    if os.path.exists(depends_on_file):
        with open(depends_on_file, "r") as f:
            dependencies = [line.strip() for line in f.readlines()]
        return dependencies
    return []


def wait_for_dependencies(dependencies, completed, failed, path):
    """
    Waits for all dependencies to be satisfied before returning.
    Now also checks for failed dependencies and raises an exception if found.
    """
    while dependencies:
        pending_dependencies = [dep for dep in dependencies if dep not in completed]
        if any(dep in failed for dep in pending_dependencies):
            failed_deps = [dep for dep in pending_dependencies if dep in failed]
            raise Exception(
                f"{path}: Cannot proceed. Dependencies failed: {failed_deps}"
            )
        if not pending_dependencies:
            return
        LOGGER.warning(
            f"{path}: Waiting for dependencies to complete: {pending_dependencies}"
        )
        time.sleep(5)


def execute_with_dependencies(paths):
    completed = set()  # Track completed paths
    errored_modules = set()  # Track paths of modules where an error occurred
    failed = {}  # Track failures

    with ThreadPoolExecutor(max_workers=terraform_threads) as executor:
        future_to_path = {}

        for path in paths:
            dependencies = read_dependencies(path)
            if dependencies:
                # Submit a future for waiting dependencies
                dep_future = executor.submit(
                    wait_for_dependencies, dependencies, completed, failed, path
                )
                # Submit the main task to be executed after dependencies are resolved
                future = executor.submit(
                    lambda p=dep_future, x=path: (
                        do_terraform(x) if p.result() is None else None
                    )
                )
            else:
                future = executor.submit(do_terraform, path)
            future_to_path[future] = path

        for future in as_completed(future_to_path):
            path = future_to_path[future]
            try:
                result = future.result()
                if (
                    result is None or not result[0]
                ):  # Assuming do_terraform returns (success: bool, message: str)
                    errored_modules.add(path)
                    failed[path] = "Failed due to internal error"  # Mark as failed
                else:
                    completed.add(path)  # Mark as completed
                LOGGER.info(f"Thread result for {path}: {result}")
            except Exception as exc:
                errored_modules.add(path)
                failed[path] = str(exc)  # Mark as failed
                LOGGER.error(f"{path} generated an exception: {exc}")

    # Check if any errors occurred and print the errored module paths
    if errored_modules:
        LOGGER.error("One or more modules failed to process correctly.")
        for module_path in errored_modules:
            LOGGER.error(f"Module with error: {module_path}")
        sys.exit(1)


# Your main execution logic here
if __name__ == "__main__":
    if tf_root_module_path:
        do_terraform(tf_root_module_path)
    else:
        paths = find_paths_with_file("provider.tf")
        for path in paths:
            LOGGER.warning(f"Discovered: {path}")
        # TODO finish later if we want multi modulue deploy in one go
        # execute_with_dependencies(paths)