142 lines
5.2 KiB
Python
Executable File
142 lines
5.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
|
||
import yaml
|
||
import sys
|
||
import os
|
||
import shutil
|
||
from jinja2 import Environment, FileSystemLoader
|
||
|
||
def load_cluster_config(config_path):
|
||
"""Load cluster configuration from YAML file"""
|
||
with open(config_path, 'r') as f:
|
||
return yaml.safe_load(f)
|
||
|
||
def generate_inventory(cluster_config, inventory_path):
|
||
"""Generate Ansible inventory file from cluster configuration"""
|
||
with open(inventory_path, 'w') as f:
|
||
f.write("# This file is auto-generated by tools/generate-inventory.py\n")
|
||
f.write("# Do not edit – your changes will be overwritten.\n")
|
||
f.write("# Edit cluster.yml and re-run the generator instead.\n\n")
|
||
|
||
# Master group
|
||
f.write("[airflow_master]\n")
|
||
for hostname, ip in cluster_config['master'].items():
|
||
f.write(f"{hostname} ansible_host={ip}\n")
|
||
|
||
f.write("\n")
|
||
|
||
# Workers group
|
||
f.write("[airflow_workers]\n")
|
||
for hostname, config in cluster_config['workers'].items():
|
||
f.write(f"{hostname} ansible_host={config['ip']}\n")
|
||
|
||
def generate_host_vars(cluster_config, host_vars_dir):
|
||
"""Generate host-specific variables"""
|
||
# Create host_vars directory if it doesn't exist
|
||
os.makedirs(host_vars_dir, exist_ok=True)
|
||
|
||
# Clear existing host_vars files to avoid stale configurations
|
||
for filename in os.listdir(host_vars_dir):
|
||
file_path = os.path.join(host_vars_dir, filename)
|
||
try:
|
||
if os.path.isfile(file_path) or os.path.islink(file_path):
|
||
os.unlink(file_path)
|
||
elif os.path.isdir(file_path):
|
||
shutil.rmtree(file_path)
|
||
except Exception as e:
|
||
print(f'Failed to delete {file_path}. Reason: {e}')
|
||
|
||
# Get master IP for Redis configuration
|
||
master_ip = list(cluster_config['master'].values())[0]
|
||
|
||
# Generate master host variables
|
||
for hostname, ip in cluster_config['master'].items():
|
||
host_vars_file = os.path.join(host_vars_dir, f"{hostname}.yml")
|
||
with open(host_vars_file, 'w') as f:
|
||
f.write("---\n")
|
||
f.write(f"# Variables for {hostname}\n")
|
||
f.write(f"master_host_ip: {ip}\n")
|
||
f.write(f"redis_port: 52909\n")
|
||
|
||
# Generate worker proxy variables
|
||
for hostname, config in cluster_config['workers'].items():
|
||
host_vars_file = os.path.join(host_vars_dir, f"{hostname}.yml")
|
||
worker_proxies = config.get('proxies', [])
|
||
|
||
with open(host_vars_file, 'w') as f:
|
||
f.write("---\n")
|
||
f.write(f"# Variables for {hostname}\n")
|
||
f.write(f"master_host_ip: {master_ip}\n")
|
||
f.write("redis_port: 52909\n")
|
||
if worker_proxies:
|
||
f.write("worker_proxies:\n")
|
||
for proxy in worker_proxies:
|
||
f.write(f" - \"{proxy}\"\n")
|
||
|
||
def generate_group_vars(cluster_config, group_vars_dir):
|
||
"""Generate group-level variables"""
|
||
# Create group_vars directory if it doesn't exist
|
||
os.makedirs(group_vars_dir, exist_ok=True)
|
||
|
||
# Create group_vars/all directory if it doesn't exist
|
||
all_vars_dir = os.path.join(group_vars_dir, "all")
|
||
os.makedirs(all_vars_dir, exist_ok=True)
|
||
|
||
# Define path for the generated file and remove it if it exists to avoid stale data.
|
||
# This is safer than removing the whole directory, which would delete vault.yml.
|
||
all_vars_file = os.path.join(all_vars_dir, "generated_vars.yml")
|
||
if os.path.exists(all_vars_file):
|
||
os.remove(all_vars_file)
|
||
external_ips = cluster_config.get('external_access_ips', [])
|
||
|
||
# Get master IP for Redis configuration
|
||
master_ip = list(cluster_config['master'].values())[0]
|
||
|
||
with open(all_vars_file, 'w') as f:
|
||
f.write("---\n")
|
||
f.write("# This file is auto-generated by tools/generate-inventory.py\n")
|
||
f.write("# Do not edit – your changes will be overwritten.\n")
|
||
f.write(f"master_host_ip: {master_ip}\n")
|
||
f.write("redis_port: 52909\n")
|
||
f.write("external_access_ips:\n")
|
||
if external_ips:
|
||
for ip in external_ips:
|
||
f.write(f" - \"{ip}\"\n")
|
||
else:
|
||
f.write(" []\n")
|
||
|
||
def main():
|
||
if len(sys.argv) != 2:
|
||
print("Usage: python3 generate-inventory.py <cluster-config-file>")
|
||
sys.exit(1)
|
||
|
||
config_path = sys.argv[1]
|
||
|
||
# Check if config file exists
|
||
if not os.path.exists(config_path):
|
||
print(f"Error: Configuration file {config_path} not found")
|
||
sys.exit(1)
|
||
|
||
# Load cluster configuration
|
||
cluster_config = load_cluster_config(config_path)
|
||
|
||
# Generate inventory file
|
||
inventory_path = "ansible/inventory.ini"
|
||
generate_inventory(cluster_config, inventory_path)
|
||
print(f"Generated {inventory_path}")
|
||
|
||
# Generate host variables
|
||
host_vars_dir = "ansible/host_vars"
|
||
generate_host_vars(cluster_config, host_vars_dir)
|
||
print(f"Generated host variables in {host_vars_dir}")
|
||
|
||
# Generate group variables
|
||
group_vars_dir = "ansible/group_vars"
|
||
generate_group_vars(cluster_config, group_vars_dir)
|
||
print(f"Generated group variables in {group_vars_dir}")
|
||
|
||
print("Inventory generation complete!")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|