This script is used to set or remove environment variables in all services in a given environment.

set_env_var_to_all_services_in_env.py

"""
This script is used to set or remove environment variables in all services in a given environment.
Usage: python set_env_var_to_all_services_in_env.py <mode> <env> <var_name> <var_value>
"""

import glob
import os
import sys
from ruamel.yaml import YAML

# Specify the path to your YAML file
yaml_file_path = 'config.yaml'


def find_yaml_files(env):
    # Define the pattern to match files
    pattern = os.path.join('overrides', 'apps', '*', '.auto-generated', 'config_vars', f'{env}.yaml')

    # Use glob to find all files that match the pattern
    files = glob.glob(pattern)

    return files

def set_env_var(yaml_file_path, var_name, var_value):
    print(f'Processing {yaml_file_path}')
    yaml = YAML()
    yaml.width = 2 ** 32  # Prevent the insertion of linebreaks in long lines
    yaml.preserve_quotes = True  # Optional: if you want to preserve quotes

    # Load the current content of the YAML file
    with open(yaml_file_path, 'r') as file:
        data = yaml.load(file)  # This maintains the order of keys

    if not data:
        data = {}

    # Check if 'envVars' key exists, if not create it as an empty dict
    if 'envVars' not in data:
        data['envVars'] = {}

    # Add or update 'var' under 'envVars'
    data['envVars'][var_name] = var_value

    # Write the updated content back to the YAML file, preserving the original order
    with open(yaml_file_path, 'w') as file:
        file.write('---\n')  # Manually write the dashes at the top of the file
        yaml.dump(data, file)

def rm_env_var(yaml_file_path, var_name):
    print(f'Processing {yaml_file_path}')
    yaml = YAML()
    yaml.width = 2 ** 32  # Prevent the insertion of linebreaks in long lines
    yaml.preserve_quotes = True  # Optional: if you want to preserve quotes

    # Load the current content of the YAML file
    with open(yaml_file_path, 'r') as file:
        data = yaml.load(file)  # This maintains the order of keys

    if not data or var_name not in data['envVars']:
        return

    del data['envVars'][var_name]

    # Write the updated content back to the YAML file, preserving the original order
    with open(yaml_file_path, 'w') as file:
        file.write('---\n')  # Manually write the dashes at the top of the file
        yaml.dump(data, file)


if __name__ == '__main__':

    mode = sys.argv[1]  # 'set' or 'rm'
    env = sys.argv[2]

    files = find_yaml_files(env)
    if not files:
        raise(ValueError(f'No files found for {env}'))

    if mode == 'set':
        var_name = sys.argv[3]
        var_value = sys.argv[4]
        for file in files:
            set_env_var(file, var_name, var_value)
            print(f'Updated {file}')
    elif mode == 'rm':
        var_name = sys.argv[3]
        for file in files:
            rm_env_var(file, var_name)
            print(f'Updated {file}')
    else:
        raise(ValueError('Invalid mode'))