from ruamel.yaml import YAML, comments from ruamel.yaml.scalarstring import DoubleQuotedScalarString as dq import os base_dir = "/home/rpetrovich/code/Repos/lmru--devops--argocd-apps/" yaml = YAML() yaml.preserve_quotes = True yaml.explicit_start = True yaml.indent(mapping=2, sequence=4, offset=2) yaml.allow_duplicate_keys = True def load_yaml(file_path): with open(file_path, 'r') as file: data = list(yaml.load_all(file)) return data def save_yaml(data, file_path): with open(file_path, 'w') as file: yaml.dump_all(data, file) def extract_cluster_info(node_path): parts = node_path.split('/') nodes_index = parts.index('nodes') env = parts[nodes_index - 2] cluster_index = parts.index('clusters') + 1 cluster_name = parts[cluster_index] subpath = '-'.join(parts[cluster_index + 2:nodes_index - 1]) if nodes_index - cluster_index > 2 else '' return cluster_name, env, subpath def add_label_if_needed(node_group): excluded_labels = [ "d8-", "opendistro", "opensearch", "otel-", "system", "mimir-", "prometheus", "frontend", "elastic", "master", "loadbalancer", "deckhouse", "router", "moon-", "sentry", "odfe-", "obs-grafana", "devops-core-argocd", "victoria-", "clickhouse", "redis", "rabbit", "loki","cassandra", "jaeger", "vm-", "memcached", "kafka" ] name = node_group['metadata']['name'] print(f"Processing NodeGroup: {name}") if any(excluded in name for excluded in excluded_labels): print(f"Excluded by name: {name}") return False # Пропускаем обработку, если имя в списке исключений node_template = node_group['spec'].setdefault('nodeTemplate', comments.CommentedMap()) labels = node_template.setdefault('labels', comments.CommentedMap()) filtered_labels = {k: v for k, v in labels.items() if not k.startswith("threshold.")} if any(excluded in label for label in filtered_labels for excluded in excluded_labels): print(f"Excluded by existing labels: {filtered_labels}") return False # Пропускаем обработку, если метки в списке исключений new_label_key = 'node-role.k8s.lmru.tech/application' # Добавляем логику: проверяем метку, но продолжаем проверку тейнтов label_already_present = new_label_key in labels if label_already_present: print(f"Label already present for {name}, but checking taints...") else: labels[new_label_key] = "" print(f"Label added for {name}: {new_label_key}") # Возвращаем True для продолжения обработки, даже если метка уже есть return True def find_cluster_file(base_dir, cluster_name, env): target_directory = f"{base_dir}/common/aqua-kube-enforcer/{env}/values/" if not os.path.exists(target_directory): print(f"Directory does not exist: {target_directory}") return None for file in os.listdir(target_directory): if cluster_name in file and env in file and file.endswith('cluster.yaml') and '-elk-' not in file: return os.path.join(target_directory, file) print(f"No matching cluster file found that contains both '{cluster_name}' and '{env}' in {target_directory}") return None def process_yaml_file(node_file): if '/elk/' in node_file: print(f"Skipping: {node_file} due to '/elk/'") return cluster_name, env, _ = extract_cluster_info(node_file) data = load_yaml(node_file) if not data: print(f"No data found in {node_file}") return changed = False tolerations = [] for doc in data: if doc is None: print(f"Skipping invalid or empty YAML document in {node_file}") continue # Проверяем метки и продолжаем даже если метка уже есть if doc.get('kind') == "NodeGroup" and add_label_if_needed(doc): changed = True taints = doc.get('spec', {}).get('nodeTemplate', {}).get('taints', []) if taints: print(f"Taints found in NodeGroup {doc['metadata']['name']}: {taints}") for taint in taints: if 'key' in taint and 'value' in taint: tolerations.append({ "key": taint['key'], "operator": "Equal", "effect": dq(taint.get('effect', '')), "value": taint['value'] }) print(f"Adding toleration for {doc['metadata']['name']}: {taint}") if changed: save_yaml(data, node_file) print(f"Changes saved to {node_file}") if tolerations: cluster_file = find_cluster_file(base_dir, cluster_name, env) if cluster_file: print(f"Found cluster file for {cluster_name} in {env}: {cluster_file}") update_cluster_file(cluster_file, tolerations) else: print(f"No cluster file found for {cluster_name} in {env}") else: print(f"No tolerations found to add for {node_file}") def update_cluster_file(cluster_file, new_tolerations): if not cluster_file: print("No matching cluster file found.") return print(f"Attempting to update cluster file: {cluster_file} with tolerations: {new_tolerations}") try: data = load_yaml(cluster_file) updated = False # Флаг для отслеживания изменений for doc in data: # Убедимся, что раздел enforcer присутствует enforcer = doc.setdefault('enforcer', comments.CommentedMap()) # Добавляем nodeSelector, если отсутствует enforcer.setdefault('nodeSelector', {})['node-role.k8s.lmru.tech/application'] = '' # Проверяем и добавляем новые tolerations existing_tolerations = enforcer.setdefault('tolerations', []) for new_tol in new_tolerations: # Проверяем, если такая же toleration уже присутствует if not any( tol['key'] == new_tol['key'] and tol['operator'] == new_tol['operator'] and tol['effect'] == new_tol['effect'] and tol['value'] == new_tol['value'] for tol in existing_tolerations ): existing_tolerations.append(new_tol) updated = True # Устанавливаем флаг, что произошли изменения if updated: save_yaml(data, cluster_file) print(f"Updated cluster file successfully: {cluster_file}") else: print(f"No updates were needed for {cluster_file}") except Exception as e: print(f"Failed to update cluster file {cluster_file}: {e}") def process_directory(directory): for root, _, files in os.walk(directory, topdown=True): if "stage" not in root: # comment if you don't want to restrict "stage" only continue for file in files: if file == 'node_groups.yaml': full_path = os.path.join(root, file) print(f"Processing: {full_path}") process_yaml_file(full_path) process_directory(base_dir)