A small script to generate TiDB cluster topology file

generate-topology.py

import sys
import subprocess
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--cluster', '-c', required=True, help="The TiDB cluster name")
args = parser.parse_args()
cluster = args.cluster
print(f"Cluster: {cluster}")

pipe = subprocess.Popen(f'fop cluster {cluster} show | grep "ap-northeast-1" | grep -v aws | sort -V', stdout=subprocess.PIPE, shell=True)
results = pipe.stdout.readlines()
pipe.wait()

def format_tikv_host(index, ip, az, host):
  return f'TiKV-{index} ansible_host={ip} labels="az={az}, host={host}"'

drainer_hosts = []
pd_hosts = []
prometheus_hosts = []
pump_hosts = []
tidb_hosts = []
tikv_hosts = []
tikv_i = {}
index = 0

for line in results:
  line = line.decode('utf-8')
  if 'Load Balancers:' in line:
    break
  if 'Drainer' in line:
    drainer_hosts.append(line.split()[2])
  if 'PD' in line:
    pd_hosts.append(line.split()[2])
  elif 'Prometheus' in line:
    prometheus_hosts.append(line.split()[2])
  elif 'Pump' in line:
    pump_hosts.append(line.split()[2])
  elif 'TiDB' in line:
    tidb_hosts.append(line.split()[2])
  elif 'TiKV' in line:
    columns = line.split()
    ip = columns[2]
    az = columns[3]
    if az in tikv_i:
      tikv_i[az] += 1
    else:
      tikv_i[az] = 1
    host = f"{az.split('-')[2]}-{tikv_i[az]}"
    tikv_hosts.append(format_tikv_host(index, ip, az, host))
    index += 1

print("# pd_servers")
for h in pd_hosts:
  print(f"{h}")

print()

print("# tidb_servers")
for h in tidb_hosts:
  print(f"{h}")

print()

print("# tikv_servers")
for h in tikv_hosts:
  print(h)

print()

print("# pump_servers")
for h in pump_hosts:
  print(f"{h}")

print()

print("# drainer_servers")
for h in drainer_hosts:
  print(f"{h}")

print()

print("# monitoring_servers")
for h in prometheus_hosts:
  print(f"{h}")

print()
print("=" * 180)

print("# tikv_servers")
for h in tikv_hosts:
  print(h.split()[1].split("=")[1])