187 lines
7.5 KiB
Python
187 lines
7.5 KiB
Python
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)
|